Relational reasoning and inductive bias in transformers and large language models

TL;DR

This paper investigates how transformer models perform relational reasoning, especially transitive inference, comparing in-weights and in-context learning mechanisms and their generalization behaviors.

Contribution

It reveals that training regimes and representation structures critically influence transformers' ability for transitive inference, with insights into mechanisms underlying reasoning.

Findings

IWL models learn linear embeddings enabling transitive inference.

ICL models generalize transitively only when data necessitates, often using match-and-copy.

Pre-training ICL on linear tasks enhances reasoning to resemble IWL behaviors.

Abstract

Transformer-based models have demonstrated remarkable reasoning abilities, but the mechanisms underlying relational reasoning remain poorly understood. We investigate how transformers perform \textit{transitive inference}, a classic relational reasoning behavior from psychology which elicits inference about indirectly related items (e.g., if $A > B$ and $B > C$, then $A > C$). We compare in-weights learning (IWL) and in-context learning (ICL) behaviors and mechanisms on these tasks, and fine profoundly different patterns of generalization. IWL models learn a linear embedding, which leads to transitive inference as well as other behavioral effects present in humans and animals. ICL models, in contrast, are capable of learning to generalize transitively, but only do so when it is necessitated by the training data, otherwise learning a match-and-copy strategy. Interestingly, pre-training ICL models on in-context linear regression tasks that provide them with a latent linear representation is sufficient to make the ICL behaviors and internal representations qualitatively and quantitatively more like IWL. In order to test whether the same inference patterns are present across in large language models, we leverage a congruency paradigm which allows us to differentially probe IWL and ICL generalization patterns without access to their training data. We indeed see IWL reasoning leads to more transitive generalization than ICL. Moreover, we find that prompting the ICL models to use a linear mental map led to increased transitive inference over different geometric prompts. Together, these results reveal that both the training regime and the geometric structure of induced representations critically determine transformers capacity for transitive inference.