IA3 › Understanding Learned Rescaling Vectors
Language AI HandbookAll chapters IA3 › Understanding Learned Rescaling VectorsOne section · Sections in this chapter3 / 11 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
Why does rescaling activations work for task adaptation? The answer requires understanding what pretrained models have learned and what information is encoded in the dimensions of key, value, and feed-forward representations.
During pretraining, a language model develops representations that encode a large variety of information simultaneously. Individual dimensions of a representation do not correspond to neat, interpretable concepts in general. However, research on mechanistic interpretability has shown that groups of dimensions consistently activate for related inputs, and that the model uses these distributed representations in organized ways. Some dimensions are more sensitive to syntactic structure; others to semantic content; still others to discourse-level phenomena like topic or sentiment. These patterns emerge from the pretraining objective, where the model must learn to predict tokens from context across billions of diverse text examples.
When you fine-tune for a specific downstream task, you often do not need the model to learn entirely new patterns. If you are fine-tuning for sentiment analysis, the model already understands positive and negative words, understands intensifiers like "very" and "absolutely," understands negation. What you need is for the model to prioritize these existing capabilities in its attention and processing. IA3's rescaling vectors learn exactly this prioritization. Think of it as a spotlight: the pretrained model has developed a wide-angle view of language, and IA3 learns to focus that view on the dimensions most needed to your task.
Consider the attention mechanism more concretely. When a token queries the context, it computes dot products with every key in the sequence. These dot products determine attention weights. By rescaling keys, IA3 can change which tokens appear most "relevant" to any given query without changing what the tokens are querying for. Suppose a dimension of the key representation reliably encodes "this is a sentiment-bearing word." Before IA3 adaptation, this dimension contributes equally to attention scores alongside dozens of other dimensions. After IA3 adaptation on a sentiment task, this dimension might be amplified by a factor of 3, making sentiment-bearing words more salient to every query in the sequence. The pretrained model already had this information; IA3 simply turned up the volume on it.
The Role of Feature Specialization
Research on large language models has found that the dimensions of their internal representations exhibit a property sometimes called "feature specialization." This means that certain groups of dimensions, while not individually interpretable, collectively encode specific types of information. When you apply IA3, you are learning to modulate these feature groups. A rescaling value far above 1 on a set of dimensions increases the contribution of whatever those dimensions collectively stand for to the model's downstream computations. A rescaling value near 0 effectively silences those dimensions, as if they were never computed at all.
This view helps explain why IA3 works well in few-shot scenarios. With only a handful of training examples, a model needs to make inferences about which existing capabilities are relevant, not learn new capabilities from scratch. If you have 50 examples of medical question answering, IA3 can learn which dimensions of the pretrained representations are most associated with accurate medical reasoning. But IA3 cannot learn new medical facts from 50 examples; it can only make the medical knowledge already present in the pretrained model easier to access. This is a feature in few-shot settings, where the risk of overfitting to scarce data is high, and a limitation when the downstream task requires knowledge absent from the pretrained model.
import numpy as np
# Simulate learned IA3 vectors (in practice these would come from training)
np.random.seed(42)
# Simulate l_k: some dimensions amplified, others suppressed
l_k_learned = np.ones(128)
# Amplify certain "sentiment-related" dimensions
l_k_learned[10:20] = np.random.uniform(1.5, 2.5, 10)
# Suppress "syntactic" dimensions
l_k_learned[50:70] = np.random.uniform(0.3, 0.7, 20)
# Add some noise to others
l_k_learned += np.random.normal(0, 0.1, 128)
# Simulate l_ff: stronger adaptation in feed-forward
l_ff_learned = np.ones(512)
l_ff_learned[100:150] = np.random.uniform(2.0, 3.0, 50) # Strongly amplified
l_ff_learned[300:400] = np.random.uniform(0.1, 0.5, 100) # Strongly suppressed
l_ff_learned += np.random.normal(0, 0.15, 512)
The histograms show how learned rescaling vectors deviate from their initial uniform value of 1.0 after task-specific training. The key rescaling vectors tend to be modestly perturbed: most dimensions remain near their initial values, and only a small subset is substantially amplified or suppressed. The feed-forward rescaling vectors often show more large deviations. This pattern is consistent with the hypothesis that feed-forward networks in transformers encode factual and semantic knowledge that may need substantial emphasis or de-emphasis for specific tasks, while attention keys primarily need fine-grained relevance adjustments.
This dimension-wise view of the feed-forward rescaling reveals that IA3 learns structured patterns rather than random noise. Consecutive dimensions often share the same sign of deviation, suggesting the model discovers real groups of features to amplify or suppress together. This structure is not enforced by the IA3 formulation; it emerges organically from gradient descent on the task loss. The emergence of such structure is additional evidence that the pretrained model's internal representations have organized geometry that IA3 learns to exploit.