IA3 › The IA3 Formulation

To understand IA3 precisely, we need to first establish where in the transformer architecture it intervenes and why those specific locations were chosen. The choice of intervention points is not arbitrary: the IA3 authors made deliberate decisions based on where learned rescaling would have the most impact on downstream task performance, and understanding their reasoning deepens your intuition for the method as a whole.

A transformer block processes input through two major computational streams. The first is multi-head self-attention, which mixes information across positions. The second is the position-wise feed-forward network, which applies learned nonlinear transformations to each position independently. IA3 intervenes at three points: the key projections of attention, the value projections of attention, and the intermediate activations of the feed-forward network. Each intervention point controls a different aspect of how information flows through the model.

Skipping the query projection is a deliberate choice. Queries encode what a position is searching for, and modifying them would change the pattern of what information each position attends to. But keys and values play complementary roles that are arguably more important for task adaptation. Keys determine which other positions appear needed to each query. Values determine what information gets retrieved when attention weights are high. By rescaling keys, IA3 can change the effective relevance of token positions without changing the queries themselves. By rescaling values, IA3 can change the magnitude of information retrieved from different positions. Together, these two interventions give IA3 fine-grained control over what information each position collects from the context. The decision to include feed-forward intermediate activations reflects the fact that feed-forward networks in transformers function as key-value memories that store factual knowledge; rescaling their intermediate activations allows IA3 to selectively activate or suppress learned knowledge pathways.

In multi-head attention, the model begins by computing queries, keys, and values through linear projections of the input:

$$\begin{aligned} Q &= XW_Q \\ K &= XW_K \\ V &= XW_V \end{aligned}$$

where:

• $Q$: the query matrix, shape $(n \times d)$
• $K$: the key matrix, shape $(n \times d)$
• $V$: the value matrix, shape $(n \times d)$
• $X$: the input sequence matrix, shape $(n \times d)$
• $W_Q$: the query projection weight, shape $(d \times d)$
• $W_K$: the key projection weight, shape $(d \times d)$
• $W_V$: the value projection weight, shape $(d \times d)$
• $n$: the number of tokens in the sequence
• $d$: the model hidden dimension

These projections then participate in scaled dot-product attention, where the query-key dot product determines attention weights and those weights are used to aggregate the values. The feed-forward network applies to each position independently after the attention output is combined and normalized, transforming the representation through a higher-dimensional space before projecting back down.

IA3 introduces three learned vectors per transformer layer:

• $\ell_K \in \mathbb{R}^{d}$ for rescaling the key projections
• $\ell_V \in \mathbb{R}^{d}$ for rescaling the value projections
• $\ell_{ff} \in \mathbb{R}^{d_{ff}}$ for rescaling feed-forward intermediate activations

The modified key and value computations become:

$$\begin{aligned} K' &= \ell_K \odot (XW_K) \\ V' &= \ell_V \odot (XW_V) \end{aligned}$$

where:

• $K'$: the rescaled key matrix, shape $(n \times d)$
• $V'$: the rescaled value matrix, shape $(n \times d)$
• $\ell_K$: the learned rescaling vector for keys, shape $(d,)$
• $\ell_V$: the learned rescaling vector for values, shape $(d,)$
• $XW_K$: the standard key projection of the input
• $XW_V$: the standard value projection of the input
• $\odot$: element-wise multiplication, broadcast across the sequence length dimension $n$

To see exactly how the broadcasting works, think about the shapes involved. The standard key projection $XW_K$ has shape $(n \times d)$: one $d$-dimensional key vector per token. The rescaling vector $\ell_K$ has shape $(d,)$. When you broadcast this multiplication, $\ell_K[i]$ multiplies the $i$-th dimension of every token's key vector simultaneously. If $\ell_K[i] = 2.0$, the $i$-th coordinate of every key in the sequence gets doubled. If $\ell_K[j] = 0.3$, the $j$-th coordinate of every key gets compressed to 30% of its original value. Every token in the sequence experiences the same rescaling, which is both a strength and a limitation: the rescaling is consistent and interpretable, but it cannot adapt differently to different positions.

For the feed-forward network, recall the standard two-layer structure. The input first passes through a linear projection to a higher-dimensional intermediate space, then through a nonlinearity, and finally through a second projection back to the original dimension:

$$\begin{aligned} h &= \gamma(XW_1) \\ \text{output} &= hW_2 \end{aligned}$$

where $\gamma$ is the activation function (typically GELU in modern transformers), $W_1$ maps from $d$ to $d_{ff}$, and $W_2$ maps from $d_{ff}$ back to $d$. The intermediate representation $h$ lives in the larger $d_{ff}$-dimensional space, which is where the model stores much of its factual and contextual knowledge.

IA3 applies its rescaling to this intermediate representation:

$$h' = \ell_{ff} \odot h$$

where:

• $h'$: the rescaled intermediate activation, shape $(n \times d_{ff})$
• $\ell_{ff}$: the learned rescaling vector for feed-forward, shape $(d_{ff},)$
• $h$: the original intermediate activation $\gamma(XW_1)$, shape $(n \times d_{ff})$
• $\odot$: element-wise multiplication, broadcast across sequence positions
• $\gamma$: the activation function, such as GELU
• $W_1$: the first feed-forward projection, shape $(d \times d_{ff})$

The full modified feed-forward computation then becomes:

$$\begin{aligned} h &= \gamma(XW_1) \\ h' &= \ell_{ff} \odot h \\ \text{output} &= h' W_2 \end{aligned}$$

The choice to rescale after the activation function rather than before is real. Before the activation, the values are linear projections with no sparsity structure. After the activation, especially with GELU or ReLU-like functions, many dimensions are near zero and a few are strongly activated. Rescaling after the activation means IA3 is working with the already-processed, nonlinearly transformed intermediate representations. It can amplify dimensions that are already highly activated (meaning the model found them relevant) or further suppress dimensions that are already near zero.

The key insight is that these rescaling operations work entirely in activation space. No weight matrix is ever modified. The model's learned linear transformations, its attention patterns in the absence of rescaling, and its feed-forward computations all remain perfectly intact. IA3's vectors form a thin layer of per-dimension "volume controls" that sit on top of the existing computations, redirecting information flow without changing the underlying computational primitives.

IA3 initializes all rescaling vectors to all-ones: $\ell_K = \ell_V = \ell_{ff} = \mathbf{1}$. This is not a minor implementation detail; it is the key to why IA3 starts training from a stable, interpretable baseline.

When all rescaling values equal 1, the operation $\ell \odot h$ simplifies to $1 \odot h = h$. Every activation passes through unchanged. The adapted model at initialization is therefore identical in behavior to the original pretrained model. There are no random perturbations, no scaling factors that might shift training dynamics, no initial distortion of the representations the model has learned over many billions of pretraining tokens. The first gradient step moves the rescaling vectors slightly away from the all-ones vector in whatever directions the task loss indicates are beneficial, and training proceeds from there.

This initialization strategy contrasts thoughtfully with LoRA's approach. In LoRA, matrix $B$ is initialized to all zeros and matrix $A$ is initialized with small random Gaussian values. The initial adaptation $BA = 0 \cdot A = 0$, so the LoRA-adapted model is also identical to the pretrained model at initialization. But LoRA reaches this through cancellation (zero times anything is zero), while IA3 reaches it through identity scaling (one times anything is itself). Both approaches reach the same goal of starting from the pretrained model, but through philosophically different mechanisms. The IA3 approach is arguably more direct: the adaptation starts as a no-op because every individual scaling value is exactly one, not because two matrices happen to multiply to zero.

The element-wise multiplication in IA3 has an elegant matrix-algebraic interpretation that reveals its basic constraints and expressivity. If we define $L_K = \text{diag}(\ell_K)$ as the $d \times d$ diagonal matrix with the rescaling vector $\ell_K$ on its main diagonal and zeros everywhere else, then the element-wise multiplication $\ell_K \odot (XW_K)$ is equivalent to right-multiplying the projected representation by this diagonal matrix:

$$\begin{aligned} K' &= \ell_K \odot (XW_K) && \text{(element-wise multiplication)} \\ &= (X W_K) L_K && \text{(rewrite as diagonal matrix multiplication)} \end{aligned}$$

where:

• $K'$: the rescaled key matrix
• $\ell_K$: the rescaling vector, shape $(d,)$
• $X$: the input sequence, shape $(n \times d)$
• $W_K$: the pretrained key projection, shape $(d \times d)$
• $L_K$: the diagonal matrix formed from $\ell_K$, shape $(d \times d)$

This reformulation exposes a key algebraic fact: we could equivalently think of IA3 as right-multiplying the weight matrix $W_K$ by the diagonal matrix $L_K$. The effective weight matrix after IA3 adaptation would be $W_K L_K$. Each column of $W_K$ gets scaled by the corresponding entry of $\ell_K$. This observation is what lets inference-time merging of IA3 into the base model, which we will cover later.

More importantly, this reformulation reveals the expressivity class that IA3 belongs to. A diagonal matrix can only stretch or compress coordinates independently along the existing axes. It cannot rotate the coordinate system. It cannot mix information between dimensions. It cannot create new directions in the representation space. Compare this to LoRA's update $BA$, where $B$ and $A$ are general $d \times r$ and $r \times d$ matrices respectively. The product $BA$ can express any linear transformation of rank at most $r$: rotations, reflections, projections onto arbitrary subspaces, and more. IA3's diagonal update is strictly less expressive than LoRA's rank-1 update for any $d > 1$. Yet IA3 often performs comparably in practice, which tells us something deep about the nature of fine-tuning.