PKV Prefix Attention in Transformers

• PKV Prefix Attention is a conditional mechanism that integrates continuous property vectors into transformer layers to preserve physical relationships.
• It projects properties into low-dimensional key/value embeddings, bypassing tokenization and reducing catastrophic forgetting during fine-tuning.
• Empirical results demonstrate improved validity, structure recovery, and robust performance in generative design for materials discovery.

Property-Key-Value (PKV) Prefix Attention is a conditional attention mechanism for transformer-based generative models that directly injects continuous property information into every layer of the network by augmenting the key and value states in self-attention. It was introduced as part of the CrystaLLM-π framework for materials discovery and inverse design, specifically addressing the limitations of discrete token-based conditioning and catastrophic forgetting during fine-tuning on property-driven tasks (Bone et al., 26 Nov 2025).

1. Motivation and Rationale for Continuous Property Injection

Standard transformers condition on discrete tokens, which forces continuous physical or functional quantities—such as band gap, stable lattice energy, or diffraction intensities—to be quantized into tokens. This process destroys their ordinal and proximity relationships, requiring large datasets for effective learning and often leading to sample inefficiency and poor interpolation. PKV Prefix Attention projects continuous property vectors directly into the attention mechanism. This approach:

• Preserves the structure and continuous nature of physical properties through low-dimensional embeddings.
• Eliminates disruption to the tokenization pipeline by sidestepping the need for sequence-level property tokens.
• Ensures every layer in the transformer “attends” to the functional requirement of the design target, greatly reducing catastrophic perturbation of pre-trained weights during conditional fine-tuning.

2. Formal Definition and Mathematical Formulation

Let $x \in \mathbb{R}^{L \times d_{\text{model}}}$ denote the sequence of token embeddings, $p \in \mathbb{R}^{P}$ be the continuous property vector, and $$W_Q, W_K, W_V \in \mathbb{R}^{d_{\text{model}} \times d_k}$$ the attention projections for $h$ heads. The PKV Prefix process is as follows:

a) Projection of Properties to Prefix Key/Value (KV) Tensors

• Compute $$h = \mathrm{ReLU}\big(\mathrm{LayerNorm}(W_{\text{hidden}}\, p)\big)$$, with $$W_{\text{hidden}} \in \mathbb{R}^{d_{\text{hidden}}\times P}$$.
• Map $h$ to stacked key-value tensors: $H_{\text{KV}} = W_{\text{KV}} \cdot h$, $$H_{\text{KV}} \in \mathbb{R}^{N \times h \times 2 \times d_k}$$, to produce per-layer, per-head $K_{\text{prefix}}$ and $V_{\text{prefix}}$.

b) Standard Sequence KV Calculation

• $Q = x W_Q \in \mathbb{R}^{L\times (h\cdot d_k)}$
• $K_{\text{seq}} = x W_K$, $$V_{\text{seq}} = x W_V \in \mathbb{R}^{L\times (h\cdot d_k)}$$

c) Concatenation of Prefix and Sequence KV

• For block $i$, head $j$: $${K'}_{i,j} = [K_{\text{prefix}}[i,j]; K_{\text{seq},i,j}] \in \mathbb{R}^{(P + L) \times d_k}$$, and likewise for $V'_{i,j}$.

d) Prefix-Conditioned Attention

• $$A_{\text{out}}^{(i,j)} = \mathrm{softmax}\left(Q^{(i,j)}{K'}_{i,j}^T / \sqrt{d_k}\right)V'_{i,j} \in \mathbb{R}^{L \times d_k}$$
• Head outputs are concatenated and linearly projected via $W_O$.

This architecture enables direct injection of continuous properties at each transformer layer.

3. Integration within Transformer Architectures

Within each transformer decoder block:

• Compute Q, K, V projections from input token embeddings.
• Retrieve per-block, per-head $K_{\text{prefix}}$, $V_{\text{prefix}}$ from the Prefix encoder, a compact MLP trained from scratch.
• Concatenate these “ghost tokens” to the sequence K/V and compute attention.
• The backbone transformer weights are typically frozen or receive a much lower learning rate, preserving unsupervised pre-training (e.g., on Crystallographic Information Files).
• Continuous properties never undergo tokenization, bypassing any change to the pre-trained text processing pipeline.

The Prefix encoder is responsible for mapping the structured property vector to the stack of KV projections per layer and head—a distinct role from traditional prompt or prefix-tuning methods.

4. Comparative Analysis: Prefix vs. Residual Conditioning

PKV Prefix Attention constitutes a form of “hard conditioning,” providing a strong steering effect:

• Prefix (hard conditioning): Extends the sequence context by appending property-derived K/V as additional positions, functionally similar to extra “ghost” tokens. This mechanism offers robust property control, especially effective in data-rich training environments.
• PKV Residual Attention (“Residual”): Implements “soft conditioning” by running a parallel attention using property-derived K/V ($K_{\text{res}}, V_{\text{res}}$) and combining its output with the base attention via a learnable scalar $\alpha$ per layer ($$A_{\text{out}} = A_{\text{base}} + \alpha A_{\text{res}}$$). This allows dynamic scaling of the property influence and graceful handling of missing properties, particularly beneficial in limited data regimes.

The table below briefly contrasts these two approaches:

Aspect	Prefix (Hard)	Residual (Soft)
Conditioning strength	Strong, structural bias	Gentle, tunable
Architectural effect	Augments K/V sequence	Parallel residual sum
Pre-training retention	Weaker in low-data regime	Stronger in low-data
Support for missing props	Not explicit	Natural masking

5. Implementation Details

• Model configuration: $d_{\text{model}}=512$, $h=8$ heads, $d_k=64$, $N\approx12$ layers.
• Prefix encoder: hidden dimension $d_{\text{hidden}}\approx1024$, ReLU, LayerNorm, dropout $0.1$.
• Parameter shapes: $W_{\text{hidden}}: (d_{\text{hidden}}\times P)$; $$W_{\text{KV}}: (2\cdot h\cdot d_k \times d_{\text{hidden}})$$; $$K_{\text{prefix}}, V_{\text{prefix}}: (N\times h\times d_k)$$.
• Initialization: Xavier uniform for linear layers; LayerNorm $\gamma=1$, $\beta=0$; Residual scalar $\alpha=0$.
• Optimization: Backbone with learning rate $\sim5\times10^{-6}$, Prefix/Residual encoder with $\sim5\times10^{-4}$ (AdamW), $0.01$ weight decay, $2\%$ warmup, $\sim50$k fine-tuning steps, early stopping.
• Data flow: Conditioning layers (Prefix or Residual) exclusively update the property-injection path, minimizing disturbance to core generative modeling capabilities.

6. Empirical Performance and Applications

PKV Prefix Attention and its residual variant have been validated across broad generative design tasks:

• Transfer learning and pre-training (MP Band Gap, 53k compounds):
  • Pre-trained models increase validity by $+20.4\%$ and VSUN quality by $+14.4\%$ versus training from scratch.
  • PKV methods retain pre-training priors more reliably than sequence-level discrete property tokenization, which fails completely in this context.
• Data scalability (“MatterGen Density,” 1k–600k samples):
  • At scale ($653$k), Prefix achieves MAE $\approx 5.26$ g/cm³, outperforming Prepend ($6.19$) and Residual ($6.35$).
  • For low data ($1$k), Residual attains higher valid structure rates ($88\%$) than Prefix ($82\%$) or Prepend ($73\%$).
• X-ray Diffraction (XRD) structure recovery:
  • Prefix: MP-20 (62.7% match, RMS-d=0.0444), improved to 69% by skipping validity filters.
  • Residual: Jarvis-DFT (66.2% match rate, RMS-d=0.0347, halving RMS-d relative to DiffractGPT).
  • Experimental structures: 44.8% match rate for unseen COD (20-consistent), MAE in lattice parameters reduced by ≈0.25 Å, $R^2$ increased by ≈0.13.
  • TiO₂ polymorphs: near-perfect structure discovery when given space group prompt.
• Photovoltaic SLME-targeted generation:
  • Prefix: $1,546$ structurally novel, stable candidates above $20\%$ predicted SLME.
  • Band-gap distribution clustered at Shockley–Queisser optimum ($1.2$–$1.4$ eV) despite absence of explicit band-gap supervision.
  • High SLME, DFT-validated structures produced, including Cs$_2$NaInAs$_2$, Cs$_2$NaGaAs$_2$, NaHfCuS$_3$, Rb$_2$(NbBr$_3$)$_3$.

A plausible implication is that PKV Prefix conditioning is preferable in data-rich settings seeking strong property control, while Residual is advantageous for delicate adaptation in lower-data or incomplete-property environments.

7. Significance and Outlook

Property-Key-Value Prefix Attention enables transformer-based generative models to perform inverse materials design by directly leveraging continuous physical property information in every layer. By bypassing sequence tokenization and preserving pre-trained generative knowledge, the technique establishes a unified, flexible, and computationally lightweight solution for structure recovery, polymorph differentiation, and functionally targeted candidate generation. The flexibility between Prefix and Residual conditioning deepens architectural adaptability, supporting robust and fine-grained control across a range of data regimes and conditioning scenarios (Bone et al., 26 Nov 2025).