Back to AI Research

AI Research

To See the Unseen: on the Generalization Ability of... | AI Research

Key Takeaways

  • To See the Unseen: on the Generalization Ability of Transformers in Symbolic Reasoning This research investigates why transformer models struggle to perform...
  • We investigate the ability of decoder-only transformer models to perform abstract symbolic reasoning; specifically solving propositional logic reasoning problems given in-context.
  • We show both theoretically and empirically that a particular representational collapse also has a crucial role: the unembeddings (last-layer weights) of unseen tokens collapse to nearly the same vector during training.
  • We support our claims with extensive controlled experiments on propositional logic reasoning problems.
  • Beyond synthetic experiments, we also observe evidence of (un)embedding collapse in the open-weight models in the Gemma 3 family, which includes 99 unused tokens reserved for downstream use.
Paper AbstractExpand

We investigate the ability of decoder-only transformer models to perform abstract symbolic reasoning; specifically solving propositional logic reasoning problems given in-context. Previous work demonstrated that models fail to generalize to problems involving variable names that were not observed during training, and it was shown that one reason behind this is the difficulty of copying (or generating) unseen tokens. We show both theoretically and empirically that a particular representational collapse also has a crucial role: the unembeddings (last-layer weights) of unseen tokens collapse to nearly the same vector during training. The collapse makes distinguishing multiple unseen variables difficult for the model (especially when the embedding and unembedding parameters are shared), and provides a mechanistic explanation for the effectiveness of existing heuristic interventions like "active forgetting", which periodically reset the token (un)embeddings. Based on these observations, we devise a combination of techniques, involving a small architecture change facilitating copying, data diversity, and freezing or resetting (un)embeddings, that achieves generalization to unseen tokens. We support our claims with extensive controlled experiments on propositional logic reasoning problems. Beyond synthetic experiments, we also observe evidence of (un)embedding collapse in the open-weight models in the Gemma 3 family, which includes 99 unused tokens reserved for downstream use. Empirically we find that the correlated embeddings of these tokens are a poor initialization for finetuning applications.

To See the Unseen: on the Generalization Ability of Transformers in Symbolic Reasoning
This research investigates why transformer models struggle to perform abstract symbolic reasoning when they encounter new variables—such as names or symbols—that were not present during their training. While these models excel at many tasks, they often fail when asked to apply logical rules to unseen tokens. The authors identify that this failure is not just a lack of training data, but a fundamental issue in how the models store and represent information, leading to a "representational collapse" that makes it impossible for the model to distinguish between different unseen variables.

The Problem of Representational Collapse

The core issue identified is that the "unembeddings"—the final layer weights that map the model's internal representations back to specific tokens—tend to collapse toward the same vector for any token the model has not seen during training. Because these weights become nearly identical, the model loses the ability to tell different unseen variables apart. This is particularly problematic in architectures where embedding and unembedding parameters are shared. The authors prove that this collapse is a natural consequence of standard training techniques like weight decay and layer normalization, which effectively push the representations of unused tokens together.

Why Existing Interventions Work

Previous research suggested that "active forgetting"—periodically resetting token embeddings during training—could improve generalization. This paper provides a mechanistic explanation for why this works: by resetting the embeddings, the model is prevented from allowing these weights to collapse into a single, indistinguishable vector. The authors’ findings confirm that this heuristic is effective because it forces the model to maintain distinct representations for tokens, rather than letting them drift toward a uniform, unusable state.

Achieving Reliable Reasoning

To solve the problem of generalization, the authors propose a multi-faceted approach. They demonstrate that models can successfully reason with unseen symbols if they combine three specific techniques:

  • Copy-enabled architecture: A small architectural change, called "copy attention," allows the model to more effectively copy input tokens directly to the output, bypassing the limitations of standard layers.

  • Data diversity: Training on a wide variety of symbolic substitutions helps the model learn the underlying structure of the reasoning task rather than relying on specific, memorized tokens.

  • Embedding management: Either freezing or periodically resetting the (un)embeddings prevents the collapse that otherwise renders unseen tokens indistinguishable.

Implications for Large Language Models

The study extends beyond synthetic experiments to examine the Gemma 3 model family, which includes 99 reserved, unused tokens. The researchers found evidence of the same (un)embedding collapse in these large, open-weight models. They conclude that because these unused tokens have highly correlated, collapsed embeddings, they serve as a poor starting point for developers looking to finetune these models for new tasks. This suggests that the way models are pre-trained can create hidden limitations that impact their future utility.

Comments (0)

No comments yet

Be the first to share your thoughts!