🤖 machine learning

Shortcut Solutions Learned by Transformers Impair Continual Compositional Reasoning

This paper investigates continual compositional reasoning in Transformers using an extended LEGO framework, revealing that while feedforward BERT models fail due to entrenched shortcut solutions, recurrent ALBERT models exhibit superior performance by learning algorithmic "for-loop" strategies that can be further enhanced through cross-experience training.

Original authors: William T. Redman, Erik C. Johnson, Brian Robinson

Published 2026-05-08

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: William T. Redman, Erik C. Johnson, Brian Robinson

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are teaching a robot to solve a series of puzzles. The puzzles aren't just random; they are variations of the same underlying logic. For example, one puzzle might ask you to rotate a triangle, the next to flip it, and the next to do both. The goal is to see if the robot can learn the logic of the first puzzle and use it to instantly solve the second and third, without forgetting how to do the first one.

This paper investigates how two popular types of AI "brains" (called BERT and ALBERT) handle this kind of learning. The researchers found that while these AI models are incredibly smart, they often take "mental shortcuts" that make them bad at learning new things later on.

Here is the breakdown of their findings using simple analogies:

1. The Two Types of Learners

The researchers tested two models:

BERT: Think of this as a photographer. It looks at a whole picture at once and tries to memorize the specific details of that single image. It's very fast and good at recognizing patterns in the moment, but it doesn't really understand the process of how the image was made.
ALBERT: Think of this as a programmer. Instead of just memorizing the picture, it learns a set of instructions (like a loop in a computer code) that can be repeated to solve the problem. It understands the mechanism behind the puzzle.

2. The "Shortcut" Trap

When the models were taught the first puzzle (let's call it "Puzzle A"), both did well. However, when the researchers introduced "Puzzle B" (which was very similar but slightly different), a problem arose.

BERT (The Photographer): It learned a shortcut. It realized, "Hey, if I just look at the very first clue in the puzzle, I can guess the answer for this specific type of puzzle." It didn't learn the actual logic; it just memorized a trick that worked for Puzzle A.
- The Result: When faced with Puzzle B, BERT got confused. Because it was relying on a trick specific to Puzzle A, it couldn't adapt. It also completely forgot how to do Puzzle A once it started learning Puzzle B. This is called catastrophic forgetting.
ALBERT (The Programmer): It learned the algorithm. It figured out, "Oh, I need to take the current state, apply a rule, and move to the next step." This is like learning a "For Loop" (a computer instruction that repeats an action).
- The Result: Because ALBERT learned the method rather than a specific trick, it could apply that method to Puzzle B much faster. It showed forward transfer, meaning it used what it learned in Puzzle A to speed up learning Puzzle B.

3. The Size Matters (But Not How You Think)

The researchers tried making the models bigger (adding more layers, like adding more rooms to a house).

For ALBERT: Bigger was better. More "rooms" meant it could refine its algorithm and get even better at transferring knowledge.
For BERT: Bigger actually made things worse or unstable. Because BERT relies on shortcuts, making the model more complex just gave it more ways to get stuck in a bad habit. It couldn't generalize its learning to new puzzles.

4. The "Memory Replay" Fix

Both models suffered from catastrophic forgetting—when they learned the new puzzle, they completely erased the memory of the old one.

To fix this, the researchers used a Replay Buffer. Imagine a student who, while studying for a new math test, keeps a few flashcards from the old test on their desk to glance at occasionally.

The Result: This worked wonders. By showing the models a tiny bit (just 1%) of the old data while they learned the new data, they stopped forgetting. Both BERT and ALBERT could remember the old puzzles while learning the new ones.

5. The Final Hurdle: Combining the Puzzles

The ultimate test of "compositional reasoning" is whether the AI can mix and match rules from different puzzles to solve a brand-new, complex puzzle that combines them all.

The Problem: Even with the "Replay Buffer" helping them remember, both models struggled to combine the rules. They could remember Puzzle A and Puzzle B separately, but they couldn't easily weave them together to solve a hybrid puzzle.
The Difference: The researchers found a way to help ALBERT succeed at this. If they taught ALBERT using a strategy where they slowly stitched the old puzzles and new puzzles together during training (rather than keeping them separate), ALBERT could learn to combine them.
The Limit: This "stitching" strategy did not work for BERT. Because BERT had already entrenched itself in its "shortcut" way of thinking during the initial training, it couldn't be taught to combine the rules later. It was too set in its ways.

The Bottom Line

The paper concludes that while AI models like BERT and ALBERT are powerful, they have a hidden flaw: they tend to learn "cheats" (shortcuts) instead of deep logic.

ALBERT is better at learning the deep logic (the "For Loop"), which helps it learn new, related tasks faster.
BERT gets stuck on surface-level tricks, which makes it bad at adapting to new situations and causes it to forget old ones.

The study suggests that to make AI truly capable of continual learning (learning forever without forgetting), we need to move away from models that rely on shortcuts and toward architectures that are built to understand the underlying "algorithms" of the world.

Technical Summary: Shortcut Solutions Learned by Transformers Impair Continual Compositional Reasoning

Problem Statement

The paper addresses a critical gap in understanding how Transformer neural networks perform Continual Learning (CL) within the context of compositional reasoning. While Transformers have demonstrated success in identifying input-specific relations and performing compositional tasks on static datasets, their ability to leverage learned representations to acquire new, related experiences without catastrophic forgetting remains understudied. Specifically, the authors investigate whether these models can develop general computational strategies (analogous to human analogical reasoning) or if they rely on non-generalizable "shortcut" solutions that hinder their ability to transfer knowledge across sequential tasks.

Methodology

The authors extend the synthetic Learning Equality and Group Operations (LEGO) framework (Zhang et al., 2022) into a Continual LEGO setting.

Task Design: The LEGO task involves solving a sequence of group operations ( $a_t = a_{t-1} * x_t$ ) within a group structure (specifically the symmetry group $D_3$ ). The task is viewed as directed graph traversal.
Continual Setup: The $D_3$ group is decomposed into three distinct "flip-flop experiences," each defined by a specific subgroup (e.g., $\{x_1, x_4\}$ , $\{x_1, x_5\}$ , $\{x_1, x_6\}$ ). Models are trained sequentially on these experiences. Crucially, once a model moves to a new experience, it receives no inputs from previous experiences, testing for forward transfer (learning new tasks faster) and performance maintenance (avoiding catastrophic forgetting).
Models: The study compares BERT (a canonical feedforward Transformer) and ALBERT (a recurrent variant with weight sharing across layers).
Architectural Analysis: The authors systematically vary model hyperparameters (number of hidden layers and attention heads) to identify minimal architectures capable of solving the task. They analyze self-attention patterns to determine the mechanistic nature of the learned solutions (e.g., algorithmic loops vs. shortcuts).
Interventions: The study evaluates the efficacy of replay buffers (saving 1% of past data) and incremental training (stitching past and current experience data together) to mitigate forgetting and enable cross-experience composition.

Key Contributions

Continual LEGO Framework: The authors introduce a controlled, synthetic paradigm for systematically investigating CL capabilities in Transformers, specifically focusing on compositional reasoning.
Architectural Divergence: The study reveals that architectural choices (feedforward vs. recurrent/weight-shared) differentially impact generalization and forward transfer. BERT models show inconsistent performance as size increases, while ALBERT models demonstrate a clearer relationship between capacity and performance.
Mechanistic Explanation of Failure/Success: Through attention pattern analysis, the authors identify that ALBERT learns a generalizable, algorithmic For loop-esque solution, whereas BERT learns shortcut solutions that fail to generalize to longer sequences or new experiences.
Limits of Replay: The paper demonstrates that while replay buffers can mitigate catastrophic forgetting, they are insufficient for enabling compositional reasoning across experiences in BERT models. Only ALBERT models benefit from training on incrementally combined data to achieve true cross-experience composition.

Results

1. Performance on Continual LEGO

Forward Transfer: ALBERT models exhibit strong forward transfer, achieving high accuracy on new flip-flop experiences with fewer training epochs. BERT models show weak forward transfer, often requiring more epochs to reach asymptotic accuracy and failing to generalize to sequence lengths longer than the training set.
Catastrophic Forgetting: Both BERT and ALBERT models exhibit near-complete catastrophic forgetting (performance drops to ~0%) on previous experiences when trained on new ones, without intervention.
Generalization: ALBERT models generalize well to longer sequences (e.g., $a_5, a_6$ ) when trained on shorter sequences ( $a_4$ ). BERT models generally fail to generalize beyond the training sequence length.

2. Mechanistic Analysis (Attention Patterns)

ALBERT (General Solution): Attention patterns in minimal ALBERT models show a "For loop-esque" signature. Tokens in a given clause predominantly attend to tokens in the preceding clause, systematically updating the state. This pattern is maintained and strengthened across experiences.
BERT (Shortcut Solution): BERT models exhibit a shortcut signature where self-attention isolates a single shared input token in the first clause for all output tokens. This solution is entrenched early in training and persists even when new experiences are introduced, preventing the model from learning the necessary iterative composition.

3. Impact of Architectural Scale

ALBERT: Performance (task accuracy, generalization, forward transfer) scales positively with the number of hidden layers (requiring at least 4 layers for success).
BERT: Increasing model size (e.g., to 12 layers) does not guarantee better performance and can lead to training instability. No clear relationship exists between model size and forward transfer or generalization for BERT.

4. Mitigation Strategies

Replay Buffers: Saving just 1% of past data successfully mitigates catastrophic forgetting for both models while preserving forward transfer. However, naive replay (saving separate examples) fails to enable composition across experiences.
Incremental/Combined Training: Training on data that stitches past and current experiences together ("incremental flip-flop") rescues performance for ALBERT, allowing it to approach the performance of a model trained on the full task. This strategy fails for BERT, where the shortcut solution becomes entrenched during initial training and cannot be unlearned or corrected by combined data.

Significance and Claims

The paper claims that the inductive bias of recurrent Transformers (specifically ALBERT's weight sharing) is better suited for continual compositional reasoning than standard feedforward architectures (BERT).

Shortcut Solutions as a Barrier: The authors argue that the "shortcut solutions" learned by BERT are not merely a performance issue but a fundamental barrier to continual learning in compositional tasks. These shortcuts become entrenched and prevent the model from acquiring the general algorithmic strategies required for forward transfer and cross-experience composition.
Recurrent Advantage: The "For loop-esque" computation learned by ALBERT provides a mechanism for generalization and forward transfer, suggesting that recurrent structures (or weight-sharing mechanisms) may be crucial for developing flexible computational strategies in neural networks.
Training Strategy Limits: The study highlights that while simple replay can fix forgetting, it does not solve the problem of composition across experiences. True compositional reasoning in CL may require training strategies that explicitly combine data across experiences, a finding that suggests generative replay methods could be a fruitful direction for future research, particularly for models like ALBERT.

The authors conclude that while Transformers can succeed on complex static tasks, their ability to continually learn and compose across new experiences is severely impaired by the specific computational strategies (shortcuts vs. algorithms) they converge upon, which are heavily influenced by architectural choices.