IA3 (infused adapters by inhibiting and amplifying activations), explained

IA3, written (IA)${}^3$ and short for "Infused Adapter by Inhibiting and Amplifying Inner Activations", is a parameter-efficient fine-tuning method that learns three vectors per transformer block and uses them to rescale activations element by element: one for the attention keys, one for the attention values, and one for the hidden layer of the feed-forward network. Every entry starts at 1, so the model is unchanged at initialisation, and training learns which channels to amplify (values above 1) and which to inhibit (values below 1). For a Llama-style 8B model that is about half a million parameters, roughly 80 times fewer than a rank-16 LoRA on all linear layers. Because each scaling sits next to a matrix multiplication, it can be folded into the weights after training for zero inference overhead, or kept separate so that different requests in one batch use different tasks. IA3 is a parameterisation, not an objective; it is trained with SFT cross-entropy or any other loss, and it belongs to the same family as LoRA and prompt tuning.

The problem and where it came from#

IA3 was introduced by Liu, Tam, Muqeeth, Mohta, Huang, Bansal and Raffel (2022) in "Few-Shot Parameter-Efficient Fine-Tuning is Better and Cheaper than In-Context Learning". The paper's question was whether few-shot in-context learning, which reprocesses every demonstration at every prediction, could be beaten on both accuracy and cost by fine-tuning a tiny number of parameters on the same few examples. They wanted a method that adds very few parameters, reaches strong accuracy from a few dozen examples, and, like in-context learning, allows mixed-task batches. Prompt tuning met the first and third requirements, but they "were unable to attain reasonable accuracy with prompt tuning" on T0-3B and observed wild swings in validation accuracy. LoRA and adapters were accurate but modify weights, which complicates mixed-task batches.

IA3 was their answer. Combined with two extra loss terms and a recipe built on T0 (called T-Few), it reached 72.4% average accuracy on T0's held-out tasks, against 66.9% for zero-shot T0 and 66.6% for few-shot in-context learning with GPT-3 175B, at about 1,000 times fewer inference FLOPs than the latter (1.1e12 against 1.4e15 per example). On the RAFT benchmark T-Few scored 75.8%, the first result above the 73.5% human baseline. Among the PEFT methods they compared on T0-3B (including BitFit, adapters, Compacter, prompt tuning, prefix tuning, FISH mask, intrinsic SAID and LoRA), IA3 was the only one to beat full-model fine-tuning.

The parameterisation, with every symbol defined#

In each transformer block, with input activations $x$, attention uses $Q=x{W}_Q$, $K=x{W}_K$, $V=x{W}_V$ for each head, and the feed-forward network computes $\gamma (x{W}_1)W_2$ with nonlinearity $\gamma$ (for gated MLPs, the gated hidden activation). IA3 introduces three learned vectors $l_k\in {\mathbb{R}}^{d_k}$, $l_v\in {\mathbb{R}}^{d_v}$ and $l_{ff}\in {\mathbb{R}}^{d_{ff}}$ and multiplies them in by broadcasting over positions:

All three are initialised to ones. Only they are trained; every pretrained weight is frozen. Each block gets its own vectors, so the parameter count is $L(d_k+d_v+d_{ff})$ for an $L$-block decoder-only model (encoder-decoder models add a second set for cross-attention). Here $d_k$ and $d_v$ are the total key and value widths across heads, which in grouped-query attention is the key-value width.

IA3 as a constrained weight update

Since $l\odot (xW)=xW\mathrm{diag}(l)$, each IA3 vector is exactly a rescaling of the output channels of $W_K$ and $W_V$, and of the input channels of $W_2$:

In the update form used for LoRA, $\Delta W_K=W_K\mathrm{diag}(l_k-1)$. This is a full-rank update in general, but it lives in a space of only $d_k$ dimensions: it can make any channel louder or quieter, and nothing else. It cannot rotate a channel or mix two channels. LoRA's update $BA$ is the mirror image: it can point anywhere, but only in $r$ directions. The two are complementary, which the toy run below measures on real weights.

The gradient of each scale is a sum over positions of the activation times its backpropagated error. For the key scale, with $K_{t,j}$ the $j$-th key channel at position $t$ and $\tilde{K}=l_k\odot K$:

so a channel is amplified when its activations correlate positively with the direction the loss wants that channel to move. Scaling keys changes how sharply each head attends (the logits for that channel are multiplied by $l_{k,j}$), scaling values changes what each head writes, and scaling the MLP hidden layer turns individual neurons up or down. That is why the authors describe it as inhibiting and amplifying.

The losses that came with it

T-Few's accuracy came from IA3 plus two loss terms for multiple-choice and classification tasks scored by rank classification. With the correct answer $y$ and $N$ incorrect choices ${\hat{y}}^{(n)}$, they add an unlikelihood loss that pushes down every token of every wrong choice, and a length-normalised softmax over choices:

where $\beta (x,y)=\frac{1}{T}{\sum }_t\log p(y_t\mid x,y_{<t})$ is the length-normalised log-probability. The total is ${\mathcal{L}}_{\text{LM}}+{\mathcal{L}}_{\text{UL}}+{\mathcal{L}}_{\text{LN}}$ with no weights to tune. With full fine-tuning of T0-3B, adding ${\mathcal{L}}_{\text{LN}}$ raised average accuracy from 60.7% to 62.71% and adding both reached 63.3%. These losses are independent of IA3 and can be used with any PEFT method; for open-ended generation they do not apply.

Side by side with its siblings

The 8B count uses 32 layers, key and value widths of 1,024 (grouped-query attention) and an MLP width of 14,336: $32\times (1024+1024+14336)=524,288$, or 0.0065% of the model. DoRA is a useful comparison: it also learns a per-channel magnitude, but adds a low-rank directional update, which is roughly "IA3 plus LoRA" on every matrix.

Worked example: scaling one head by hand#

Take one head with $d_k=d_v=2$, no $\sqrt{d_k}$ scaling, a query $q=(1,1)$, and two content positions with keys $k_1=(1,0)$, $k_2=(0,1)$ and values $v_1=(2,0)$, $v_2=(0,2)$. Without IA3, the logits are $q\cdot k_1=1$ and $q\cdot k_2=1$, the weights are $(0.5,0.5)$ and the output is $(1,1)$.

Now apply $l_k=(2,0.5)$: the scaled keys are $(2,0)$ and $(0,0.5)$, the logits become 2 and 0.5, and the softmax gives $e^2/(e^2+e^{0.5})=7.389/9.038=0.818$ and 0.182. The output is $0.818v_1+0.182v_2=(1.635,0.365)$. The head now prefers position 1. Within a single head this is a change a prefix cannot make, because a prefix leaves the relative weights of content positions untouched, while IA3 reweights the content logits channel by channel. Add $l_v=(1,0)$, which silences value channel 2: the output becomes $(1.635,0)$. Four numbers changed the head's routing and what it writes. Folding into weights is immediate: $W_K^{\prime }=W_K\mathrm{diag}(2,0.5)$ and $W_V^{\prime }=W_V\mathrm{diag}(1,0)$ give the same outputs with no extra operations.

What IA3 cannot do with this head: make the output $(0,2)$ when attention is on position 1, because that requires mapping value channel 1 into channel 2, a rotation. A rank-1 LoRA on $W_V$ could.

Toy runs: what channel scaling can and cannot represent#

Two experiments, one on real weights and one on a small trained transformer, both seeded PyTorch.

What the update space contains. We took the key projection of layer 12 of Qwen2.5-0.5B-Instruct, a $128\times 896$ matrix (two key-value heads of width 64), and asked how much of a target change $\Delta W^{\star }$ each parameterisation can represent at best. IA3 on this matrix has 128 numbers and represents $\mathrm{diag}(l-1)W$ (row scaling in the $[\text{out},\text{in}]$ layout); the best $l$ has a closed form per row, $l_i-1=⟨W_i,\Delta W_i^{\star }⟩/\Vert W_i{\Vert }^2$. LoRA at rank $r$ represents the best rank-$r$ approximation of $\Delta W^{\star }$. We used two targets of equal norm: a random rescaling of the rows (scales drawn from $\mathcal{N}(0,0.2^2)$), which is exactly IA3-shaped, and a random rank-1 matrix, which is exactly LoRA-shaped.

IA3 captured all of the rescaling target and 0.13% of the rank-1 target. LoRA captured all of the rank-1 target at any rank, but only 16.2% of the rescaling target at rank 1, 32.2% at rank 4, 65.4% at rank 16 and 97.4% at rank 64 (65,536 parameters, 512 times IA3's count). A change that looks trivial to IA3 is high rank, because scaling the rows of a full-rank matrix produces a full-rank change. The two methods are not points on one efficiency curve; they cover different subspaces of possible updates. Which one fits a task depends on whether the task needs channels turned up and down or new directions.

A trained model, an old skill and a new one. We pretrained a 2-layer, width-64, 4-head decoder on digit sequences with a task token in front (copy, reverse or sort six digits), reaching 99.9%, 99.95% and 96.8% exact-match accuracy. Without the task token it reversed only 1.0% of inputs correctly. We then froze it and trained only adapters for 1,500 steps on two tasks with no task token: reverse (a skill the model has) and rotate-by-one (output the digits shifted left by one position, never seen in pretraining; 0.05% before adaptation). IA3 with 768 parameters ($l_k$, $l_v$ of width 64 and $l_{ff}$ of width 256 in each of 2 blocks) reached 100% on reverse at every learning rate we tried (1e-2, 3e-2, 1e-1) and 99.95% to 100% on rotate. So did rank-4 LoRA with 9,216 parameters. For comparison, a 10-token soft prompt with 640 parameters also reached 100% on reverse but at most 30.9% on rotate across a sweep of learning rates and training lengths (details in the prompt tuning article). Rescaling channels was enough to re-route attention for a new positional pattern in this small model, which a soft prompt at the input could not do. The learned scales were far from 1: the largest $|l-1|$ was 4.3 to 9.0 on rotate against 1.8 to 2.5 on reverse, a reminder that "small parameter count" does not mean "small change".

The algorithm and training loop#

freeze all base weights
for each block: l_k = ones(d_kv); l_v = ones(d_kv); l_ff = ones(d_ff)       # trainable
    patch: K = (x W_K) * l_k ; V = (x W_V) * l_v ; h = (gamma(x W1) * l_ff) W2
optional: initialise l_k, l_v, l_ff from vectors pre-trained on a multitask mixture

optimizer = Adafactor or Adam on all l vectors, lr ~ 3e-3, warmup 60 steps, linear decay
for step in range(1000):                         # T-Few: 1,000 steps, batch 8
    loss = L_LM + L_UL + L_LN                     # or plain SFT cross-entropy for generation
    loss.backward(); optimizer.step(); optimizer.zero_grad()

deploy: fold  W_K *= l_k (column-wise on outputs), W_V *= l_v, W2 *= l_ff (row-wise on inputs)
or keep vectors per task and multiply activations per example (mixed-task batches)

The T-Few recipe used T0-3B, 1,000 steps at batch size 8 (so 8,000 example passes, which for 20 to 70 training examples is 100 or more epochs), Adafactor at learning rate 3e-3 with a 60-step warmup and linear decay, and prompt templates from P3. It also pre-trained the IA3 vectors themselves for 100,000 steps on T0's multitask mixture before few-shot fine-tuning, which lifted average accuracy from 64.6% to 65.8%. That step generalises: if you will train many IA3 task vectors on one base, first train a shared vector on a broad mixture and start every task from it.

Data needs are small by design. The method was validated on few-shot datasets (tens of examples per task) and is most useful from a few dozen to a few thousand examples. Larger datasets usually justify LoRA's greater capacity.

Hyperparameters that matter#

Tune the learning rate first. If loss falls quickly and then plateaus well above a LoRA baseline, you are capacity-limited; adding modules (queries, up-projections) helps a little, but the real fix is a method with directional capacity.

Compute and memory cost#

Memory follows the frozen-base pattern: $M\approx 2P+16P_{\text{IA3}}+M_{\text{act}}$, where $P_{\text{IA3}}$ is so small (0.5M for 8B) that its optimiser state is a few megabytes. The backward pass must still flow through every frozen layer, so activation memory and backward compute match LoRA and prompt tuning, roughly $4PD$ FLOPs for $D$ tokens, against $6PD$ for full fine-tuning. The forward overhead is three element-wise multiplies per block, which is negligible. After folding, inference is exactly the base model's cost. A task checkpoint for an 8B model is about 1 MB in bf16, against 84 MB for rank-16 LoRA on all linear layers.

The cost comparison that motivated the paper is with in-context learning: T-Few's inference cost on T0-3B was 1.1e12 FLOPs per example, while 175B in-context learning with demonstrations cost 1.4e15. Fine-tuning the T-Few model cost 2.7e16 FLOPs, about as much as in-context learning with GPT-3 175B on 20 examples.

Failure modes and how to detect them#

Capacity ceiling

Training loss stops falling at a level a rank-4 LoRA beats easily. The task needs channels to mix, not just change volume. The toy below shows the extreme case: on a random rank-1 target, the best IA3 fit captured 0.13% of it. Compare against a small LoRA early.

Channels driven to zero

Some entries of $l$ collapse toward 0 (or go negative), silencing heads or neurons; general capability on unrelated prompts degrades. Monitor the distribution of $l$: mean, minimum, and the fraction outside $[0.5,2]$. Remove weight decay toward 0, lower the learning rate, or clamp.

Overfitting tiny datasets

With 30 examples and 1,000 steps, the model sees each example dozens of times. Validation accuracy peaks and then declines, or varies strongly between prompt templates. Use held-out templates and early stopping, and report the median over templates as the paper did.

Wrong module mapping

Applying the feed-forward scale to the output of the down-projection instead of its input scales the residual stream rather than the MLP's neurons, a different and much weaker parameterisation. Check which side of each matrix your library scales (PEFT has an explicit feedforward_modules list for this).

When to use IA3, and what to use instead#

Minimal implementation#

IA3 for a Hugging Face Llama-family model in plain PyTorch, then the PEFT equivalent.

import torch, torch.nn as nn

class IA3Out(nn.Module):           # scales the OUTPUT channels of a linear layer (keys, values)
    def __init__(self, base: nn.Linear):
        super().__init__(); self.base = base
        self.l = nn.Parameter(torch.ones(base.out_features))
    def forward(self, x):
        return self.base(x) * self.l.to(x.dtype)
    @torch.no_grad()
    def fold(self):                 # W' = diag(l) W  (rows of the [out, in] weight)
        self.base.weight.mul_(self.l[:, None].to(self.base.weight.dtype)); return self.base

class IA3In(nn.Module):            # scales the INPUT channels (MLP hidden units before down_proj)
    def __init__(self, base: nn.Linear):
        super().__init__(); self.base = base
        self.l = nn.Parameter(torch.ones(base.in_features))
    def forward(self, x):
        return self.base(x * self.l.to(x.dtype))
    @torch.no_grad()
    def fold(self):                 # W' = W diag(l)  (columns)
        self.base.weight.mul_(self.l[None, :].to(self.base.weight.dtype)); return self.base

def add_ia3(model):
    for p in model.parameters():
        p.requires_grad_(False)
    for layer in model.model.layers:
        layer.self_attn.k_proj = IA3Out(layer.self_attn.k_proj)
        layer.self_attn.v_proj = IA3Out(layer.self_attn.v_proj)
        layer.mlp.down_proj = IA3In(layer.mlp.down_proj)
    return [m.l for m in model.modules() if isinstance(m, (IA3Out, IA3In))]

# params = add_ia3(model); opt = torch.optim.AdamW(params, lr=3e-3, weight_decay=0.0)
# with PEFT:
# from peft import IA3Config, get_peft_model
# cfg = IA3Config(task_type="CAUSAL_LM", target_modules=["k_proj", "v_proj", "down_proj"],
#                 feedforward_modules=["down_proj"])
# model = get_peft_model(model, cfg); model.print_trainable_parameters()

Keep the vectors in fp32 during training; they move by small relative amounts around 1, which bf16 represents with only about three significant digits. In POFE, IA3 is available as a PEFT option on SFT recipes for the narrow, high-cardinality case, and the recipe reports the scale distribution alongside the loss.

Further reading#

Liu et al. (2022) introduces IA3 and T-Few; Sections 3.2 and 3.3 cover the losses and the method, and Table 1 the cost comparison with in-context learning. Hu et al. (2021) is the LoRA baseline. Lester et al. (2021) and Li and Liang (2021) are the prompt and prefix methods IA3 was compared against. He et al. (2021) places these methods in one framework, and Liu et al. (2024, DoRA) is the closest modern relative with a per-channel magnitude.

Frequently asked questions#

References#

1. Haokun Liu, Derek Tam, Mohammed Muqeeth, Jay Mohta, Tenghao Huang, Mohit Bansal, Colin Raffel (2022). Few-Shot Parameter-Efficient Fine-Tuning is Better and Cheaper than In-Context Learning. arxiv.org/abs/2205.05638
2. Edward J. Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, Weizhu Chen (2021). LoRA: Low-Rank Adaptation of Large Language Models. arxiv.org/abs/2106.09685
3. Brian Lester, Rami Al-Rfou, Noah Constant (2021). The Power of Scale for Parameter-Efficient Prompt Tuning. arxiv.org/abs/2104.08691
4. Xiang Lisa Li, Percy Liang (2021). Prefix-Tuning: Optimizing Continuous Prompts for Generation. arxiv.org/abs/2101.00190
5. Junxian He, Chunting Zhou, Xuezhe Ma, Taylor Berg-Kirkpatrick, Graham Neubig (2021). Towards a Unified View of Parameter-Efficient Transfer Learning. arxiv.org/abs/2110.04366
6. Shih-Yang Liu, Chien-Yi Wang, Hongxu Yin, Pavlo Molchanov, Yu-Chiang Frank Wang, Kwang-Ting Cheng, Min-Hung Chen (2024). DoRA: Weight-Decomposed Low-Rank Adaptation. arxiv.org/abs/2402.09353

Related reading