🧬 biology

Relational reasoning and inductive bias in transformers and large language models

This paper investigates how transformers perform transitive inference, revealing that while in-weights learning naturally develops linear embeddings for robust generalization, in-context learning typically relies on copy strategies unless pre-trained on linear tasks or explicitly prompted to construct linear mental maps, highlighting that both training regimes and representation geometry critically determine relational reasoning capabilities.

Original authors: Jesse Geerts, Andrew Liu, Stephanie Chan, Claudia Clopath, Kimberly Stachenfeld

Published 2026-05-12

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Jesse Geerts, Andrew Liu, Stephanie Chan, Claudia Clopath, Kimberly Stachenfeld

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ⚕️ This is an AI-generated explanation of a preprint that has not been peer-reviewed. It is not medical advice. Do not make health decisions based on this content. Read full disclaimer

Imagine you are teaching a robot how to understand the world. You want to know if it can figure out that if A is bigger than B, and B is bigger than C, then A must be bigger than C. This is called "transitive inference," a classic logic puzzle that humans and even monkeys solve naturally.

This paper investigates how two different types of "learning" in AI models (specifically, Transformers) handle this puzzle. The authors found that how the robot learns changes what it learns.

Here is the breakdown of their findings using simple analogies:

1. The Two Learning Styles: The "Hard Drive" vs. The "Flashcard"

The study compares two ways a model learns:

In-Weights Learning (IWL) – The "Hard Drive" Approach:
Imagine you are studying for a test by memorizing facts until they are permanently etched into your brain (or a hard drive). Once the training is done, the robot "knows" the rules.
- The Result: When the robot learned this way, it didn't just memorize the specific pairs it saw (A>B, B>C). It built a mental "number line." It understood that everything has a rank. Because it had this mental map, it could easily guess the relationship between items it had never seen together (like A and C). It acted just like a human or a monkey, getting better at the puzzle the further apart the items were on the list.
In-Context Learning (ICL) – The "Flashcard" Approach:
Imagine you are taking a test, but you are allowed to look at a cheat sheet (the context) right in front of you. You don't memorize the rules; you just look at the examples provided in the moment to solve the specific question.
- The Result: When the robot learned this way, it was very clever at copying what it saw on the cheat sheet. If the cheat sheet said "A > B," and the question was "A > B," it got it right. But if the question was "A > C" (which wasn't on the sheet), it failed. It didn't build a mental map; it just played a game of "Match and Copy." It only solved the puzzle if the answer was literally written in the prompt.

2. The "Aha!" Moment: Teaching the Robot to Think in Lines

The researchers asked: Can we trick the "Flashcard" robot into thinking like the "Hard Drive" robot?

They tried pre-training the "Flashcard" robot on a different task: Linear Regression. Think of this as teaching the robot to draw a straight line through a set of dots.

The Magic: Once the robot learned to draw straight lines (a linear relationship), it suddenly started solving the transitive inference puzzle perfectly, even when using the "Flashcard" method.
Why? Because drawing a line forces you to understand that numbers have an order (1, 2, 3...). Once the robot understood this "line" concept, it could apply it to the new puzzle. It stopped just copying and started reasoning.

3. The Real-World Test: Asking Big Language Models

Finally, the team tested this on huge, real-world AI models (Large Language Models like the ones you might chat with). They gave these models logic puzzles where the facts were either:

True to reality (e.g., "A whale is bigger than a fish").
False to reality (e.g., "A fish is bigger than a whale" – forcing the model to ignore its memory and look only at the new facts).

They gave the models a "hint" about how to visualize the problem:

Hint A (The Number Line): "Imagine these items are on a straight line from smallest to biggest."
Hint B (The Circle): "Imagine these items are on a circle."

The Finding:

When the models were told to use a straight line, they solved the logic puzzles much better, especially the tricky ones where they had to ignore their prior knowledge.
When they were told to use a circle, they struggled. Why? Because on a circle, the rules of "bigger than" break down (if you go around the circle, A can be bigger than B, B bigger than C, but C can wrap around and be bigger than A). The "circle" geometry confused the logic.

The Big Takeaway

The paper concludes that geometry matters.

If an AI model learns to store facts in its "brain" (weights), it naturally builds a straight-line map of the world, allowing it to reason logically.
If an AI model relies only on the information in front of it (context), it tends to just copy patterns unless it is specifically trained to think in "lines."
Even for giant, smart AI models, simply telling them to "imagine a straight line" helps them reason better than telling them to "imagine a circle."

In short: To make AI reason like a human, you don't just need to give it more data; you need to give it the right shape of thinking.

Technical Summary: Relational Reasoning and Inductive Bias in Transformers and Large Language Models

Problem Statement

Transformer-based models have demonstrated remarkable reasoning capabilities, yet the mechanisms underlying relational reasoning—specifically the ability to infer relationships between indirectly related items—remain poorly understood. A critical gap exists in distinguishing how models perform transitive inference (e.g., if $A > B$ and $B > C$ , then $A > C$ ) under two distinct learning regimes:

In-Weights Learning (IWL): Information is stored in the model's parameters during training.
In-Context Learning (ICL): Information is provided dynamically within the input sequence at inference time.

While ICL has been studied in classification and regression, its behavior in relational reasoning tasks is unclear. It is difficult to determine whether large language models (LLMs) rely on prior knowledge encoded in weights or are performing true in-context inference, largely due to limited control over training data in pre-trained models. This paper investigates how these learning modes differ in their inductive biases and generalization patterns using transitive inference as a prototypical task.

Methodology

The authors employed a multi-faceted approach combining controlled experiments on small, trained-from-scratch transformers with probing experiments on large language models (LLMs).

1. Transitive Inference Task Setup

The core task involves a hierarchy of items (e.g., $A > B > C > \dots > G$ ).

Input: Pairs of items (tokens) with labels indicating their relative order (+1 if first > second, -1 otherwise).
Training Data: Models are trained exclusively on adjacent pairs (symbolic distance = 1, e.g., $A>B$ , $B>C$ ).
Evaluation: Models are tested on non-adjacent pairs (symbolic distance > 1, e.g., $A>C$ ) to measure transitive generalization.
Architecture: A two-layer, attention-only transformer (omitting MLPs for interpretability), using ResNet18 embeddings of Omniglot images as input features.

2. Experimental Conditions

The study compares three specific training regimes:

IWL (Fixed Hierarchy): The hierarchy is fixed across all training sequences. Context examples are random and irrelevant, forcing the model to learn the ranking structure into its weights.
Standard ICL (Dynamic Hierarchy): A new random hierarchy is assigned for every training sequence. The context contains all adjacent pairs for that specific sequence. The model must infer the ranking from the context, preventing weight storage of specific hierarchies.
ICL with Linear Regression Pre-training: Models are first pre-trained on an in-context linear regression task (predicting $y = Wx$) where weights $W$ change per sequence. This forces the model to learn a latent linear representation via ICL before being evaluated on the transitive inference task.

3. Mechanistic Analysis

To understand how models solve the task, the authors utilized:

Principal Component Analysis (PCA): To visualize the geometry of hidden layer representations.
Attention Pattern Analysis: To identify specific circuit mechanisms, such as "induction heads" (previous-token attention + induction heads).
Ablation Studies: Systematically removing attention heads to test causal importance.

4. Large Language Model (LLM) Probing

To test generalizability, the authors used the ReCogLab dataset with pre-trained LLMs (Gemini and Gemma families). They employed a congruency paradigm:

Congruent: Relationships align with real-world knowledge (e.g., Whale > Dolphin).
Incongruent: Relationships contradict real-world knowledge (e.g., Goldfish > Dolphin).
Geometric Scaffolding: Prompts were augmented with instructions to imagine items on a linear number line (supports transitivity) versus a circle (violates transitivity due to wrap-around).

Key Results

1. Divergent Behaviors in IWL vs. Standard ICL

IWL Models: Naturally exhibit transitive inference despite training only on adjacent pairs. They display human-like behavioral effects:
- Symbolic Distance Effect: Accuracy increases as the distance between items in the hierarchy increases.
- Terminal Item Effect: Higher performance on items at the ends of the hierarchy.
- Representations: PCA reveals a continuous, linear embedding where items are arranged by signed symbolic distance.
Standard ICL Models: Fail to generalize transitively when trained only on adjacent pairs.
- They achieve high accuracy on adjacent pairs (memorization) and terminal items but perform at chance on internal non-adjacent pairs.
- Mechanism: They rely on an induction circuit (match-and-copy strategy) rather than learning relational structures. PCA shows no distance-based representations; the model simply copies the label from the context if the pair is present.
- Note: If ICL models are explicitly trained on non-adjacent pairs, they can learn transitivity, but they do not exhibit the symbolic distance effect seen in IWL.

2. The Role of Linear Pre-training

Pre-training ICL models on in-context linear regression fundamentally alters their behavior.
These models do exhibit transitive inference, symbolic distance effects, and terminal item effects on the subsequent TI task.
Mechanism: The induction circuit disappears. Instead, the model develops distributed representations that capture hierarchical relationships, similar to the IWL models. The linear regression task induces an additive representational structure (scalar ranks) where transitivity emerges mathematically ( $rank(A) > rank(B) \land rank(B) > rank(C) \implies rank(A) > rank(C)$ ).

3. Findings in Large Language Models

Congruency: LLMs perform better on congruent relational reasoning probes (leveraging stored IWL-like knowledge) than incongruent ones (requiring pure ICL).
Geometric Scaffolding: When prompted with incongruent tasks, linear geometry prompts ("Imagine a number line") significantly improve transitive inference compared to circular prompts ("Imagine a circle").
This effect is most pronounced in incongruent conditions and scales with model size (in Gemma models), suggesting that LLMs possess latent linear representations that can be activated by appropriate geometric framing.

Key Contributions

Characterization of Inductive Biases: The paper demonstrates that IWL and ICL lead to qualitatively different computational strategies. IWL naturally induces a linear, transitive bias, whereas standard ICL defaults to a match-and-copy strategy unless the training data or pre-training explicitly necessitates relational generalization.
Mechanistic Explanation: The authors identify that the lack of transitive generalization in standard ICL is due to the formation of induction circuits (pattern matching) rather than the learning of relational structures. Conversely, linear regression pre-training suppresses these circuits in favor of distributed linear representations.
Bridging Small Models and LLMs: By using geometric prompting, the study shows that the findings from controlled, small-scale experiments generalize to modern LLMs. It provides evidence that LLMs' ability to perform transitive reasoning depends on accessing appropriately structured representations, either stored in weights or activated via context.

Significance and Claims

The paper claims that the capacity for transitive inference in transformers is not a fixed capability but is critically determined by:

The Training Regime: Whether information is learned in-weights or in-context.
The Geometric Structure of Representations: Whether the induced representations are linear (supporting transitivity) or rely on pattern matching.

The authors conclude that while transformers can perform transitive inference, standard ICL on adjacent pairs does not automatically yield this ability. Instead, it requires either specific pre-training that induces linear representations (like linear regression) or the activation of latent linear structures through geometric prompting in LLMs. This work highlights that "reasoning" in LLMs may often be a reflection of the geometric structure of the representations they are forced to use or access, rather than an inherent logical deduction mechanism.