IA3 › Implementation
Language AI HandbookAll chapters IA3 › ImplementationOne section · Sections in this chapter6 / 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
Let's implement IA3 from scratch to see exactly how the rescaling mechanism integrates into standard transformer components. We will build incrementally, starting with the attention module and then combining everything into a complete transformer block.
The implementation has three natural layers of complexity. First, the attention module, which must apply rescaling to the key and value projections at the right point in the computation. Second, the feed-forward module, which must apply rescaling to the intermediate activations after the first linear layer and activation function. Third, the complete transformer block, which orchestrates both components and exposes utilities for freezing the pretrained parameters and accessing only the IA3 vectors.
import math
import torch
import torch.nn as nn
import torch.nn.functional as F
class IA3Attention(nn.Module):
"""Multi-head attention with IA3 rescaling on keys and values."""
def __init__(self, d_model: int, num_heads: int):
super().__init__()
self.d_model = d_model
self.num_heads = num_heads
self.head_dim = d_model // num_heads
# Standard attention projections (frozen during IA3 training)
self.W_q = nn.Linear(d_model, d_model, bias=False)
self.W_k = nn.Linear(d_model, d_model, bias=False)
self.W_v = nn.Linear(d_model, d_model, bias=False)
self.W_o = nn.Linear(d_model, d_model, bias=False)
# IA3 rescaling vectors (these are the trainable parameters)
# Initialize to ones so model starts with original behavior
self.l_k = nn.Parameter(torch.ones(d_model))
self.l_v = nn.Parameter(torch.ones(d_model))
def forward(self, x: torch.Tensor, mask: torch.Tensor = None):
batch_size, seq_len, _ = x.shape
# Compute Q, K, V projections
Q = self.W_q(x)
K = self.W_k(x)
V = self.W_v(x)
# Apply IA3 rescaling to keys and values
# Rescaling is applied before reshaping to multi-head format
K = (
self.l_k * K
) # Element-wise multiplication, broadcast over batch and seq
V = self.l_v * V
# Reshape for multi-head attention
Q = Q.view(
batch_size, seq_len, self.num_heads, self.head_dim
).transpose(1, 2)
K = K.view(
batch_size, seq_len, self.num_heads, self.head_dim
).transpose(1, 2)
V = V.view(
batch_size, seq_len, self.num_heads, self.head_dim
).transpose(1, 2)
# Scaled dot-product attention
scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.head_dim)
if mask is not None:
scores = scores.masked_fill(mask == 0, float("-inf"))
attn_weights = F.softmax(scores, dim=-1)
attn_output = torch.matmul(attn_weights, V)
# Reshape and project output
attn_output = (
attn_output.transpose(1, 2)
.contiguous()
.view(batch_size, seq_len, self.d_model)
)
return self.W_o(attn_output)The necessary lines in this implementation are K = self.l_k * K and V = self.l_v * V. These two operations are the entirety of what distinguishes IA3-adapted attention from standard attention. PyTorch's broadcasting handles the shape mismatch automatically: self.l_k has shape (d_model,) and K has shape (batch_size, seq_len, d_model). The multiplication broadcasts the rescaling vector across both the batch and sequence dimensions, applying the same per-dimension scaling to every token's key vector in the batch.
An important implementation detail: the rescaling is applied before the multi-head reshape, which means each IA3 vector controls all heads uniformly. The -dimensional rescaling vector is shared across all attention heads. Dimension of rescales the -th dimension of the full key representation, which may belong to the -th attention head. This means IA3 cannot apply different rescaling to different heads, which is a design simplification but also a constraint on expressivity compared to a head-specific rescaling scheme.
class IA3FeedForward(nn.Module):
"""Feed-forward network with IA3 rescaling on intermediate activations."""
def __init__(self, d_model: int, d_ff: int):
super().__init__()
# Standard FFN layers (frozen during IA3 training)
self.W_1 = nn.Linear(d_model, d_ff, bias=False)
self.W_2 = nn.Linear(d_ff, d_model, bias=False)
# IA3 rescaling vector for intermediate activations
self.l_ff = nn.Parameter(torch.ones(d_ff))
def forward(self, x: torch.Tensor):
# First projection and activation
h = F.gelu(self.W_1(x))
# IA3 rescaling of intermediate activations
h = self.l_ff * h
# Second projection
return self.W_2(h)The feed-forward implementation mirrors the same pattern. The rescaling h = self.l_ff * h occurs between the two linear layers, after the nonlinearity has been applied. This placement is deliberate: by rescaling after the activation function, IA3 operates on the already-processed intermediate representation where the nonlinear structure of the model's learned features is already in place. The scalar values in self.l_ff individually control how much each of the feed-forward network's intermediate "neurons" contributes to the final output projection.
class IA3TransformerBlock(nn.Module):
"""Complete transformer block with IA3 adaptation."""
def __init__(self, d_model: int, num_heads: int, d_ff: int):
super().__init__()
self.attention = IA3Attention(d_model, num_heads)
self.ffn = IA3FeedForward(d_model, d_ff)
self.norm1 = nn.LayerNorm(d_model)
self.norm2 = nn.LayerNorm(d_model)
def forward(self, x: torch.Tensor, mask: torch.Tensor = None):
# Pre-norm architecture
x = x + self.attention(self.norm1(x), mask)
x = x + self.ffn(self.norm2(x))
return x
def ia3_parameters(self):
"""Return only the IA3 rescaling parameters."""
return [self.attention.l_k, self.attention.l_v, self.ffn.l_ff]
def freeze_pretrained(self):
"""Freeze all parameters except IA3 vectors."""
for name, param in self.named_parameters():
if "l_k" not in name and "l_v" not in name and "l_ff" not in name:
param.requires_grad = FalseThe ia3_parameters method gives clean access to just the three trainable rescaling vectors. In practice, you would pass these to your optimizer rather than the full parameter set. The freeze_pretrained method implements the important training setup where all original weights are frozen and only the IA3 vectors receive gradients. This simple name-based check, looking for l_k, l_v, or l_ff in the parameter name, is straightforward and reliable for this implementation, though in production you would typically use a more systematic approach to track which parameters belong to the adapter versus the base model.
# Create a model and examine IA3 parameters
d_model = 512
num_heads = 8
d_ff = 2048
block = IA3TransformerBlock(d_model, num_heads, d_ff)
block.freeze_pretrained()
# Count trainable vs frozen parameters
trainable = sum(p.numel() for p in block.parameters() if p.requires_grad)
frozen = sum(p.numel() for p in block.parameters() if not p.requires_grad)IA3 Transformer Block Parameter Analysis ============================================= Hidden dimension: 512 Feed-forward dimension: 2048 Trainable (IA3 vectors): 3,072 Frozen (pretrained): 3,147,776 Trainable percentage: 0.0975% IA3 vector breakdown: l_k (key rescaling): 512 params l_v (value rescaling): 512 params l_ff (FFN rescaling): 2048 params
With a 512-dimensional model and 2048-dimensional feed-forward, IA3 adds just 3,072 trainable parameters to a block containing over 2 million frozen parameters. That is less than 0.15% of the block's total parameter count. Training becomes sharply cheaper: gradients are only computed for 3,072 values, not 2 million.
Setting Up IA3 Training
In practice, training an IA3-adapted model requires three setup steps beyond the standard fine-tuning workflow: freezing the base model parameters, creating an optimizer that only updates the IA3 vectors, and loading/saving only the adapter vectors for deployment.
# Demonstrate training setup for IA3
# Step 1: Freeze base model
block_train = IA3TransformerBlock(d_model=512, num_heads=8, d_ff=2048)
block_train.freeze_pretrained()
# Step 2: Create optimizer with only IA3 parameters
ia3_params_list = block_train.ia3_parameters()
optimizer = torch.optim.AdamW(ia3_params_list, lr=3e-3)
# Step 3: Verify gradient flow
x = torch.randn(2, 10, 512) # batch=2, seq=10, d_model=512
output = block_train(x)
loss = output.sum()
loss.backward()
# Check which params got gradients
params_with_grad = [
name for name, p in block_train.named_parameters() if p.grad is not None
]
params_without_grad = [
name for name, p in block_train.named_parameters() if p.grad is None
]Gradient Flow Verification: Parameters with gradients: 3 Parameters without gradients: 10 IA3 parameters receiving gradients: attention.l_k: shape (512,), grad norm 3.4190 attention.l_v: shape (512,), grad norm 29.8552 ffn.l_ff: shape (2048,), grad norm 36.2504
The gradient flow verification confirms that only the three IA3 vectors receive gradients. All other parameters, representing the pretrained transformer block, have grad = None. This is exactly what we want: the optimizer only updates the tiny set of rescaling vectors, leaving the pretrained knowledge intact.