IA3 › Worked Example: Computing IA3-Adapted Attention

Let's work through a concrete numerical example to solidify the mechanics. We will trace exactly how IA3 modifies a small attention computation, showing every step of the calculation.

Consider a tiny transformer with hidden dimension $d = 4$, processing a sequence of 3 tokens. The input matrix $X$ and frozen weight matrices $W_K$ and $W_V$ are:

$$X = \begin{bmatrix} 1 & 0 & 1 & 0 \\ 0 & 1 & 0 & 1 \\ 1 & 1 & 0 & 0 \end{bmatrix} \quad W_K = \begin{bmatrix} 0.5 & 0.1 & 0.2 & 0.0 \\ 0.0 & 0.5 & 0.1 & 0.3 \\ 0.2 & 0.0 & 0.5 & 0.1 \\ 0.1 & 0.3 & 0.0 & 0.5 \end{bmatrix} \quad W_V = I_4$$

For simplicity, we set $W_V = I_4$ (the identity matrix), so value vectors equal the input embeddings. The IA3 rescaling vector for keys is:

$$\ell_K = \begin{bmatrix} 2.0 & 0.5 & 1.5 & 0.3 \end{bmatrix}$$

Step 1: Compute the standard key matrix. Multiply $X$ by $W_K$:

$$K = X W_K = \begin{bmatrix} 0.7 & 0.1 & 0.7 & 0.1 \\ 0.1 & 0.8 & 0.1 & 0.8 \\ 0.5 & 0.6 & 0.3 & 0.4 \end{bmatrix}$$

Each row is the key vector for one of the 3 tokens.

Step 2: Apply IA3 rescaling. Multiply each element of $K$ by the corresponding entry of $\ell_K$. Since $\ell_K = [2.0, 0.5, 1.5, 0.3]$:

$$K' = \ell_K \odot K = \begin{bmatrix} 2.0 \times 0.7 & 0.5 \times 0.1 & 1.5 \times 0.7 & 0.3 \times 0.1 \\ 2.0 \times 0.1 & 0.5 \times 0.8 & 1.5 \times 0.1 & 0.3 \times 0.8 \\ 2.0 \times 0.5 & 0.5 \times 0.6 & 1.5 \times 0.3 & 0.3 \times 0.4 \end{bmatrix} = \begin{bmatrix} 1.4 & 0.05 & 1.05 & 0.03 \\ 0.2 & 0.4 & 0.15 & 0.24 \\ 1.0 & 0.3 & 0.45 & 0.12 \end{bmatrix}$$

Notice the effect: dimensions 1 and 3 (indices 0 and 2), which had large initial values, are amplified further. Dimensions 2 and 4 (indices 1 and 3) are compressed. The first token's key has been pulled strongly toward dimensions 1 and 3, while its contributions from dimensions 2 and 4 are nearly suppressed.

Step 3: Observe the effect on attention scores. Suppose we compute a simple self-attention score for the first token querying itself (ignoring the multi-head structure for clarity). If the query for token 1 is $q = [0.7, 0.1, 0.7, 0.1]$ (the same as its unscaled key vector for this example), then:

Without IA3: $q \cdot k_1 = 0.7^2 + 0.1^2 + 0.7^2 + 0.1^2 = 0.49 + 0.01 + 0.49 + 0.01 = 1.00$

With IA3: $q \cdot k_1' = 0.7 \times 1.4 + 0.1 \times 0.05 + 0.7 \times 1.05 + 0.1 \times 0.03 = 0.98 + 0.005 + 0.735 + 0.003 = 1.723$

The amplified dimensions 1 and 3 contribute much more to the attention score after IA3. This larger score makes token 1 appear more needed to itself. Meanwhile, the other tokens with less activation in dimensions 1 and 3 will see their keys relatively less scaled, changing the relative attention distribution.

Step 4: Count the parameters. For this toy example with $d = 4$ and $d_{ff} = 16$ (assuming the $4\times$ expansion), IA3 introduces:

• $|\ell_K| = 4$ parameters
• $|\ell_V| = 4$ parameters
• $|\ell_{ff}| = 16$ parameters
• Total: 24 parameters

The frozen weight matrices $W_K$ and $W_V$ together contain $4 \times 4 + 4 \times 4 = 32$ parameters, plus $W_Q$ and the output projection $W_O$ would add another 32. In this toy scale, IA3 adds 24 parameters to a block with 64 frozen parameters: 37% overhead. At real scales (LLaMA-7B with $d = 4096$), this ratio shrinks sharply. The frozen weights per layer include four $4096 \times 4096$ matrices totaling 67 million parameters, while IA3 adds only about 19,000 parameters, which is 0.028% overhead.