Technical Summary: SSA: Improving Performance With a Better Scoring Function

Problem Statement

While Transformer models demonstrate strong in-context learning (ICL) capabilities, they exhibit systematic failures when generalizing to distribution shifts—scenarios where test-time inputs or algorithmic structures differ from the training distribution. Recent literature has identified that Large Language Models (LLMs) struggle in "low-probability settings," failing to generalize even when the underlying rules remain unchanged.

The authors investigate the architectural origins of these failures using controlled ICL tasks. They identify the Softmax scoring function within the attention mechanism as a primary contributor to these generalization failures. Specifically, they observe that when input sequences contain values with significantly larger magnitudes than those seen during training, Softmax causes attention saturation. This phenomenon forces the attention weights to collapse almost entirely onto the single token with the maximum value, effectively ignoring all other context tokens crucial for the task. This "attention collapse" prevents the model from integrating information across the sequence, leading to systematic errors in tasks requiring logical aggregation or functional inference.

Methodology

Controlled Experimental Setup

To isolate architectural effects from pre-training artifacts, the authors trained Transformer models from scratch on two synthetic ICL tasks:

1. Quantification Task: Evaluating the model's ability to interpret logical quantifiers ("Every" and "Some") over sequences of numbers.
2. Linear Function Task: Evaluating the model's ability to infer an affine function $f(x) = ax + b$ from in-context examples.

The models were trained on data drawn from a standard normal distribution $N(0, 1)$ and tested under distribution shifts, including longer sequence lengths and inputs drawn from distributions with higher variance ($N(0, \sigma)$ where $\sigma \ge 1$).

The Proposed Solution: Scaled Signed Averaging (SSA)

The authors propose Scaled Signed Averaging (SSA) as a trainable alternative to Softmax. SSA replaces the exponential scoring function with a parameterized function applied element-wise:
$$x \mapsto (1 + b|x|)^{\text{sgn}(x)n}$$
where $b > 0$ and $n \ge 1$ are trainable scalar parameters.

Key Theoretical Properties:

• Polynomial vs. Exponential Growth: Unlike Softmax, which grows exponentially and rapidly saturates to a "Hardmax" distribution when logit differences increase, SSA grows polynomially.
• Theoretical Guarantees: The authors prove via Theorem 1 and Theorem 2 that under global logit magnification or inner disproportion (where one token's score vastly exceeds others), Softmax weights converge to 1 for the winner and 0 for others exponentially fast. In contrast, SSA converges only polynomially, preserving a non-zero "contextual mass" for secondary tokens. This allows the model to maintain a balanced view of the context even in the presence of extreme values.

Evaluation Scope

The study evaluates SSA across three distinct settings:

1. Controlled ICL Tasks: Small Transformers (12 layers, 8 heads) trained from scratch.
2. Decoder-Only Benchmarks: A 114M parameter Nemotron-style model trained from scratch on the FineWeb corpus.
3. Encoder-Only Architectures: Variants of BabyBERTa trained on child-directed speech (AO-CHILDES corpus) and evaluated on a grammatical probing suite.

The authors also systematically tested existing alternatives to Softmax, including Sparsemax, Entmax, Temperature-scaled Softmax, CosFormer, and SA-Softmax, finding none to provide consistent improvements over standard Softmax on their specific tasks.

Key Results

1. Controlled ICL Tasks

• Generalization under Shift: SSA significantly outperforms Softmax in both the quantification and linear function tasks when tested on out-of-distribution (OOD) data.
• Robustness to Magnitude: In the quantification task, Softmax models fail when a single large value appears in the sequence, predicting based solely on that value. SSA models correctly aggregate information across the entire sequence, maintaining low error rates even with extreme input values.
• Linear Function Inference: SSA models achieve substantially lower Mean Squared Error (MSE) when inferring functions with coefficients drawn from shifted distributions compared to Softmax baselines.

2. Decoder-Only Benchmarks

• NLP Performance: A 114M parameter decoder-only model trained with SSA achieved lower perplexity than its Softmax counterpart on both in-distribution (FineWeb) and out-of-distribution (Wikipedia) text.
• Benchmark Gains: On standard NLP benchmarks (e.g., ARC, BoolQ, HellaSwag), the SSA-based model consistently outperformed the Softmax baseline across zero-shot evaluations, with notable improvements on tasks like arc_easy and boolq.

3. Encoder-Only Models

• Grammatical Probing: SSA-based BabyBERTa models showed improved performance on multiple grammatical probing tasks (e.g., subject-verb agreement, argument structure) compared to Softmax-based models, demonstrating that the benefits extend to masked language modeling architectures.

Significance and Claims

The paper claims that Softmax saturation is a fundamental architectural limitation that causes attention mechanisms to collapse onto single tokens when faced with large-magnitude inputs, regardless of whether those tokens are relevant to the task. This failure mode is identified as a root cause of poor generalization in ICL under distribution shifts.

The significance of the work lies in:

1. Identifying a Concrete Mechanism: Pinpointing the Softmax scoring function, rather than model scale or data complexity, as a primary driver of generalization failure in specific ICL scenarios.
2. A Simple, Effective Alternative: Introducing SSA as a "drop-in" replacement for Softmax that requires only two additional scalar parameters per attention head.
3. Broad Applicability: Demonstrating that SSA improves performance across both decoder-only and encoder-only architectures, as well as on synthetic and real-world NLP tasks.

The authors maintain a modest stance, acknowledging that while SSA mitigates failures caused by Softmax inflexibility, it does not resolve all generalization challenges, particularly in scenarios involving extreme distribution shifts where both input and function distributions differ significantly from training. They conclude that the underlying mathematical structure of attention, which conflates token magnitude with relevance, remains an open area for future research.