In-Context Occam’s Razor: How Transformers Prefer Simpler Hypotheses on the Fly
Authors: Puneesh Deora
Bhavya Vasudeva
Tina Behnia
Christos Thrampoulidis
Venue: COLM 2025 (2025)
Keywords: In-context learning
transformers
linear regression
Markov chains
Bayesian Occam's razor
TL;DR: Transformers exhibit Bayesian Occam's razor in-context

Abstract: In-context learning (ICL) enables transformers to adapt to new tasks through contextual examples without parameter updates. While existing research has typically studied ICL in fixed-complexity setups, real-world language models encounter tasks of diverse complexity levels. This paper investigates how transformers navigate hierarchical task structures where higher-complexity categories can perfectly represent any pattern generated by simpler ones. 
We design testbeds based on Markov chains and linear regression that reveal transformers not only identify the correct complexity level for each task but also accurately infer the corresponding parameters—even when the in-context examples fit multiple complexity hypotheses. Notably, when presented with data generated by simpler processes, transformers consistently favor the least complex sufficient explanation. We theoretically explain this behavior through a Bayesian framework, demonstrating that transformers effectively implement an in-context Bayesian Occam's razor by balancing model fit against complexity penalties.

Published as a conference paper at COLM 2025
In-Context Occam's Razor: How Transformers Prefer Simpler Hypotheses on the Fly
Puneesh Deora', Bhavya Vasudeva', Tina Behnia', Christos Thrampoulidis!
1 University of British Columbia, Vancouver
2 University of Southern California, Los Angeles
(puneeshdeora, tina. behnia, cthrampo &ece. ubc. ca, bvasudev@usc.edu
Abstract
In-context learning (ICL) enables transformers to adapt to new tasks through contextual examples without parameter updates. While existing research has typically studied ICL in cixel-compeexity setups, real-world language models encounter tasks of diverse complexity levels. This paper investigates how transformers navigate hierarchical task structures where eegher-complexity categories can perfectly represent any pattern generated by simpler ones. We design testbeds based on Markov chains and linear regression that reveal transformers not only identify the correct complexity level for each task but also accurately infer the corresponding parame- even when the in-context examples fit multiple complexity hypothe- ses. Notably, when presented with data generated by simpler processes, transformers consistently favor the least complex sufficient explanation. We theoretically explain this behavior through a Bayesian framework, demon- strating that transformers effectively implement an in-context Bayesian Occam's 's razor by balancing model fit against complexity penalties.
Introduction
In-context learning (ICL) has emerged fundamental capability of large language models (LLMs), enabling them to adapt to novel tasks through contextual examples without param- eter updates (Brown et al., 2020). This remarkable ability has drawn significant attention in interpretability research, with studies examining both commercial-scale LLMs (Elhage et al., 2021; Wang et al., 2023; Min et al., 2022) and controlled, synthetic environments. The latter approach trains transformers from scratch on well-defined tasks-such as linear regression (Garg et al., 2022; Akyürek al., 2023; von Oswald et al., 2022; Zhang et al., 2023), discrete functions (Bhattamishra et al., 2024), and Markov processes (Edelman et al., 2024; Rajaraman et al., 2024; Park et al., 2025) offeringg precise control over training distributions and enabling direct comparisons with known algorithms. However, these synthetic studies typically restrict analysis to tasks of fixed complexity, diverging from real-world scenarios where LLMs encounter diverse tasks spanning multiple complexity levels.
To bridge this gap, we investigate ICL in setups with hierarchical task complexity structures. We design task categories with distinct complexity levels, where tigher-complexity categories form strict supersets of lower ones: specifically, the nigher-complextie category contains hypotheses capable of emulating any pattern produced by simpler categories. For instance, consider a transformer trained on next-token prediction for sequences generated by both order-1 and order-3 Markov chains. Since any order-l chain can be represented as a special case of an order-3 chain, an inherent ambiguity arises during inference: when presented with a sequence genuinely generated by an order-1 process, can the transformer identify the true complexity class, or will it default to the most expressive hypothesis in its repertoire?
Can transformers effectively differentiate between tasks of varying complexity during in-context learning? When presented with data compatible with multiple hypothesis classes, do they accurately identify the simplest sufficient hypothesis, or do they systematically default to the most complex available hypothesis?
1

---

Published as a conference paper at COLM 2025
(Pretraining) Tasks of Hierarchical Complexity
Inference
LLM Evaluation
Seq. App Spup? 3, Mig)
Prompt with ambiguous esamples
Prompt
co complex examples
Arb (Simlla)
Smpple Teok
Vlaye Vospptl
Darp Iuursko
x aepgle
~ K Japal
Model(X) is HXX) X Soogle
Eompples
11001
13010
2075
01100 -
02011
W complex
Complex Model(X) ss MXX), Ronppks
Query
01111
01111
Roopples
st. y2 MXx)) i
Model output
Figure 1: (Left) We train transformers on a hierarchy of task families in which the complex class strictly contains the simpler one (the simple category). (Middle) After training, the models use the in-context examples to infer both the appropriate complexity level and the task parameters that best fit the context, on the fly. (Right) We probe the same inductive bias in off-the-shelf LLMs: prompted with ambiguous examples that admit both simple and complex explanation (copy the first bit or take majority over three bits (second, fourth and fifth), respectively), the model performs in-context Occam's 's razor by consistently following the simple hypothesis (see Section 3.4 for details.).
Our systematic investigation answers this question affirmatively. See Fig. for an overview of our framework. Returning to our Markov chain example, Fig. 2 demonstrates that at inference time, the transformer successfully recognizes the true order of the generating process. When presented with sequence from an order-1 chain, (left) appropriately employs bigram statistics; when presented with an order-3 chain, (right) it switches to tetragram statistics. This confirms that transformers can distinguish between different complexity categories and adapt their predictions to the context, rather than defaulting to the highest-complextty hypothesis available.
Our contributions are follows:
We introduce framework for studying ICL across hierarchical complexity levels, where higher-complexity tasks form strict supersets of simpler ones, creating inherent ambigu- ity in hypothesis selection in context.
We empirically demonstrate in two representative testbeds, a Markov chain and a linear regression setting, that transformers correctly identify the simplest sufficient hypothesis rather than defaulting to the more expressive one.
We provide a theoretical justification through Bayesian lens, showing that transformers naturally implement Bayesian Occam's razor in-context by balancing data likelihood against model complexity.
Additionally, going beyond the two testbeds, we demonstrate the same behaviour on transformers trained on probabilistic context-free grammars (PCFGs) with varying complexity, and on pretrained LLMs (GPT-4) tested on representative setting with Boolean-function tasks.
2 Background and Previous Work
To understand the mechanisms underlying ICL, researchers have developed controlled synthetic environments with well-defined tasks. In just the last couple of years, there has been a proliferation of work on this topic. Here, we focus on reviewing two key synthetic setups: Markov chains and linear regression, introduced by Edelman et al. (2024) and Garg et al. (2022), respectively, which form the basis for our investigation into tasks of varying complexity. We then discuss follow-up studies most relevant to our analysis. For comprehensive review and additional references, we refer readers to Dong et al. (2024).
2.1 Markov Chains and Linear Regression as Testbeds for ICL
ICL of Markov Chains. Markov chains were introduced as a testbed for studying ICL by Edelman et al. (2024). Specifically, they showed that transformers trained sequences generated from random order-1 Markov chains learn to do in-context inference on unseen Markov chains of order-1 by learning to output bigram statistics of the context.
2

---

Published as a conference paper at COLM 2025
Uniform
Unigram Order
Digrem
Trigrm
Tetregram
KL Div(Diatriautioo Model)
Order
0.2
200
400
600
800
200
400
600
800
Sequences Seen (in thousands)
Sequences Seen (in thousends)
We train a transformer for next-token prediction on sequences generated by random order-l and order-3 Markov chains. During inference, we assess its ICL ability by evaluating performance on sequences derived from unseen order-1 and order-3 chains. (Left) shows the distance between the model's 's output distribution on the last token and n- gram statistics of the context for order-1 inference, as a function the number of sequences seen during training. (Right) presents analogous results for order-3 inference. Notably, the trained transformer can identify the true order (1 or 3) of the context, and then predict using either bigram tetragram statistics accordingly.
Figure 2:
To establish the foundation for our study, we now formalize their key findings. Let x
denote the f-th symbol from finite vocabulary [V] [1, V} of size V generated by k-th
order Markov chain, ie., p(x: = a|x, 11..,x1) = p( = 2|x 1,--X; k 1) for all E [V]
and let P E RV*xV denote the row-stochastic transition matrix with these conditional distributions as rows. Sample transition matrix at random: for example, sample each of its rows according x1, X2. XT of T > k symbols from this transition matrix by letting the first k entries [x1,..,,Xy] be drawn uniformly at random. A transformer model Mş parameterized by is trained to rege-regressively predict the next-symbol X1+1 using the previous f symbols Xge by (approximately) minimizing the expected loss, with respect to cross-entropy &:
Dirichlet prior with parameter a. Generate a sequence X := Xx7
L(0)= = (E-
Σ E--1 [==
((Mo(X=))
(1)
where D1r(1) @V denotes V independent draws from the Dirichlet prior. At inference, draw transition matrix Pinf ~ Dir(1) %V and let X Pinf denote sequence of length T generated from it that is given as prompt to the model. Edelman et al. (2024) show that the Kullback-Leibler (KI.) divergence of the model's output probability to the birgram statistics
E 1=2 1 XT, = v) ET 1x--1 xx), V U (V)
(2)
averaged over many realizations of X, Pinfs is (almost) zero. Here, 1() is the indicator function that equals 1 when the condition is true and 0 otherwise. Thus, the transformer looks back at the sequence to compute the bigram probabilities, and predict the next token E [V] using these probabilities.
ICL of Linear Regression. Before Markov chains, the first synthetic playground for ICL was introduced by Garg et al. (2022), who demonstrated that a transformer trained on random examples from linear regression tasks can learn to perform in-context inference on unseen tasks by computing the least-squares solution. In this setup, sequences consist of interleaved input-output pairs (x,y) y) of the form = [x1, y1, X2, V2, X1, y/ll,, with feature vectors E RR and labels E R. The feature vectors X; are sampled i.i.d from a standard Gaussian distribution, i.e., X; ~ (O, I,). Each sequence is characterized by a task-specific weight vector w (0, Ia), such that labels are generated as yt
x,, VI E [T].
Similar to the Markov chain setup, the transformer Ma is trained to predict label y1+1 using the previous pairs (x, y) along with the new input X (collectively denoted by X by
TTo be precise, the model input is [xt yy..., N,.,..,,,,,, yt), where y; Rdd with zeros in the first d - entries and yÚ as the last entry.
3

---

Published as a conference paper at COLM 2025
minimizing the expected loss over sequences with T pairs:
I(0)= Nix.
E (MMo(x()))11)))
(3)
where l is the squared loss. Garg et al. (2022) demonstrated that when presented with an inference-time sequence X = (1, V2,..., -1,, generated from unseen task vector Winf ~ N I2,, the transformer effectively infers Winf from the in-context examples. Specifically, the squared error between the model's prediction and the least- squares prediction Uis- = x wis where WLS A'y approaches zero when averaged over many realizations of X and Winf- Here, A E R(T-1)xd represents the feature matrix with columns X2,,X,,... x y I is the vector of context labels [yl.... ......,,yy-11)) and At denotes the Moore-Penrose pseudoinverse.
Other Related Work. Raventos et al. (2023) and Park et al. (2025) study linear-regression and Markov-chain settings, respectively, training transformers on sequences drawn from finite task set and show that limited task diversity induces a trade-off between task retrieval and task learning modes of ICL (Pan et al., 2023). On the theoretical side, Xie et al. (2022) give Bayesian account of ICL, and a complementary line of work provides explicit transformer weight constructions that implement functionalities observed to facilitate linear regression (Li et al., 2024; 2023; Ahn et al., 2023; von Oswald et al., 2022; Fu al., 2024) and Markov chains (Edelman et al., 2024, Rajaraman et al., 2024; Chen et al., 2024b). All of these studies focus on a single task family; none tackle settings where multiple task cate- gories are linked by a hierarchy of complexity. We bridge this gap by extending the Bayesian lens to hierarchically related tasks and showing that transformers consistently adopt the simplest hypothesis that explains the context in both regression and Markov-chain domains. Transformers have also been shown to perform algorithm selection: Lin & Lee (2024) reveal setrivaal-eersus-leamme modes in Gaussian-mixture regression, and Bai et al. (2023) show context-depent deeeedend switching among linear, noisy-linear, and classification rules. A broader simplicity view links ICL to prequential coding (Elmoznino et al., 2025), while Xiong al. (2024) study prompts containing several unrelated tasks in superposition. However, none of these studies tackle hierarchically related task categories, where a higher-complexity class can fully explain the context produced by any simpler class. We consider this challenging regime and show that transformers resolve the ambiguity through a Bayesian Occam's 's razor, consistently defaulting to the simplest hypothesis that still explains the context. See Appendix for a detailed discussion.
3 In-Context Learning Across Hierarchical Complexity Levels
Previous research has extensively examined ICL for fixed complexity tasks. But how do transformers behave when trained on tasks spanning multiple well-defined complexity levels? Can they accurately identify both the complexity category of given task and infer its underlying parameters in-context? Here, we introduce two cestbeds-xxteniing the Markov chain and linear regression paradigms- to systematically investigate these questions. Additionally, we verify the conclusions on probabilistic grammar setting.
of
3.1 Markov Chains
Modeling Tasks of Varying Complexity. To form task categories of varying coeptexity, train on Markov chains of different orders, where we group all chains of a fixed order into one task category. The category of order-k; has lower complexity than any order-k2 category for k2 k due fewer degrees of freedom. Importantly, any order-k; chain can be perfectly represented as special case of an order-k2 chain, making higher-order categories strict supersets of loder-order ones.
For concreteness, we train the transformer model Ma on sequences generated from two different task categories?. We consider order-1 and order-k categories, where k > 1. During training, we first sample a transition matrix P3 of order-s, which is then used generate
Multiple task categories can be treated in similar manner without substantial further insights. Thus, we focus on two categories for simplicity of exposition.
4

---

Published as a conference paper at COLM 2025
sequence X = X2, of length T. The order is sampled uniformly at random from the set ord = {1,k}. Similar to Eq. (1), we train autoregressively to minimize the (expected) loss, this time over Markov chains of both orders {1,k}, again with cross-entropy loss e, iee,
L(0)= == IE Bx--R-DDr(1)ev s~Unni((rd)
Σ T 1 (Xa()))))))1))
At inference, we prompt the model with a sequence generated from either order-1 (simple) or order-k (complex) chain. To evaluate the ICL ability of the model, we compare (with respect to KL divergence) the model's output probability for the token following the prompt to either the order-1 statistics of Eq. (2) or the order-k statistics
Σ 1x XT, X--2 " k+1 XT Xt = v)
O E [V].
(5)
1 1 = XT, XT 1,.",x k = k)
When prompted with an order-k chain, the transformer must recognize the higher-order dependencies that cannot be captured by order-1 statistics. However, when prompted with an order-1 chain, the transformer faces a more subtle decision, as any order-1 transition matrix can be perfectly represented as special case of an order-k transition matrix. This latter scenario directly tests whether the transformer defaults to the most complex hypothesis or correctly identifies the simplest sufficient one.
Transformers Distinguish Between Task Categories In-context. We train a GPT-2 type decoder-only transformer for all the experiments (Karpathy, 2023). Please refer to Ap- pendix C for specific training details.
In Fig. we train transformer on se- quences drawn from order-1 and order- 3 Markov chains (i.e., Eq. (4) with [1,3]). We then measure the KL diver- gence between the transformer's output distribution and several well-defined sta- tistical strategies when the input is taken from unseen order-1 or order-3 Markov chains: uniform, bigram, trigram, tetra- gram statistics (Eq. (5) for k .).
Untoom
Unigram
Bigram
Tigram
KL-DVIDierooeloo Model)
01
0.0
200
402
620
500
200
402
620
Sequencess Seen (in thousends)
Sequences Seen (in thousends)
We see clearly that with sufficient train- ing, the transformer consistently learns to distinguish between order-1 and order- 3 sequences, applying bigram and tetra- gram statistics appropriately to each case.
Figure A transformer trained on sequences generated by random order-1 and order-2 Markov chains can infer whether the context is order-1 or order-2, then generate predictions us- ing the corresponding bigram or trigram statis- tics. (Left) shows the distance between the model's output distribution on the last token and well-defined context strategies for order-1 inference, as a function of the number of se- quences seen during training. (Right) presents analogous results for order-2 inference.
We validate this finding in Fig. by re- peating the experiment with order-1 and order-2 Markov chains, again observing that the model accurately infers the un- derlying order. Crucially, in both experi- ments, the transformer does not default to the highest-order statistics (tetragram or trigram) for all inputs, despite these statistics having sufficient expressive power to model order-1 sequences. Instead, it selects the appropriate complexity level based on the in-context sequence. This demonstrates that the transformer effectively implements Occam's 's razor, identifying the simplest hypothesis adequately explaining prompt.
In these experiments, the context length T is kept fixed at T = 300 (Fig, 2) and T = 200 (Fig. 3) for both training and inference. Here, we focus on how the transformer learns to distinguish between the two categories as function of training iterations (equivalently, sequences processed). In Section 4, we provide additional experiments examining how performance varies with context length.
Bayesian Interpretation of the Selection Mechanism. We adapt a Bayesian viewpoint of ICL. (Panwar et al., 2024; Xie et al., 2022; Lin & Lee, 2024) to explain how transformers
5

---

Published as a conference paper at COLM 2025
choose the simplest task that explains the context. In this respect, we assume sufficient training data such that the transformer implements the Bayes optimal minimizer. Recall the mixture model described in Eq. (4), and let X [x],...,xT] XT] be an observed sequence for > k. Then, the Bayes' optimal predictive distribution for XT given is the mixture
x|
Σ
p(s X)
P X, s),
UE [V),
(6)
model output
calegory posterior
S-gram statistics of X
where p(s X) denotes the posterior probability of the order-s chain given the sequence.
Further, (xT=+1 = X, s) denotes the Bayes' optimal predictive distribution computed with respect to only order-s Markov chains. Specifically, for the Dirichlet prior Dir(1) @ it can be shown that p(xT+1 = o| x ,) which is a smoothed version of the order-s statistics in Eq. (5); e.g. see Edelman et al. (2024, Eq. (3)). When T > k, where both statistics are well-defined for both orders 1 and k, the smoothing effect becomes negligible (particularly since the smoothing is essentially independent of s).
From Eq. (6), the model's output (LHS) when trained with chains of all orders s can be seen as concex mixture of what would have been the model's output if the transformer was trained only on a single order. This means the model's output crucially depends
express as: p(s X) p(X s) (ss) p(X s') ((s'). Here, (Ts) is the prior on orders, which is set to be uniform in Eq. (4) and thus cancels out in the ratio, and p(X {s) is the marginal likelihood of the prompt (context) given the order s. Consequently, which category's statistics (s-gram) the transformer's 's outputs depends which marginal likelihood [X| dominates.
on the posterior probability coefficients p(s X), which we
Asymptotics of Likelihood. For large > k, standard combined with a Bayesian Information Criteria (BIC) style Laplace approximation (Schwarz, 1978; Ghosal & van der Vaart, 2017; Csiszár & Talata, 2006) gives the following useful approximation to the marginal likelihood (Schwarz, 1978):
conjugacy
s) Σ, log px(x; ",,
V2 (V-1) logT.
Here, Px denotes the empirical conditional probabilities derived from the prompt X. See Appendix B.1 for details. Thus, the marginal likelihood decomposes into an empirical likelihood term and model complexity penalty term
Bayesian Occam's Razor. Suppose the data is generated by an order-s* Markov chain, repre- senting simple category. We wish to compute the ratio p(X $ pX s*) for any complex category of order-s with s > s" Note that for T >> K, for any s 2 s* the empirical probabili- ties Px 14-1...., xt-s) approach the true data transition probabilities derived from the simple category s". Thus, the empirical likelihood terms in Eqn. (7) are (approximately) equal for any > s* This reflects the fundamental property that Markov chain of order s* can be perfectly represented as Markov chain of order s where the transition probabilities are invariant to the additional history. However, the model complexity penalty term being proportional to ys strongly favors the smallers s. Together, for T > k it holds:
p(X s) P
Eqn 6)
ss > 5*
RJ 0
+ p(s x)
p( = X Ps+
2 5" gram
(8)
statistics
model output conditioned on prompt from simple calegory
The model implements Bayesian Occam's razor in favor of the true lower-order model s".
Saturation to the Higher Order. Now, suppose that the true process is of order-k, rather than (say) order-1, where k > 1. The model order complexity term in Eq. (7) always favors the lower order (here 1). However, in this case the empirical likelihood terms differ to each other when evaluated for s = k vs 1. Previous works (Ghosal & van der Vaart, 2017; Csiszár & Talata, 2006) have shown that the per-sample improvement in log likelihood under the correct, higher-order model accumulates linearly in T. Concretely, for large T and some universal constant c > 0,
log p( k) log x s= x cT (v+ 1 1-V(VV-VV- log T
T
p(s = k x) 1
2
6

---

Published as a conference paper at COLM 2025
Thus, the posterior probability correctly saturates to the true higher order once T is large. In Appendix D, we construct a 2-layer transformer that can compute the empirical conditional probabilities px, which are key to compute the posterior p(s X) needed to realize the Bayes' optimal solution.
3.2 Linear Regression
Modeling Tasks of Varying Complexity We construct similar setup of two task cat- egories with varying complexity for ICL, of linear regression. We define two categories with hierarchical complexity: "complex" category parameterized by using the full feature space, and "simple" category where w lies on a lower-dimensional subspace of Rd with the structure category corresponds to sparse linear regression where half the features have no impact the output. The tasks in the complex category have higher degrees of freedom (d vs. d/2) and thus greater expressive power. Importantly, the complex category forms strict superset of the simple category in the sense that any sequence X generated from the simple in
Wa (V [4)
[ where Wa/2 NV La/2) This simple
category can be perfectly explained by some parameter from the complex category. The training details are analogous to Eq. (3) only now first sample uniformly from the set dim /d/2, d), then generate W, (O, I[.)) and generate the labels y; w x; for E [T] where for = d/2, we effectively use w [Ws, 0] (zero-padded to dimension d). We refer the reader to Appendix B.2 for details.
It
At inference time, we generate prompts using either wa or Wa/2 as the underlying task parameters. Let X denote a sequence with T in-context examples (xt, yt), and Xtest the query vector for which we want to predict label. We compare the transformer's 's output to two least-squares (I.S) benchmark estimates: ÿ wLS where w is either the full d- dimensional solution w := Aby or the restricted d /-mmennnsio solution w Ab/2y.
Here, A₂ == (x,, xx]] is the feature matrix, y [yı....,, ,y1....,y7 .r]] is the label vector, and A4/2 Ar(Ha/2 04/2]" is the projection of the feature matrix onto the first d/2 dimensions. When context length T < d or T < d for the respective estimates, the LS solution corresponds to the minimum 22-norm interpolating solution.
When prompted with a sequence generated by the complex regressor Wa the full- dimensional solution w/s naturally provides a better fit than w which is constrained to a d/2-dimensional subspace. However, the critical question arises when we prompt the model with a sequence generated by the simple regressor Wa/2: Does the transformer default to the more expressive full-dimensional solution, or does it correctly identify the simpler hypothesis? Specifically when d > 2 d/2, both solutions perfectly interpolate the context data, yet they generally differ in their predictions. This creates ideal test for whether transformers implement a form of Occam's razor when selecting in-context between equally valid explanations of varying complexity.
Transformers Distinguish Between Task Categories In-context Fig. 4 demonstrates how transformers distinguish between task complexities in linear regression. When presented with sequences generated by the complex regressor Wa, the transformer's 's predictions consistently align with the full-dimensional solution w!S rather than wTS. as expected since the restricted solution lacks sufficient degrees of freedom. The most interesting finding emerges when the transformer is prompted with sequences from the simple regressor Wa/2 in the regime where d T > d/2. Here, despite both w/5 and w perfectly interpolating LS the context data, the transformer's predictions significantly favor w This demonstrates that the transformer does not default to the more expressive full-dimensional solution but rather correctly identifies the simpler hypothesis- clear demonstration of Occam's razor at work. As context length increases into the T 2 regime, we observe w d/2 Wa/2 and the transformer essentially recovers the true regressor.
'The specific choice of d/2 the dimensionality and the alignment of the subspace with the coordinate axes are made for simplicity, without affecting the generality of our findings.
7

---

Published as a conference paper at COLM 2025
(a) d = 10
(b) d
20
Wa/2 0
W
W1/2 0
W
0.20
0.20
0.13
L5
47245
01.5
80
MSEd
0.15
0.15
80
d15
0.15
0.10
0.10
0.05
0.05
0.05
0.05
0.00
0.00
0.00
0.00
10
15
20
10
15
10
10
Connext Length
Contect Lengtth
Contest Length
Conteat Length
Figure 4: A transformer trained on T 39-long sequences generated using random (a) d 10 or (b) 20 dimensional regressors. We plot (M,([XXXtest) w Xtest) ) where refers to the two benchmark least-squares solutions w w in Section 3. When d/2 prompted with unseen sequences of (left) d/2 or (right) d dimensional regressors, the predictions align most closely with the corresponding LS solution w
d/2 or w/s. In the right figures, (orange, dot) is off-scale, demonstrating that wis is bad fit for sequences generated using Wa- See for evaluation over train time.
3.3 Probabilistic Grammars
In this section, we validate Occam's razor like inductive bias in transformers trained on probabilistic context-free grammars (PCFGs), which are stochastic systems designed to generate sequences symbols from specified vocabularies (Collins, 2013) (see Appendix B.4 for formal description). In our setup, non-terminal symbols N == {, B, C,C,, terminal symbols := (a, eos}, where eos corresponds to the end-of-string symbol. The 'complex' PCFGs use the following set of probabilistic productions:
$ > ABC BAC AAC BBC
with probabilities P1, P3, P4,
A + a,
B 3
(9)
C-> S eos
with probabilities 91,92
The 'simple' PCFGs are subset of the 'complex' PCFGs, where enforce p3 P4 0.
To generate sequences from this set of PCFGs, we first draw the probabilities pr and 9; randomly from a uniform Dirich- let distribution. Then, we keep generat- ing strings from this PCFG (starting from and following the production rules, till it returns eos) and concatenating them till
KL-DivDDst II Model)
Simple
Simple
Complex
Complex
0.0
the total sequence length matches the con- text length T. To control the complexity of generated sequences and prevent infinite recursion, we impose maximum deriva- tion depth dmaxi once this depth is reached, expansions of the nonterminal are re- stricted to produce the terminal symbol eos. For generating training sequences, we set dmax = 10, whereas for test sequences, is set as 0, i.e, we only get strings from the set {ab eos, basos, a b b with different probabilities.
75
Seq. Seen (in thousands)
Seq. Seen (in thousands)
Figure A transformer trained on sequences from random 'complex' and 'simple' PCFGs (see Eq. (9)) can infer the type of PCFG the context comes from. Here, the test sequences are of the form (...,a) a} or {...,bb, where ev- ery string before the last element is from the set { (ab eos, b eos, a a eass, bb eos}. (Left) compares the KL divergence between model output probabilities and the probabilities of the 'simple' PCFG. (Right) compares the KL diver- gence between model output probabilities and the probabilities of the ground truth 'complex' PCFG as well as those of the 'simple' PCFG.
In Fig. 5, we train transformer on se- quences drawn from the two types PCFGs. The test sequences are of the form {... a) or {... b). For sequences from the 'complex' grammar, we measure the KL divergence between the transformer's out- put distribution with the ground truth probabilities of the grammar, as well as those of the 'simple' PCFG. Note that since p3 p4 0 for the 'simple' PCFG, P(ala) P(blb) O. We
8

---

Published as a conference paper at COLM 2025
see that the transformer output distribution is closer to the 'complex' PCFG. For sequences from the 'simple' grammar, we find that the output distribution closely matches the 'simple' PCFG. This result shows that the transformer infers the type and parameters of the PCFG from which the context was generated. Here, we use the same context length for train and test sequences; please see Appendix A for test results with different context lengths.
3.4 Results Pre-trained LLM
In this section, we conduct an experiment with Boolean functions to validate in-context Occam's razor in a pre-trained LLM (GPT-4). We consider Boolean input {0, 1) and two Boolean functions: a 'simple' function where label 1/s = x(0) and a 'complex' function where label ye Maj f2, is]), where (2,,
if I z 3/2 and 0 otherwise, for z = x /2 /3]. For instance, let d = 5, and 11, 12, 13 be 1,3, respectively. Then, some example sequence-labe pairs are shown in Fig.
[d] and Maj(z)
We consider prompts with n examples (x, yi) followed by a query + in the format
(Simple)
Amb (Complex)
Comp (Simple)
(Complex)
6-5
1.00
shown in Fig. 1. To con- struct these examples, we first sample the indices i, i2, 13 ran- domly from [d] (this set is fixed for one prompt). sider two types of prompts, one with ambiguous examples, where the labels predicted by both functions are the same, i.e,, y,, ye, to check which function the model prefers, and the other with examples from the complex function, with labels y/c * ys, to check if the model can in-context learn the complex function as well. In both cases, the query unambiguous, i.e., ys + ye, to check whether the model uses the simple or the complex function to predict the final output.
0.75
0.50
0.25
0.00
12
16
4
12
16
4
8
32
16
Number Examples
Number of Examples
Number of Examples
We con-
Figure Results with GPT-4 for in-context learning Boolean functions with two types of prompts. We con- sider Vsimple = x(0 and Уcomplex = Maj(x(1,3,4)). Com- paring accuracy with respect to the simple and complex functions on unambiguous queries, we see that the model exhibits in-context Occam's 's razor: it uses the simple func- tion when given ambiguous examples and the complex function when the context can only be explained by it.
We prompt GPT-4 with these prompts and compare the output's accuracy (where the output is the token with the largest log probability) with respect to the simple and the complex functions, over 50 prompts. The results for different numbers of examples are shown in Fig. 6. We consider d 5,6,7. We observe that prompts with ambiguous examples which can be explained by both the functions, the model outputs on the unambiguous query match with the simple function. On the other hand, on prompts with examples that can be explained only by the complex function, it gets higher agreement with the complex function on the query. These results confirm that GPT-4 exhibits in-context Occam's razor in this setting with Boolean functions - when two Boolean functions explain the context, the model uses the simpler one to make predictions on the query input.
4 Additional Experiments
In this section, we present additional results with our Markov-chain setup by (i) examining the inference behaviour of trained transformer on sequences of varying length and (ii) testing whether LSTMss show similar Occam's 's razor-like bias. Please see Appendix A for further analyses: (i) transformers trained only on the complex task do not exhibit this bias; (ii) the effect of model scale; and (iii) the effect of the training-task mixture proportion.
Testing with Different Orders and Context Lengths. In this section, we test how varying the context length affects trained transformer's behaviour on sequences from different task categories. In Fig. 7, we take a transformer trained only on order-1 and order-3 chains with sequence length 400, and evaluate it orders-1, 2, 3, and sequences. For each evaluation, we measure the KL divergence of the model's ' output compared to uniform/ bigram/ trigram/ tetragram/ pentagram statistics for varying prompt sequence lengths.
9

---

Published as a conference paper at COLM 2025
Uniform
Unigram
Bigram Order
Thiramm
Tetragram Eval. Order
Pentagram
1C. (Dhttbbbttonnn mooett
Eval. Order
Eval. Order
100
150
203
220
200
250
402
100
150
200
350
300
2.0
-00
300
150
200
250
303
1.0
400
120
353
200
2.0
700
352
130
Seqpartce Length
Sequerrion Langeh
Sequerice Lerggh
Sequueice Langth
Figure 7: Effect of context length on inference behaviour. Figure shows trained trans- former's performance when prompted with sequences from different-order chains as function of the sequence length. This transformer was trained on sequences from order-1 and order-3 chains with sequence length 400. Here, we also evaluate its performance on order-2 and order-4 sequences, which were not seen during training.
On order-1 sequences, KL(bigram || model) stays essentially flat and low across different context lengths, indicating a consistent fit the simplest statistic. In contrast, (I.tttragram || model) is large for short contexts (where higher and lower-order statistics are less similar), and steadily drops as the context grows (as expected by Eq. (5)). The clear gap for short contexts further validates our hypothesis that the model prioritizes the simpler explanation.
results for different orders for se- quence length 2 200, we find that consistent with the main results in Section 3, for order-1 sequences, the KL divergence to bigram is the lowest, and for order-3, the KL to tetragram is the lowest. For
Comparing
Model Comparison Divergence
Count)
0
(C.06M)
LSTI (7.03M)
LSTM (LJ.62M)
LSTM (26.50M)
754!
200;
order-4, we find that the model does not learn pentagram statistics, but rather returns output based on tetragram statistics. However, for order- 2 sequences, which the model was never explic- itly trained on, the behavior more nuanced while the model avoids the simplest explanation (bigram), it does not strongly commit to either the more complex explanation (tetragram) or a specific intermediate complexity it has not been directly trained on (trigram). See App. A.3 for further discussion and additional results.
CON)
2.5
Kl. (N) || model) for onder:) seg. (x10-1)
Comparison with LSTMs. Comparison between LSTMs and trans- formers of different capacities, trained on a mixture of order-1 and order-3 chains, and tested on unseen sequences at the end of training. We observe that LSTM models require significantly larger capacity than transformers to in- fer the true order (1 or 3) and predict using the corresponding statistics (bi- gram or tetragram) at test time.
Figure 8:
Results for LSTMs. In this section, we train stacked LSTMs to investigate if they exhibit in- context Occam's 's razor similar to transformers. From Fig. 8, we find that LSTMs exhibit a weaker form of Occam's is razor-like bias, as they require significantly more capacity to consistently select the correct underlying Markov chain order from the input sequence at inference time compared to transformers. See App. A.6 for inference results with training time for the models in Fig. 8, and App. C for training details.
5 Conclusion
This work demonstrates that transformers trained on hierarchical task complexity structures effectively implement a form of Bayesian Occam's 's razor, consistently identifying the simplest sufficient hypothesis without defaulting more expressive model classes. Our finding raises important questions: From a mechanistic perspective, how do transformers internally implement this complexity selection? From an optimization standpoint, how do gradient descent dynamics lead to the emergence of these Bayesian selection principles during training? Future work could investigate more complex hierarchical structures beyond simple dimensionality differences, such as models with different structural constraints (e.g., various sparsity patterns in regression) or tasks where complexity varies along multiple axes simultaneously. Our findings suggest that principled hypothesis selection may be an inherent inductive bias property of transformers trained on diverse task distributions, potentially contributing to their remarkable generalization capabilities in real-world settings.
10

---

Published as a conference paper at COLM 2025
Acknowledgments
This work was partially funded by the NSERC Discovery Grant No. 2021-03677, the Alliance Grant ALLRP 581098-22, and CIFAR AI Catalyst Grant. PD and TB were also supported by the UBC 4YF Doctoral Fellowship. BV was supported by NSF CAREER Award CCF-2239265. This work was done in part while BV was visiting the Simons Institute for the Theory of Computing. The authors also acknowledge use of the Sockeye cluster by UBC Advanced Research Computing and the Discovery cluster by USC CARC. The authors thank Core Francisco Park for useful discussions and the anonymous reviewers for helpful feedback.
References
Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra. Transformers learn to implement preconditioned gradient descent for in-context learning. In Thirty-seventh Conference on Neural Information Processing Systems, 2023. URL https://openreview net/
Forum?id-LziniAXE1S
Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learning algorithm in-context learning? Investigations with linear models. In Inf. Conference on Learning Representations, 2023. URL https://openrewiew. goppeneeiie..nt/oorum//.. net/forummid 0g0X4H8yN41.
Yu Bai, Fan Chen, Huan Wang, Caiming Xiong, and Song Mei. Transformers statisticians: Provable in-context learning with in-context algorithm selection. In Thirly-seventh Confer- ence Neural Information Processing Systems, 2023. URL https::/oppnrreiieeeee for. id=1iMSqUuVg9.
Satwik Bhattamishra, Arkil Patel, Phil Blunsom, and Varun Kanade. Understanding in- context learning in transformers and LLMs by learning to learn discrete functions. In The Twelfth International Conference on Learning Representations, 2024. URL httpttt//opeerreiieeeeee/ net/forum? (ettforrmmdd=eeeyyeeffff
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini
Agarwal, Ariel Herbert-Vore Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel Ziegler, Jeffrey Wu, Clemens Winter, Chris Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, llya Sutskever, and Dario Amodei. Language models are few-shot learners. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 1877-1901. Curran Associates, Inc., 2020. URL httpps:/ File/1147ccdd5bfcc4447774 Bbwb8ac142684482Ba2200
neurips.... wttpsperrrpppp..pcpppppppp wtwpc//pppeeeeeee//eeeeeeeeeeeee pdf..
T Tony Cai, Tengyuan Liang, and Harrison H Zhou. Law of log determinant of sample covariance matrix and optimal estimation of differential entropy for nigh-dimensiioal gaussian distributions. Journal of Multivariate Analysis, 137:161-172, 2015. 2015.
Siyu Chen, Heejune Sheen, Tianhao Wang, and Zhuoran Yang. Training dynamics of multi- head softmax attention for in-context learning: Emergence, convergence, and optimality, 2024a. URL https: pertps://arxxe http://awxvvvv..gr..a wbtp:o////bbba/agggoss 19442.
Siyu Chen, Heejune Sheen, Tianhao Wang, and Zhuoran Yang. Unveiling induction heads: Provable training dynamics and feature learning in transformers. arXiv preprint arXiv:2409.10553 2024b.
Michael Collins. Probabilistic context/free grammars (pcfgs). Lecture Notes, 2013.
Imre Csiszár and Zsolt Talata. Context tree estimation for not necessarily finite memory processes, via bic and mdl. IEEE Transactions on Information Theory, 52(3):1007-1016, 2006. doi: 10.1109/TTI NTTT.20058844227
11

---

Published as a conference paper at COLM 2025
Qingxiu Dong, Lei Li, Damai Dai, Ce Zheng, Jingyuan Ma, Rui Li, Heming Xia, Jingjing Xu, Zhiyong Wu, Baobao Chang, Xu Sun, Lei Li, and Zhifang Sui. A survey on in-context learning. In Yaser Al-Onaizan, Mohit Bansal, and Yun-Nung Chen (eds.), Proceedings of the 2024 Conference on Empirical Methods in Natural Language Processing, pp. 1107-1128, Miami, Florida, USA, November 2024. Association for Computational Linguistics. doi: 10. URL https: / /awlaaaacaannoyyyyaoaaooo org/ 2024. emnlp-main. 64/.
Ezra Edelman, Nikolaos Tsilivis, Benjamin L. Edelman, Eran Malach, and Surbhi Goel.
The evolution of statistical induction heads: In-context learning markov chains. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024. URL
https://openrevieww wwt..e/orummiioagTTT6000000
N. Elhage, N. Nanda, C. Olsson, T. Henighan, N. Joseph, B. Mann, A. Askell, Bai, A. Chen, T. Conerly, N. DasSarma, D. Drain, D. Ganguli, Z. Hatfield-Dodds, D. Hernandez, A. Jones, J. Kernion, L. Lovitt, K. Ndousse, D. Amodei, Brown, J. Clark, J. Kaplan, S. McCandlish, and C. Olah. A mathematical framework for transformer circuits. Transformer Circuits Thread, 2021. URL
cirrimttrrriiitssseieeeiecsss puwp http:/222//f/aarrre /ttp///wwww///////////////////// siteewook/iidx.... index. html.
Eric Elmoznino, Tom Marty, Tejas Kasetty, Leo Gagnon, Sarthak Mittal, Dhanya Sridhar, and Guillaume Lajoie. In-context learning and occam's 's razor, 2025. URL htttps://opereevieee
Deqing Fu, Tian qi Chen, Robin Jia, and Vatsal Sharan. Transformers learn to seehineeessconde order convergence rates for in-context linear regression. In The Thirty-eighth Anmual Conference on Neural Information Processing Systems, 2024. URL https: wttpss//ee/eeeeee.eeee forum?2d=LLhhcooccb
Shivam Garg, Dimitris Tsipras, Percy Liang, and Gregory Valiant. What can transformers learn in-context? a case study of simple function classes. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in Neural Information Processing Systems, 2022. URL
Subhashis Ghosal and Aad van der Vaart. Fundamentals of Nonparametric Bayesian Inference, volume 44 of Cambridge Series in Statistical and Probabilistic Mathematics. Cambridge University Press, 2017. doi: 10.1017/988113909995
A. Karpathy. Mingpt. https://gtts / http://hww....mmmaaaaa///////e/e
2023.
Yingcong Li, M. Emrullah Iidiz, Dimitris Papailiopoulos, and Samet Oymak. Transformers as algorithms: generalization and stability in in-context learning. In Proceedings of the 40th International Conference on Machine Learning, ICML'23. JMLR.org, 2023.
Yingcong Li, Ankit Singh Rawat, and Samet Oymak. Fine-grained analysis of in-context linear estimation: Data, architecture, and beyond. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024. URL https: /ttps///////eeeeeeee/mmo net/forum?id= 1YPAYmfQqm.
Ziqian Lin and Kangwook Lee. Dual operating modes of in-context learning. In Proceedings of the 4'sst International Conference on Machine Learning, ICML/24, JMLR.org, 2024.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization, 2019.
Yue M. Lu, Mary I. Letey, Jacob A. Zavatone-Vett Anindita Maiti, and Cengiz Pehlevan. Asymptotic theory of in-context learning by linear attention, 2025. URL https://arxiv: org/abs/2405. 11751.
Sewon Min, Xinxi Lyu, Ari Holtzman, Mikel Artetxe, Mike Lewis, Hannaneh Hajishirzi, and Luke Zettlemoyer. Rethinking the role of demonstrations: What makes in-context learning work? In Yoav Goldberg, Zornitsa Kozareva, and Yue Zhang (eds.), Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, pp. 11048- 11064, Abu Dhabi, United Arab Emirates, December 2022. Association for Computational Linguistics. doi: 2022. wmw.pmmain. 759/.
URL https: http://aaaaanttlyy..ooo org/
12

---

Published as a conference paper at COLM 2025
Jane Pan, Tianyu Gao, Howard Chen, and Danqi Chen. What in-context learning "learns" in- context: Disentangling task recognition and task learning. In Anna Rogers, Jordan Boyd- Graber, and Naoaki Okazaki (eds.), Findings of the Association for Computational Linguistics: ACI. 2023, pp. 8298-83988888 8319, Toronto, Canada, July 2023. Association for Computational Linguistics. doi: 10.18653/v/ htt//22.iiningssaaddle 2023. findings-aclll... 527/.
URL https://aclanttoology org/
Madhur Panwar, Kabir Ahuja, and Navin Goyal. In-context learning through the bayesian prism, 2024. URL https://arxxi.. hwwpaaawwwiii..rabaara/aaaaaaaee org/abs/2306. 04891.
Core Francisco Park, Ekdeep Singh Lubana, and Hidenori Tanaka. Algorithmic phases of in-context learning. In The Thirteenth International Conference on Learning Representations, 2025. URL https: /ttp:///ppeeeeeieee//m/ meetforrumidd-ggtt
Nived Rajaraman, Marco Bondaschi, Ashok Vardhan Makkuva, Kannan Ramchandran, and Michael Gastpar. Transformers on markov data: Constant depth suffices. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024. URL https:
//openreview..m wwwpeenrvviee..eeffouu/muaaaee nett/foruumidd-Sottt2222
Allan Raventos, Mansheej Paul, Feng Chen, and Surya Ganguli. Pretraining task diversity and the emergence of non-bayesian in-context learning for regression. In Thirly-seventh Conference on Neural Information Processing Systems, 2023. URL https://openreviert net/ forum?1d-BtA2455XDE
Gideon Schwarz. Estimating the dimension of a model. The Annals of Statistics, 6(2):461-464, 1978.
Johannes von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordvintsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. ar Xiv preprint ar Xtt::2212.07677, 2022.
Lean Wang, Lei Li, Damai Dai, Deli Chen, Hao Zhou, Fandong Meng, Jie Zhou, and Xu Sun. Label words are anchors: An information flow perspective for understanding in-context learning. In Houda Bouamor, Juan Pino, and Kalika Bali (eds.), Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, Pp. 9849-9849888 Singapore, December 2023. Association for Computational Linguistics. doi: 10.18653/v1/200 emnlp-main.609. URL battps://aclantiiee hwtpooolooyyoooog2222222 1ttp2222..mmmmmmmll....mm emnlp-main. m00.pp.mai...........
Sang Michael Xie, Aditi Raghunathan, Percy Liang, and Tengyu Ma. An explanation of in-context learning as implicit bayesian inference. In International Conference on Learning Representations, 2022. URL https:://opeereeieew.ttttt wtt.:eeet//oormmmiiddddddte.
Zheyang Xiong, Ziyang Cai, John Cooper, Albert Ge, Vasilis Papageorgiou, Zack Sifakis, Angeliki Giannou, Ziqian Lin, Liu Yang, Saurabh Agarwal, Grigorios G Chrysos, Samet Oymak, Kangwook Lee, and Dimitris Papailiopoulos. Everything everywhere all at once: Lims can in-context learn multiple tasks in superposition, 2024. URL /ttps://arxiv.oor org/ abs/2410. 005.2003...
Ruiqi Zhang, Spencer Frei, and Peter L. Bartlett. Trained transformers learn linear models in-context, 2023.
Ruiqi Zhang, Jingfeng Wu, and Peter Bartlett. In-context learning of linear transformer block: benefits of the mlp component and one-step gd initialization. Advances in Neural Information Processing Systems, 37:18310-18361. 2024.
Yedi Zhang, Aaditya K. Singh, Peter E. Latham, and Andrew Saxe. Training dynamics of in-context learning in linear attention, 2025. URL
wrgp/br///////1111// 16265.
Piotr Zwiernik, Caroline Uhler, and Donald Richards. Maximum likelihood estimation for linear gaussian covariance models, 2016. URL https://arxiv..og//b///4818 56044
13

---

Published as a conference paper at COLM 2025
Appendix
A Additional Experimental Results
14
A.1 Effect of Curriculum Training
14
A.2 Training with Only the Complex Task
15
A.3 Testing with different context lengths
15
A.4 Effect of Model Scale
17
A.5 Effect of Training Mixture Proportion
18
A.6 Results for LSTMs
19
Additional Details for Section 3
19
B.1 Markov chains: Omitted details for asymptotics of likelihood
19
B.2
Linear Regression: Omitted training details
20
B.3 Bayes' Optimal for Linear Regression
20
B.4
PCFGs
23
C Details of Experimental Settings and Code
23
D Construction for Markov Chains of Varying Complexity
24
Related Work
26
A
Additional Experimental Results
A.1 Effect Curriculum Training
We used curriculum learning in the linear regression setting, following previous linear- regression ICL studies (e.g., Garg et (2022) and follow ups). To test its impact, we reran the experiment without any curriculum (Fig. 9, left). In that setting, the transformer eventually converges to the d //2-dimensional least-squares solution, but it takes substantially longer than with curriculum scheduling (Fig. 9, right).
Without Curriculum training
With Curriculum training
0.6
w E 2d2 d/2
Cunreuium Eniss
w E 9d22 d52
2ạ2 dLS
0.4
0.4
MSEld
MSEId
0.2
0.0
0
500 1000
500
1000
Sequences Seen (x:10))
Sequences Seen (x10*)
Figure 9: Normalized MSE (normalized by d) between the transformer's 's predictions and the two benchmark solutions-d/ 2-LSS (simple) and d-LS (complex)-over training steps. 20 (same setting as Fig. 4) 15. (Left) Model trained without curriculum
Models are trained with context length T 40, with and evaluated for sequences with Tisst scheduling on equal proportions of d/2- and d-dimensional regressors in every mini-batch. (Right) Curriculum schedule used in Fig. 4.
14

---

Published as a conference paper at COLM 2025
Specifically, we consider the following curriculum scheduling. Training begins with quences from hoomd/4-dimm 4-dimensional regressors. Every 2000 steps, the dimensionality of regressors in the entire mini-batch is increased by 2 until it reaches d/2. From that point onward, the dimensionality of regressors in half of each batch remains at d/2, while the other half continues to increase by every 2000 steps until it reaches d; thereafter, batches contain equal proportions of d/2 and d. This curriculum markedly accelerates convergence to the simpler d/2S predictor.
A.2 Training with Only the Complex Task
In this section, we train the transformer on fixed order-k Markov chains (following Edelman et al. (2024); Park et (2025)) to examine whether at inference time, can infer the correct order when presented with sequences generated from a lower-order chain (< k). In Fig. 10, we test transformers trained on only order-3 (left) and order-2 (right) chain sequences on sequences from order-1 chains and find that the model predictions don't match bigram statistics most closely. This finding is crucial, as it suggests that transformer trained on order-k chain learns only the order-k statistics of the context and does not generalize to lower-order statistics, We repeat this experiment for the case of linear regression (see Fig. 11) and observe similar behaviour.
Uniform
Unigram
Bigram
Trigram
Tetragram
Trained with order
Trained with order 2
Div| Model)
Order 1
Order 3
Order 1
Order 2
0.4
0.2
0.0
200
400
600
800
0
200
400
600
800
200
400
600
800
0
200
400
600
800
Seq. Seen (In thousands)
Seq, Seen (in thousands)
Seq. Seen (In thousands]
5eq. Seen (in thousands)
Figure 10:
Transformer trained on only the complex task does not exhibit Occam's razor- like inductive bias. A transformer trained on random order-3 (left column) order-2 (right column) Markov chains can only predict based on order-3 (left column) or order-2 (right column) statistics, and fails to predict based on order-1 statistics when given order-1 in-context sequences.
We also do the same experiment for the linear regression setup and train the transformer of sequences generated with only d-dimensional regressors, and similarly observe that the transformer does not exhibit an Occam's razor-like inductive bias.
A.3 Testing with different context lengths
In this section, test how varying the context length affects trained transformer's behaviour on sequences from different task categories in the Markov chain setting and the PCFG setting.
In Fig. 7, we take transformer trained only on order-1 and order-3 chains with sequence length 400, and evaluate it orders-1, 2, 3, and 4 sequences. For concrete quantification serving as a baseline for behaviors under order-2 and order-4, if we fix sequence length 300, the relative gap between the lowest KL values and the second lowest (a larger gap indicates more confident prediction) are 1.618 for order-1 and 1.074 for order-3 sequences. For order-4 sequences, we find that the model does not learn pentagram statistics, but rather returns output based on tetragram statistics. Now, moving to order-2 sequences, which again the model was never explicitly trained on, the relative gap of KL distances to trigram and tetragram is 0.1073, significantly smaller than the gaps reported above for order-1 and order-3.
In Fig. 12, evaluate the model in Fig. 3 on different context lengths at the end of training. We find that the model predictions most closely match with bigram for order-1 and trigram for order-2 sequences, respectively, expected. We also find that the gap between bigram
15

---

Published as a conference paper at COLM 2025
0.25
WER
0.20
WERM. dLs
MSEdd
0.15
WERE, CS
WEAt dLs
10
20
30
40
Length
Figure 11: transformer trained on X 11, 11, X2, 92. XT, yrl sequences generated using fixed complexity, random d-dimensional regressors Wi (see Section for details). We plot (Ma (M,,(XXxtss))
w Xtest) where Wbench refers to two benchmark least-squares solutions in d or 2 dimensional space described Section 3. We see that the transformer's
predictions align most closely with w no matter the type of inference-time regressor used
to generate the sequences (Wa or wa/2). This indicates that the transformer doesn't learn to estimate the lower-complexity solution w unlike the case in Fig. 4. Here, T 39 and d 10. Also the first curve (blue, dot) is out of the plotted range.
wss
-
Uniform
Unigram
---
Bigram
Trigram
Tetragram
Pentagram
00 (Dierituooon modett
Eval. Order
Eval. Order
Eval. Order
Eval. Order
0.3
0.0
100
130
302 Seque Langth
200
250
300
430
100
250
252
230
300
352
400
350
130
200
230
300
350
402
100
330
202
150
300
350
400
Sequence Length
Sequerce Langth
Sequarice Langth
Figure 12: Figure illustrates how varying the context length affects trained transformer's performance when prompted with sequences from different-order chains inference. This transformer was trained on sequences from order-1 and order-2 chains with sequence length 200. The results are averaged over sequences each from 250 transition matrices.
Simple
Complex
(isst Model)
Simple
Complex
0.2
1.0
0.1
0.5
0.0
0.0
20
50
10
20
40
Context Length
Context Length
Figure 13: Figure illustrates the performance of a trained transformer when prompted with sequences of varying context lengths from the two grammars. The context length during training was 50 in this case. We find that the transformer can identify the correct type of grammar for slightly shorter context lengths, but the performance starts deteriorating for very short context lengths.
and trigram order-1 sequences increases context length becomes smaller. This is consistent with the observations in Fig. 7.
In Fig. 13, we evaluate the model in Fig. (trained on PCFGs) different context lengths at the end of training. We find that for both simple and complex grammars, the performance deteriorates when the context length is (very) small.
16

---

Published as a conference paper at COLM 2025
Results Order-1 Sequences
Results Order-3 Sequances
Orees1
Bigram
Order-2 vs Totragram
Order3 KL Bigram
Order:3
Tetragram
Enfbocie Mr
152
51
230 Saven (x122))
(-107)
(x309)
Sean (x107)
Figure 14: Effect of model scale on inference behaviour. Each heatmap shows the KL divergence between the model's output distribution and the indicated n-gram baseline as training progresses (horizontal axis), across different embedding dimensions (vertical axis). We report results for models with 6 layers, 6 heads. We observe that with embedding dimension 30, the model successfully learns bigram statistics on order-1 sequences, but fails to learn tetragram statistics order-3 sequences. In contrast, all models with embeddimg dimension eventually learn both. We also observe a consistent trend: as overparame- terization increases, the training time required to acquire both order-1 and order-3 statistics decreases.
Results Order-1 Sequences
Results Orden-3 Sequences
Layers 6 Hoads
203
Order:1 KL vs Bigram
Order-3 KL vs Tetragram
Order-3 KL Bigram
Order-3 K. vs Tetragram
0.1
MG
192-
03
120
493
0.
(x10))
Secuenom Tetragram
Seen (x131)
Seen txx09))
Layers 8 Heads
vs Bigram
Bigram
KL Tetragram
Enposdiig Der
Sequenoie Sean (x10)
Secuerom Semn (x00))
Segarces Seen Ex131
Sequinoes Sean Lx:00
12 Layers 12 Heads
Orden-
Orden-1
Order-3 KL Bigram
Order
STONME Den
540 KED
200
700
200
240
400
520
032
Saque Sean (+107
Secueroo (x009
Seojarces (x-27)
Sequanies Span (x:27)
Figure 15: Effect of model scale on inference behaviour. Each heatmap shows the KL divergence between the model's output distribution and the indicated 1t-gram baseline as training progresses (horizontal axis), across different embedding dimensions (vertical axis). We report results for models with layers, heads (top row, used to report main results in paper), 8 layers, 8 heads (middle row), and 12 layers, 12 heads (bottom row). Darker/blus regions indicate lower KL divergence; brighter yellow regions indicate higher KL. In each row, the left two heatmaps show evaluation on order-1 test sequences, and the right two on order-3 test sequences. For the 6-layer, 6-head model (top), we observe that with embedding dimension 30, the model successfully learns bigram statistics on order-1 sequences, but fails learn tetragram statistics on order-3 sequences. In contrast, all models with embeddimg dimension > 30 eventually learn both. We also observe a consistent trend: as overparameterization increases, the time required to acquire both order-1 and order-3 statistics decreases. This trend also appears, though less prominently, for the deeper and wider models (8L, 8H and 12L, 12H).
A.4 Effect of Model Scale
In Fig. 14, we present results for larger transformer models by increasing the embedding dimension in layer, head transformer model. We find that larger models also exhibit
17

---

Published as a conference paper at COLM 2025
912
Order-1 KL vs Bigram
Order-1
KL vs Tetragram
Order-3 KL Bigram
Order:3 KL vs Tetragram
304
9.
220
256
03
203 200
403
520
600
SrtueRow: Seen (430))
500u8920r Seen 14104)
50qu99006 Scon (x121
Soguencre Sece 2411
Figure 16: Repetition of the experiment in Fig. when training on a mixture of order-1 and order-3 chains with layer transformer, with и heads in the first layer, 1 head in the second layer, and embedding dimension dim = 16n. We observe that for order-3 sequences, the model output probabilities are not closely aligned to tetragram statistics as in deeper models in Fig. 15. However, increasing the model size via scaling up the number of heads (and embedding dimension) similarly speeds up the time required to learn order-1 and order-3 statistics.
Urden1 KL vs Bigram
Orden-1 KL vs trigram
Orden-2 KL vs Bigram
Orden2 KL vs Trigram
320
356
bricchi Die
392
338
305 420
100
202
300
402 550
150
100
203
300
402
520
100 400 530
902
Secustom Seen (-32
Sequarcie Swan (4321)
SequaTOR Swan (x127)
Sequrrsas (4327)
Figure 17: Repetition of the experiment in Fig. when training on a mixture of order-1 and order-2 chains with 2 layer transformer, with и heads in the first layer, head in the second layer, and embedding dimension dim = 16n. We observe that increasing the model size speeds up the time required to learn order-1 and order-2 statistics, similar to the experiments with deeper models in Fig. 15 and two layer models in Fig. 16.
the Occam's razor-like inductive bias, and converge faster the number of parameters is increased. We also observe this trend in the next few experiments where we vary the number of layers, attention heads, and the embedding dimension.
In Fig. 15 we repeat the experiments in Fig. 14 for 8 layer, 8 heads and 12 layer, 12 heads models by varying the embedding dimension. In Figs. 16 and 17, we repeat the experiments in Figs. 2 and 3, respectively, with a 2 layer transformer model, varying the embedding dimension and number of heads in the first layer, while keeping the heads the second layer fixed to 1. We find that across all settings, the transformer infers the correct Markov chain order from the input context. We also observe that scaling the model size speeds up convergence.
A.5 Effect of Training Mixture Proportion
In Fig. 18, we train transformers while systematically varying the fraction of order-1 and order-3 sequences in each mini-batch (order-1 fraction (.91, For each setting, we evaluate at inference time: (i) (bigramm|/modee) and (ii) [L(terragamm||modl model) on both order-1 and order-3 test sequences.
We find that the transformer reliably learns bigram statistics for order-1 sequences, even when order-1 examples make up a small fraction the training mix. In contrast, learning tetragram statistics for order-3 sequences becomes increasingly difficult as the fraction of order-1 sequences rises - under such skew, the model tends to default to bigram- behaviour. Moreover, the number of training steps required to acquire tetragram statistics increases as the training distribution becomes more imbalanced.
Increasing the batch size helps offset this effect: larger batches ensure that more order-3 sequences are seen per step in absolute terms, which raises the threshold at which the model can still learn the correct higher-order statistics.
18

---

Published as a conference paper at COLM 2025
Reeults on Order-1 Bequencee
Results on Order-3 Sequences
Sigaas
Orden1 letroyram
Order Bigamm
KLVS letrogam
Batch sizee 52
O Frantoes
Ordenn Fradten
02
0
200
400
8C0
235
430
933
200
500
600
700
610
62
0.9
Crdend 10 vn Bicram
Orden1 MD us Seireguem
Drdee-3 10. vs Bigram
Orden1 XL Tetragram
Betch size 64
0.8
Coaa Thactinn
OWIS Rodion
Cettent Fractinn
2.5
01
22
200
820
400
900
420
ECO
400
Sequeroms (4:10)
Sequances (4001))
Sequereas (x.00)
SequaTOR
(X10/)
Figure 18: Influence of training mixture proportion on inference behaviour. Each heat- map shows the KL divergence between the model's output distribution and the indicated n-gram baseline as training progresses (horizontal axis), for different fractions of order-1 sequences in the mini-batches (vertical axis) using batch size 32 (top row) and 64 (bottom row). We find that even when this fraction is very low, the model learns bigram statistics for order-1 sequences. Increasing this fraction makes the learning of tetragram statistics for order-3 sequences increasingly difficult, but increasing the batch size helps offset this effect.
A.6 Results for LSTMs
Fig. 19 shows training time evaluation of the LSTM models shown in Fig. 8 and comparison with transformer models.
LSTM
dim C.86M paramns
Al,
7AM parmes
61, parees
di,din parters
Transformer
61. dH1 dim 90, 0.75M
6L 6H1, dim 120, 1,26M
6L. 6H, dim 192, 3.02M
Figure 19: The plots show the KL divergence between the model's output and various n-gram statistics for order-1 (left) and order-3 (right) sequences, as function of training time, across LSTM models of varying depth and width reported in Fig. 8. For direct comparison, results for a 6-layer, 6-head transformer (dim= 192) are shown in the bottom- right (reproduced from Fig. 2).
B Additional Details for Section 3
B.1 Markov chains: Omitted details for asymptotics of likelihood
Using a BIC style Laplace approximation to the marginal likelihood (Schwarz, 1978; Ghosal & van der Vaart, 2017; Csiszár & Talata, 2006), we have
log p( s)
log
4 log
Here, P3 denotes the parameters of the order-s chain that maximize the likelihood of the observed sequence X. dg is the effective dimension of an order chain, concretely, ds ys (V 1). For an order-s chain, know that p(x [X Ps) -II Pp, (xi X4 1 Xt-s) and
19

---

Published as a conference paper at COLM 2025
the parameters P3 that maximize this are the empirical conditional probabiiittes derived from the observed sequence, i.e. PP, (xt Xt-1,..,, Xt-s) = (xx; Xt, for each context of length Formally, px(xx/ X s) Mit-gr- (x)
- 1(8) where Mx (X) denotes the
number of symbol transitions x; ... > X, in the sequence X.
B.2 Linear Regression: Omitted training details
Training sequences with T in-context examples are generated as follows: we first sample uniformly from the set dim = {d/2,d) then generate w, (O) For task dimension $ d/2, we effectively use W = [ws, 0] (zero-padded to dimension d). We then sample feature vectors X ~ (0, I and generate the labels yt w X for f E [T], forming sequences Xst [1, X2, 1/2, X,, yr]- We
this to train transformer Ma to auto- regressively predict the label V1+1 using the first f pairs of examples Xet by minimizing the population square-loss (analogous Eqn. (3)).
L(0) := (o.L)
(X Xrest) y141),
(10)
w...V((11.)) s~UUnil(dim)
where & the squared loss.
B.3 Bayes' Optimal for Linear Regression
We can explain the preference for the simpler hypothesis using a similar Bayesian viewpoint as done in Section 3.1 for Markov chains. Specifically, the Bayes' optimal prediction on the test query Xtest given the in-context demonstrations is y W Xtests where WBO is the mixture: WBO = [de(d/2,4) ((d') Aa. w/s Here, p(d' Aa,y) denotes the posterior probability of the d'-dimensional regressor given the matrix vector A, and the labels y, and wis [w]s, 0). Using Bayes' theorem, we express the posterior probabilities in terms of the likelihoods (assuming a uniform prior on regressor dimensionality): p(d' Aary) p(y (da///)2) p(y A), For Gaussian priors over Wa' we can derive analytical expressions for these likelihoods. Similar to our Markov chain analysis, the likelihood incorporates a complexity penalty that favors lower-dimensional models. Consequently, when the underlying regressor generating sequence X is w* Wa/2 and d > T>> d/2, we find that (d//2 y) >> p(d) y), resulting in WBO w 1/2' Conversely, when w War the improved fit outweighs the complexi complexity penalty, giving WBo w/S. This Bayesian mechanism explains why transformers implement Occam's 's razor by selecting the simplest hypothesis in-context.
L.S.
Calculating the Bayes Optimal. In order to find the Bayes optimal, which is the posterior mean of conditioned on the sequence X, we need the conditional distribution p(w Aa,y). By Bayes' rule,
p(y Ag, w) p(w) p(w Aa,y) = f ppy X, w') p(w") dw"
(11)
Under uniform prior on the regressor dimensionality, the overall prior p(w) - 1 (w) +
(w),,, where (w/z(w), (w() ~ (0, denote the respective priors over the d/2 and d dimensional regressors. Here, Pa/2 (w) corresponds to Gaussian (0 11/2) over the first d/2 components of w with zeroes in the remaining d/2. Plugging this into Eq. (11), and using the fact the d/2 and w!S are the Bayes' optimal solutions under their respective priors, we get mixture of solutions
E W
Ag,y] =
La/2
I +
La L.s.
L + La
d/2 La/2 + Ln
20

---

Published as a conference paper at COLM 2025
where Lo p(y]Aa, d') denote the marginal likelihood of y under the regressor of dimen- L sionality d'. Notice the term E.T, p(d" Aa, using Bayes' rule and uniform prior over the dimensionality of regressor. Thus, the Bayes' optimal is the following wosterior-weighted combination of the respective LS solutions
p(d' Aary) d'E(d/2,
Σ
wis. we
WBO =
Expressions for L1 and L2. The marginal likelihood Lg under the noiseless linear regres- sion setting is
Lr === =
b(y 6 Ad w) Par (w) dw.
This can be interpreted as the "density" of when viewed as random variable A w for Par Concretely, if the sequence generating regressor w Pa(w), then L (a.s.), as no w Pd/2 solves y A,w. Hence, WBO w.S.
Bayesian Occam's Razor. Now consider the other case when w ~ Pa/2(w). Here, regres- sors from either prior can "explain" the context and will have non-zero marginal likelihood. Consider the regime d > T d/2, the density of y under Pa/2 is
La/2 (271)d/4 v / A4/2)
1
exp - || 1S 1/2))
exp (pp(-1 |w/ 11))
La (2T) /det(A, (a AT)
In order find the ratio of the above marginal likelihoods, we need to compute the difference
A == log logge((A,A
logdet( ////2///////1/2))
In order to compute Aa, first look at its expectation. Later, we use concentration argument to get high probability bound on 4
Expectation of 44- As Aa has i.idd Gaussian entries, both A44 and A A//2 follow Wishart distributions. It is known (e.g,, Zwiernik et al. (2016, in Prop. A.1)) that
+ log det( (A4A]) ( Tlog,,
E log ( A4/2) = ////2(T/2) (d/2) 1og
where up is the multivariate digamma function. InWe will first calculate the difference It further holds that
T (r(d/2) 44/2 (T/2) Σ 4 (
4/2
Σ
(12)
=
where op is the digamma function. The expaasoticeppansion of the digamma function gives us
(z) = (z) + 0 2),
larg(z)) < 72, >> co.
(13)
Letz; - and 2/ 1- Using Eq. (13) in Eq. (12) when T,d co, and ignoring the second order terms, we have
d/2
(d(2)) (1122(T/22)))
Σ log(z;)
Σ
log(z/:
20/
22
d/2
= logg
+
7
Σ
2
(T d/2) log2
(14)
Term-1
Term-2
21

---

Published as a conference paper at COLM 2025
Term-1 can be calculated as
1h 2e = log 11
d'(T dd(TT d/2)! T![d-T]!
log
d llog(d) - d (T - d d/2) log( (T d/2) (T d/2) Tlog (T) + T (d- T) log(d - T) d - T
(dlogd)(c- 1/2) + + ((c- (c11/2)(log(e- 1/2) 1) cloge (1 c) log c))
g(c)
where the second equality follows by using Stirling's approximation, and neglecting the logarithmic terms. The third equality follows by using T = cd, where c E (1/2,1). It can be shown that g(c) bounded function in (1/2, We can easily see that it monotonically decreasing function
gc) log (c 1/2)(1 6) 0, C (1/2,1).
Using this we have, g(1) S (c) S g(122), where g(122) log2 and g(1) - 2 + 1). Therefore, we can write Term-l as
log
(15)
Term-2 can be bounded as
d/2
1
1
1I
2- -c) = O(1),
1/2
2
1
2
12
d72
27
d/2
$ 7-1
2
Σ
12
1 =
1
T-o+172) 1-4477)=00(1
where we use T = cd, where c E (1/2,1).
We can simplify Eq. (14) by using Eq. (15) with only the dominant d log d term and ignoring (11) Term-2. Then, using this difference in the multivariate digamma functions, we have
E E[Aa] IE log d EE log det (et(AA//AA/2) (d do
(16)
High probability bound on Az. Cai et al. (2015, Thm. 1) establish a central-limit theorem (CLT) for the log-determinant of high-dimensional sample covariance matrix. Specialised to our (un-normalised) Gram matrix A their result gives
log (eg(AAA))))
[E[log det( [e
d-+00
(0,,)),
2a = Σ w (T+ttl
U,TTT
where 11 is the trigamma function. Exactly the same argument applied to the A 1 Ad/2 yields
log det (Att(A//4442) Ellog det( (442
d/2
T-t00 N(,,1),
T
Σ 4h
(T
TT,422
Each centred log-determinant is sub-Gaussian with variance proxy dRT T,A/2 O(1). More concretely,
Pr [log det( ) - E log det(A,A]]]
1) < 2exp(-ct2) and likewise for
Choosing Clogd d and applying union bound gives, for suitable constant C > 0,
log [ - E. Fog det( (A4AT) S Clogd]
and { (ogdet(AA//Aa/2) Elog det (et(A < Clogd)) > 1-
22

---

Published as a conference paper at COLM 2025
Hence, the difference satisfies
A - [[Ad]] < 2Clog with probability
(17)
Combining Eq. (17) with the expectation difference computed in Eq. (16) yields
)et(A,4 = exp(E(A.)) exp(O(logı)) 2 ddda
det( A A//
with probability least 22/d and some constant c' which is dependent on the ratio = T/d E 1).
At the same time, it is easy to show that [| |
LS d/2 12 0. Put together, we have
La/2/Lo 2 (2x) q''d
which is > for large and T d/2. This implies that, pdd/2 Adry) ~ 1, and WBO LL
2/2'
B.4 PCFGs
Formally, a PCFG is described by the tuple S, where N is a set of nonterminal symbols, is a set of terminal symbols, and R is collection of production rules that dictate how nonterminals are expanded into sequences of terminals and nonterminals. A special start symbol initiates the generation process, and assigns probabilities to these production rules, indicating their likelihood of selection during sequence generation.
C Details of Experimental Settings and Code
For experiments Markov chains, linear regression and PCFGs, train GPT-2 type decoder-only transformer (Karpathy, 2023). We use AdamW (Loshchilov & Hutter, 2019) optimizer in all experiments with learning rate of 10-4, unless stated otherwise. We set the batch size as 32.
Markov chains. For all Markov chain experiments, vocab size V 3. For the order-1 and order-3 experiments in Fig. 2, context length T for all sequences was set 300. We used a layer, 6 head transformer, with embedding dimension set to 192. For the order-1 and order-2 experiments in Fig. 3, context length T was set to 200. We used a layer transformer with 20 heads, and embedding dimension encoding in all the experiments. For the fixed-order experiments in Fig. 10, we used the setup from the corresponding variable order experiment.
set to 320. Additionally, we used relative position
Linear regression. We used a 12 layer, head transformer with embedding dimension 256 similar Garg et al. (2022). We set 10 and 20 (for Fig. 4), and context length T 39.
PCFGs. We used 4 layer, head transformer with embedding dimension 128, and context length of 50.
LSTMs. All LSTM models were trained with the AdamW optimizer (learning rate 2 x 10-3 6s equal to (0.9, 0.99)) and cosine learning rate decay with a minimum learning rate of 10-5 In all experiments, we set hidden dimension be 4 times the embedding dimension.
Evaluation details. The results in Fig. 7 and Fig. 12 are averaged over 15 sequences each from 250 transition matrices reported with 95% bootstrapped confidence intervals. The results in Fig. 4, Fig. 11 are averaged across 2000 contexts, each generated with inde- pendently sampled ground-truth regressor. Shaded bands show 95% confidence intervals obtained from 100 bootstrap trials.
23

---

Published as a conference paper at COLM 2025
Code.
The code is available at attpppp//gttt..ccc/////eeeereeep
D Construction for Markov Chains of Varying Complexity
Recall from Eq. (6) that the model learns tasks different complexity by implementing the Bayes optimal strategy. To achieve this, the model must compute several key quantities, including the s-gram statistics p(- x77 and the category posterior.
Previous work has shown that two-layer attention-only model with heads in the first layer can compute the relevant s-gram statistics for the last token (Edelman et al., 2024; Rajaraman et al., 2024). To extend those results, for order one statistics, we demonstrate that by using V heads in the second layer, the model can represent all the conditional distributions p(u v) for any pair of vocabulary tokens simultaneously. These conditional distributions are the building blocks that produce the log-likelihood used to compute the posterior. This approach can be extended higher-order Markov chains by using V heads in the second layer.
Lemma 1. A two-layer attention-only transformer with one attention head in the first layer, and V attention heads in the second layer can output all the empirical conditional distributions
1 -x
p(u v)
for any u, E [V].
Σ 1
for any imput sequence X.
Proof. We outline how to set the transformer's parameters so that, for each pair (v, the model's final output can encode the conditional probability p(u u v).
Recall that the input to the model sequence [x,, xT] of length with x; [V]. Let zt E R2 be the input embedding of the f-th element of the sequence xt, and Z 2r. be the sequence of input embeddings.
We write the two-layer transformer with V heads the second layer as
f(x)
W. (Z) || ... hy),,
where "I| denotes concatenation, W, E IR RRxdVV is the final linear projection, and each (Z) is the output of the v-th attention head in the second layer. Specifically,
ho (Z) = (Attn + f) o (Attn + 1) (Z),
with
(Zt = softmax(mask(AA ZWy, where Wy is the value matrix, and the attention map E R TxT is defined as
(2. Wo + (i Wxz
Aij
vd
Here Wo, Wx, are the key and query matrix, Fk is the relative positional embedding, and mask(.) enforces causal masking so that position can only attend to positions 1,...,i.
We will show how to choose the model's weights so that, for each the embedding from the w-th head of the second layer, (Z),, produces vectors which-wher multiplied by W,- W,-yield the conditional distribution p(- v) across the sequence X. For simplicity, we will omit the superscript & in Attn, from now on.
Layer 1: Isolating the previous token. We first show how to configure the weights of the first layer so that each position copies the immediately preceding token z; For this construction need the embedding dimension of 3V.
24

---

Published as a conference paper at COLM 2025
Set the embedding matrix E [I,, lv, Vxa In this case, the input embedding 2; is le,, e,,,0 OOT where ep is the v-th basis vector.
Now let the key, query matrix and the positional embeddings be as follows.
cc Lvvv
We
WK
0
R == =e2.11
0
where c is a large constant. With these weights, the attention matrix A is zero for all entries except when j = that is, Aij if * 1. For the entry where j = 1, have
1 = 1 - c. As approaches infinity, the softmax function saturates. In other words, if we define
S softmax (Aas
then Si = when j = and for all other j. This ensures that the model only attends to the previous position in the first layer.
Now choosing
(
0
Wy
0
we have
T T Wv 10, TOOTT
Attn (Z);, Σ S1zz
This extracts and carries along the "middle" V components of the embedding at position i-1 Putting things together, the embedding of after the first layer will be
z = (Attn, + 1) (Z); 2 27 + 2
27
Layer 2 v-th head: Computing conditional probabilities p(- v). We now use V heads in the second layer to encode, for each possible previous token how likely the next token is 1/ E V.
In this layer, we set the relative positional embedding r 0 and
(0
0
0
We
e T
0
Wx
- ep e
Eg e
0
0
0
With this weight configuration, we have
AT 2τ T WoW,
fc
if 1 =
0 o.w.
Then, as C + 00,
1 1(x-11 = v)
ST.1
Σ
1 1 1 v)
Now setting
0 ly
Wy =
0
0
25

---

Published as a conference paper at COLM 2025
we have
I 12) = E S
Σ [1
[0,0,e,] + +27
NE[V]
= Σ
p(u 1,)00, +2
=c[v]
Finally, with Pvv3v = [0,0,1], setting Wo as
0
0 P
0
Wo- =
0
0
P
we have
f(Z) = Wol M1 (2) || ... hv (2)
Ep(- 10- 1)ex
...
Σ 10- v= [VVe]]
11
11
E Related Work
Task Diversity: Learning vs Retrieval. In the setups investigated by Garg et al. (2022); Edelman et al. (2024), the transformer is trained on sequences generated from fresh tasks drawn from the appropriate distributions each optimization iteration. Raventos et al. (2023); Park et al. (2025) have further extended this study for linear regression and Markov chains respectively, examining scenarios where the model trained on sequences generated from a finite number of tasks, thus modeling and analyzing the impact of task diversity in ICL and revealing tradeoff between two distinct modes of ICL: task retrieval and task learning (I'an et al., 2023). See also Lu et al. (2025) for high-dimensional statistical approach that theoretically justifies some of these empirical findings. In this work, we focus on the learning mechanism of ICL and maintain the vanilla setting of training with fresh tasks at each iteration. However, in our settings these tasks originate from a finite number of complexity categories rather than single one.
Theoretical Explanations: Linear Regression and Markov Chains. One of the first works investigating the theoretical foundations of ICL is Xie al. (2022), which proposed Bayesian perspective that has informed numerous subsequent studies, e.g., Raventos et al. (2023). Most relevant to our work, Edelman et al. (2024) demonstrated that the bigram statistics rule implemented by transformers trained on order-1 Markov chains is the Bayes' optimal solution (given the context) under the Dirichlet prior. Similarly, for linear regression, the least-squares solution learned by transformers is the Bayes' optimal predictor (given the context) under Gaussian priors in the setting of Garg et al. (2022). A complementary explicit weight constructions that implement
approach to explaining ICL has focused functionalities observed to facilitate ICL, such as gradient descent on the context for linear regression (Li et al., 2024; 2023; Ahn et al., 2023; von Oswald et al., 2022; Fu et al., 2024) and statistical induction heads for Markov chains (Edelman et al., 2024; Rajaraman et al., 2024; Chen et al., 2024b). However, relatively few studies have investigated how-or whether- these weight configurations can actually be reached during training (Zhang et al., 2023; 2024; Chen et al., 2024a; Zhang et al., 2025). None of these theoretical frameworks has addressed the specific setting of multiple task categories with hierarchical complexity relationships. In this work, we extend the Bayesian viewpoint to explain how transformers select the simplest sufficient hypothesis in both linear regression and Markov chain settings. While these two domains differ in several respects- the format of the prompt (pairs of inputs vs. stream of symbols), the optimization objective (squared loss vs. next-token prediction), and the underlying mechanisms (gradient descent vs. statistical induction heads) we investigate both settings under unifying prism.
26

---

Published as a conference paper at COLM 2025
Algorithm Selection. Most closely related to our work are Lin & Lee (2024); Bai et al. (2023) investigating ICL regression across multiple hypotheses. Lin & Lee (2024) considers ICL of linear regression where features x and tasks are drawn from Gaussian mixtures, generalizing the setting of Garg et al. (2022). This richer framework enables principled study of the task retrieval versus learning modes through Bayesian perspective. Similarly, Bai et al. (2023) investigates ICL under tasks from multiple categories, such as linear regres- sion, noisy linear regression, and linear classification. They demonstrate that transformers implement algorithm selection in context, applying the most appropriate algorithm for the provided context at inference time. However, these works do not address tasks from different categories with clearly defined hierarchical complexity relationships, where higher-complexity category can fully explain the context generated by simpler categories. In this particularly challenging setting, we show that algorithm selection is implemented via form of Bayesian Occam's razor, allowing transformers to identify and apply the simplest sufficient hypothesis. We demonstrate this both for regression and Markov-chatt settings. More recent work by Elmoznino et al. (2025) posits that a meta objective of ICL is linked to prequential coding algorithm that simultaneously minimizes both algorithm complexity and prediction error on in-context demonstrations. While this intriguing interpretation reveals form of simplicity preference in ICL, it differs fundamentally from our focus on hypothesis selection across explicitly defined hierarchical complexity structures. Additionally, Xiong et al. (2024) explore how pretrained LLMs perform ICL when several tasks are presented in superposition within single context. Not only does their framework differ by embedding multiple tasks in the same context simultaneously, but also these tasks are not organized into hierarchically related complexity classes.
27


Reviews:
[{"review_id": "mV3J5g6i36", "reviewer": "Reviewer_Qdck", "summary": "The authors formulate hierarchical tasks where higher complexity tasks can strictly represent lower complexity tasks (example: k-order markov chains), they then examine whether transformer architectures do ICL to learn orders of task complexity, and choose the simplest possible interpretation of task given in the prompt at inference time. The authors empirically assert that transformer models can discover different complexity levels for hierarchical tasks, and will prefer the simplest possible explanation. In practice the evaluation is accomplished by measuring the kl-divergence between model distribution and different order statistics in order to measure whether models can switch to different complexity levels as a function of input complexity. This work builds on related approaches which focus on either specific tasks or finite sets of tasks, to examine tasks which are explicitly hierarchical -- specifically predicting continuations of k-order markov chains. The authors present a Bayesian interpretation of results, showing that the transformer architecture can learn to select the simplest hypothesis when trained on tasks that enable hierarchical factorization. \n\nThe paper is clearly written with a good theoretical exposition, but most of the content is focused upon the problem and theoretical framing rather than upon anything specific to the transformer architecture, or to LLMs in general. The key takeaways of the paper aren't really clear, but the insights are still probably interesting for a subset of the COLM audience interested in bertology-related topics.", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "- clear exposition and theoretical motivation of synthetic tasks that measure the transformer architecture's ability to learn hierarchical task representations \n- good experimental setup that should be straightforward to replicate and extend to new tasks", "reasons_to_reject": "- low applicability in the current presentation to real world LLM evaluations, no examination of how results may generalize to i.e. foundation models \n- no examination of other possible architectures such as LSTMs -- is this a unique characteristic of transformer models? Do they need to be a certain size / dimension etc as a function of task complexity, etc ... \n- the takeaway for readers aren't really clear beyond \"here's a cool observation about transformer models in this specific setting\"", "questions_to_authors": "for the COLM audience, suggest moving more of the mathematical exposition to the appendix and adding more analysis of the transformer architecture itself", "extra_scores": {}}, {"review_id": "MsjhWEoaeS", "reviewer": "Reviewer_BYfu", "summary": "The paper studies how autoregressive transformers trained for in-context learning (ICL) handle tasks that have a hierarchical structure, where simpler instances can be perfectly represented by more complex ones. The authors consider 3 synthetic tasks: Markov chains with order k (increasing), dense or sparse linear regression, and probabilistic context-free grammars (PCFGs). Through experiments, they empirically show that models trained on data drawn from a mixture of simple and complex tasks (each separately shown during training) select the simplest hypothesis that fits the context during inference. They also provide some Bayesian based arguments, arguing that this observation is a consequence of marginal-likelihood trade-offs between data fit and model complexity.", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "- Investigating whether transformers exhibit an implicit “Occam’s razor” during ICL is an interesting research problem that was relatively unexplored from previous studies; this work sheds more lights in models’ inductive biases.\n    \n- The authors present strong empirical evidence supporting the hypothesis that LLMs trained for ICL on a mixture of simple and complex task instances (where the latter are able to fully represent the former) select the simplest one during inference. In fact, I think that the experiment in Appendix B, where the authors show that a model trained only on a complex instance does not infer the statistic of a simpler one is particularly interesting, and thus does not generalize well.\n    \n- The Bayesian arguments that the authors present to support their empirical observations seem convincing, and rely on reasonable assumptions that agree with their experimental setting.", "reasons_to_reject": "- Although the authors present some ablations on model scale in Appendix D (especially considering models with fewer parameters), I think that exploring how the main observations change at different (especially higher) scales would be an insightful addition.\n    \n- The paper never varies the proportion or ordering of simple vs. complex instances during training. Thus, it is unclear whether the simplicity bias depends on the specific mixture or could be amplified/mitigated by curriculum-style scheduling as the amount of training data grows, or the complexity of the instances is slowly increased.", "questions_to_authors": "What do you see as the key implications of your findings for real-world downstream tasks?", "extra_scores": {}}, {"review_id": "kN6joTWnBh", "reviewer": "Reviewer_peHZ", "summary": "While most prior work on in-context learning (ICL) has focused on tasks with fixed complexity, this paper shifts attention to scenarios involving varying levels of task complexity. The authors investigate how transformers handle hierarchical task structures, where more complex categories can fully represent patterns generated by simpler ones. To explore this, the authors outline two illustrative testbeds—Markov chains and linear regression—and demonstrate that, with sufficient training, transformers can distinguish between sequences of varying complexity and apply them appropriately. Additionally, they offer an explanation for the selection mechanism, grounding it in the principles of Occam’s Razor. However, the validation of this theory is somewhat limited, with a brief discussion of a probabilistic grammar setting that lacks in-depth analysis. This raises questions about how the simplified example relates to real-world applications.", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "This is a highly theoretical paper that addresses an important topic, with two testbeds carefully described to explore the issue at hand.", "reasons_to_reject": "It is not clear that the results fully reflect real-world environments, which is the primary motivation of this paper. The experiments themselves are not thoroughly described, and the use of four non-terminal and tree terminal systems in a probabilistic context-free grammar does not provide a convincing approximation of language complexity. As a result, the findings may not generalize well to more realistic, real-world settings.", "questions_to_authors": "Why did you select GPT2? Would the behavior change with different models or different architectures?", "extra_scores": {}}, {"review_id": "3zBcUaQOOh", "reviewer": "Reviewer_neZS", "summary": "The paper studies the performance of in-context-learning in transformers on tasks with varying complexity levels. It investigates if a simpler solution is preferred for a simple task even if the transformer  has access to a more complex solution. There are both theoretical and empirical arguments presented to conclude that indeed this is the case. \n\nThe markov-chain experiment’s results do not seem to support the hypothesis clearly, but are not presented as such and thus I recommend rejection.", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "- The paper presents a novel hypothesis and proposes clearexperiments to test it.\n- It conducts precise experiments in three different settings (markov chain, linear regression, and synthetic grammars) to study the claims.\n- There is well-written theory to support the empirical evidence.", "reasons_to_reject": "The markov chain experiment is the expanded on and studied the most in the main body of the paper, but the results seem to not clearly support the claim that on order-1 chains the model resorts to using the simpler hypothesis. \n\nIn both Figure 1 and 2, it is correct that the model’s output has a clear separation between the order-1 and order-$k$ statistics on the higher order sequence and chooses the more complex solution. But on the left side for the simpler lower order sequences, the KL divergence is equal for both the order-1 and order-$k$ statistics. This makes sense because in equation 5, the order-$k$ statistics calculated on an order-1 sequence (for a long enough sequence or enough samples), should be equal to the order-1 statistics.This is obvious in figure 1 where higher order statistics all perform similarly on the left. In figure 2 left trigram performs similarly as well.\n\nThus this experiment does not provide convincing support of the hypothesis, and can be indicative of the model still using order-k solutions on an order-1 sequence. Since this is the main experiment of the paper, I recommend rejection.", "questions_to_authors": "- In figure 2, the legend and plots do not match up. Does the lime yellow plot correspond to uniform or tetragram?\n- What are the results of the markov chain experiment on order-$k’$ sequences where $1<k’<k$?", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper investigates an important and novel hypothesis: Transformers' ICL inherently favor simpler hypotheses when given tasks structured hierarchically by complexity. The authors provide clear theoretical framing based on Bayesian principles and experimentally demonstrate this "Occam's razor"-like inductive bias using well-designed synthetic testbeds (Markov chains, linear regression, and probabilistic grammars).

All reviewers agree on the clarity, theoretical depth, and novelty of the hypothesis, along with the strength of the controlled experiments. Initial concerns centered around limited experiments regarding real-world applicability, model scale generalization, and potential ambiguities in interpreting Markov chain results. Authors' detailed rebuttals successfully addressed most major concerns through additional experiments: verifying the inductive bias in larger models (GPT-4), exploring model architectures (LSTMs), clarifying the Markov chain experiment interpretations, and examining curriculum and task mixture effects.

That said, I am surprised to see that the paper currently lacks a discussion on structural risk minimization (SRM), an essential theoretical context for understanding inductive biases and model complexity trade-offs. The authors must explicitly address this relationship in their revision to fully situate their findings within the broader theoretical literature.

Overall, the paper effectively establishes ICL's inductive biases toward simpler explanations in hierarchical contexts, with empirical evidence. The improvements promised in the authors’ rebuttals, along with a needed discussion on SRM, would enhance the paper's clarity, applicability, and generalizability.
Decision: Accept