To See the Unseen: on the Generalization Ability of Transformers in Symbolic Reasoning

1. Problem Statement

The paper investigates the ability of decoder-only transformer models to perform abstract symbolic reasoning, specifically solving propositional logic problems via in-context learning (ICL). The core challenge is generalization to unseen tokens: models must execute reasoning algorithms (like forward-chaining) using variable names (symbols) that were never observed during training.

Previous research established that transformers fail at this task, often unable to copy or generate unseen tokens. While some works suggested architectural changes (e.g., Boix-Adserà et al., 2024) or "active forgetting" (Anand et al., 2025) as solutions, the underlying mechanistic cause of this failure remained unclear. This paper posits that the failure is not merely a lack of copying capability but is fundamentally driven by a representational collapse of the unembeddings (last-layer weights) of unseen tokens.

2. Methodology

Theoretical Analysis

The authors provide a theoretical proof regarding the behavior of softmax classifiers (and by extension, transformers with tied embeddings) trained with $\ell_2$-regularized Gradient Descent (GD) or Stochastic Gradient Descent (SGD).

• Lemma 4.1: They prove that for any two labels (tokens) $i$ and $j$ that do not appear in the training data, the distance between their weight vectors ($W_i$ and $W_j$) contracts (decreases) at every step of optimization, provided the input features have bounded norms (ensured by LayerNorm) and the model fits the data well.
• Implication: This leads to the collapse of unseen token unembeddings into nearly identical vectors, making it mathematically impossible for the model to distinguish between different unseen variables.

Empirical Investigation

• Synthetic Experiments: The authors trained small decoder-only transformers (60M–75M parameters) on synthetic propositional logic problems (Horn clauses) with Chain-of-Thought (CoT) reasoning traces.
  • Vocabulary Control: They used a vocabulary with 100 deliberately held-out (unseen) tokens.
  • Variables: They manipulated vocabulary diversity ($\|P\|$), the presence of Copy Attention heads, and the training regime for embeddings (trainable, frozen, or periodically reinitialized).
• Large-Scale Validation: They analyzed the Gemma 3 family of open-weight models (1B to 27B parameters). They examined the 99 "unused" tokens reserved in the tokenizer and measured the cosine similarity of their embeddings compared to used tokens. They also performed fine-tuning experiments on Gemma 3 1B using these unused tokens as symbols.

Proposed Solutions

The paper proposes a combination of techniques to mitigate the collapse:

1. Copy Attention: An architectural modification where a specific attention head attends directly to input embeddings ($X W_E$) rather than hidden states, facilitating the copying of tokens.
2. Data Diversity: Increasing the size of the training predicate vocabulary to ensure the model encounters a wide variety of symbolic patterns.
3. Embedding Management:
   • Freezing: Keeping embeddings fixed at initialization (where they are orthogonal).
   • Active Forgetting: Periodically resetting (reinitializing) the embeddings during training to prevent them from collapsing.

3. Key Contributions

1. Identification of Representational Collapse: The paper identifies and proves that the unembeddings of unseen tokens collapse to nearly the same vector during training. This collapse is the primary reason models cannot distinguish between different unseen variables, rendering them indistinguishable to the model.
2. Mechanistic Explanation for Existing Heuristics: The collapse phenomenon explains why "active forgetting" (periodic resetting of embeddings) works: it prevents the weights from converging to the collapsed state.
3. Resolution of Architectural Limitations: The authors demonstrate that the "copying" architecture proposed by Boix-Adserà et al. (2024) is insufficient on its own when multiple unseen variables are present. Without preventing embedding collapse, the model cannot distinguish between the multiple unseen tokens it is trying to copy.
4. Practical Solution: They propose a robust recipe for generalization: Copy Attention + High Data Diversity + (Frozen or Periodically Reset) Embeddings.
5. Real-World Evidence: They provide empirical evidence that this collapse occurs in large, real-world models (Gemma 3), specifically affecting the 99 unused tokens, which results in poor initialization for downstream fine-tuning tasks.

4. Results

Synthetic Logic Experiments

• Standard Transformers: Fail completely when query or all predicates are unseen. They either ignore the token or hallucinate a seen token.
• Copy Attention Alone: Improves performance on single unseen tokens but fails when all variables are unseen because the collapsed embeddings make the variables indistinguishable.
• Frozen Embeddings: Models with frozen (orthogonal) embeddings generalize perfectly to unseen tokens, even without copy attention, though copy attention improves accuracy further.
• Active Forgetting: Periodically resetting embeddings allows standard models to generalize, provided there is sufficient data diversity.
• Data Diversity: A threshold of diversity is required. Models with small vocabularies ($\|P\|=150$) fail to maintain the inductive bias for generalization, while larger vocabularies ($\|P\|=9600$) succeed.

Gemma 3 Analysis

• Embedding Collapse: In all Gemma 3 model sizes, the 99 unused tokens exhibit significantly higher mean cosine similarity (e.g., 0.78 in 1B, 0.24 in 27B) compared to used tokens (e.g., 0.09 in 1B). The effect is more pronounced in smaller models.
• Fine-tuning Efficiency: When fine-tuning Gemma 3 1B on logic problems:
  • Using English adjectives/letters as symbols: Reaches ~90% accuracy in <500 steps.
  • Using unused tokens as symbols: Reaches 90% accuracy in **5000 steps**.
  • This 10x slowdown is attributed to the model needing to "un-collapse" or learn to distinguish the highly correlated unused token embeddings.

5. Significance and Conclusion

This paper fundamentally shifts the understanding of why transformers fail at abstract symbolic reasoning. It moves beyond the narrative of "lack of copying capability" to a deeper optimization dynamic issue: the collapse of unseen token representations.

• Theoretical Impact: It provides a rigorous proof that $\ell_2$ regularization and LayerNorm inevitably cause unseen token weights to contract, explaining a specific failure mode in deep learning.
• Practical Impact: It offers a concrete, actionable framework for training models that can generalize to new symbols (crucial for reasoning, code generation, and scientific discovery). It suggests that simply adding more data is not enough; one must actively manage the embedding space (via freezing or resetting) and use specific architectural heads (Copy Attention) to maintain distinguishability.
• Broader Implications: The findings suggest that "unused" tokens in large language models (like Gemma) are not neutral; they are highly correlated and potentially detrimental for downstream tasks requiring novel symbol usage, necessitating careful initialization or fine-tuning strategies.

In summary, the paper argues that symbolic reasoning requires the preservation of token distinctiveness, which is naturally eroded by standard training dynamics, and proposes specific architectural and training interventions to preserve it.