Super Weights in LLMs and the Failure of Selective Training

Technical Summary§

This paper investigates the concept of "Super Weights"—individual parameters whose removal degrades LLM performance—and tests their trainability. The authors find that:

• Pruning Super Weights does not universally degrade all LLMs.
• Training only Super Weight coordinates (100 to 8,192 parameters) in isolation yields random-guessing accuracy on OLMo-1B and OLMo-7B. Expanding to local neighborhoods (up to 36K parameters) does not help.
• Training an equal number of random coordinates in the same down_proj layers improves over baseline, confirming the failure is specific to Super Weight positions, not sparsity.
• Vanilla LoRA (low-rank decomposition on attention weight matrices) with only 0.16% of parameters succeeds. Applying LoRA to down_proj also works.
• Constraining LoRA updates to Super Weight coordinates yields indistinguishable results from random coordinates (10-seed ablation).

Key Equations§

Let $W \in \mathbb{R}^{d \times k}$ be a weight matrix. LoRA updates it via: $$ W' = W + \Delta W, \quad \Delta W = BA, $$ where $B \in \mathbb{R}^{d \times r}$, $A \in \mathbb{R}^{r \times k}$ with rank $r \ll \min(d, k)$.

For selective training of Super Weights at coordinates $\mathcal{S}$, the update is: $$ \Delta W_{ij} = \begin{cases} \theta_{ij} & \text{if } (i,j) \in \mathcal{S} \\ 0 & \text{otherwise} \end{cases} $$ where $\theta_{ij}$ are learnable scalars.

Implementation Details§

• Models: OLMo-1B, OLMo-7B.
• Downstream task: multiple-choice QA (e.g., MMLU).
• Training: 1-5 epochs, AdamW, learning rate 1e-4.
• LoRA rank: $r=8$ (for 0.16% parameters in attention).
• Super Weights identified by weight magnitude after pretraining.

Code Snippet (PyTorch-style)§

# LoRA update for a linear layer (down_proj example)
class LoRALayer(nn.Module):
    def __init__(self, original_layer, rank=8):
        super().__init__()
        self.original = original_layer
        self.A = nn.Parameter(torch.randn(rank, original_layer.in_features) * 0.01)
        self.B = nn.Parameter(torch.zeros(original_layer.out_features, rank))
        
    def forward(self, x):
        return self.original(x) + (x @ self.A.T) @ self.B.T

# Selective training of Super Weights (only coordinates in S)
class SelectiveUpdate(nn.Module):
    def __init__(self, original_layer, super_mask):
        super().__init__()
        self.original = original_layer
        self.super_weights = nn.Parameter(original_layer.weight[super_mask].detach().clone())
        self.mask = super_mask
        
    def forward(self, x):
        # Apply update only at masked positions
        delta = torch.zeros_like(self.original.weight)
        delta[self.mask] = self.super_weights - self.original.weight[self.mask].detach()
        return nn.functional.linear(x, self.original.weight + delta, self.original.bias)

The key insight: Super Weights are not individually trainable. Structured, low-rank updates over entire matrices enable effective fine-tuning, while sparse scalar updates at important positions fail. This suggests that parameter importance does not imply trainability in isolation; gradient information is diluted when updates are confined to a few coordinates.