Can Mamba Learn How to Learn? A Comparative Study on In-Context Learning Tasks

• Prior work links ICL to demonstrations without weight fine-tuning and reports improving Transformer ICL as training loss decreases.The emergence mechanism remains under investigation.Source: §1 · p. 1
• The study examines whether attention-free models can perform ICL, a question of interest as alternatives to attention address its quadratic cost.Mamba is highlighted for efficiency and near-state-of-the-art pretraining performance at smaller scales, while ICL capabilities usually emerge beyond 3 billion parameters.Source: §1 · p. 2
• SSMs often match Transformers across ICL tasks; Mamba struggles with decision trees and retrieval but outperforms Transformers on sparse parity.Table 1 also reports that Transformer fails on sparse parity and Mamba struggles to retrieve vector-valued MQAR values.Source: §1 · p. 2, Table 1 · p. 2
• MambaFormer combines Mamba and attention, performing well across evaluated tasks and learning both sparse parity and retrieval.The architecture replaces Transformer MLP blocks with Mamba blocks and eliminates positional encodings.Source: §1 · p. 2
• The study focuses on non-language ICL tasks and smaller models, so findings may differ for actual language settings at higher parameter counts.The authors also report synthetic formal-language ICL experiments as evidence of potential language capabilities.Source: §1 · pp. 2–3

• Transformer ICL research includes explicit meta-learning and proposals that attention mimics optimization; linear-attention analyses connect the objective to preconditioned gradient descent.Ahn et al. (2023) and Mahankali et al. (2023) prove this result for one layer of linear attention.Source: §2 · p. 3
• Prior work connects retrieval with ICL through frozen GPT-2 nearest-neighbor behavior and induction heads that correlate with ICL behavior.The induction-head finding is attributed to Olsson et al. (2022).Source: §2 · p. 3
• S4 uses input-independent state-space parameters and a convolutional form, while Mamba makes its state-space parameters input-dependent for sequence mixing.S4’s convolutional expression enables near-linear complexity using the Fast Fourier Transform.Source: §2 · pp. 3–4
• Sub-quadratic models have been proposed to avoid attention’s quadratic sequence-length cost, but prior work found them behind attention on multi-query recall.The recall task generalizes the induction-head task.Source: §2 · pp. 3–4
• Language-based ICL benchmarks report differing model comparisons: LSTMs outperform Transformers on GINC, while Transformers outperform sub-quadratic models on random-finite-automata formal languages.Source: §2 · p. 4

3.1 Model Training for In-context Learning

• Models are trained from scratch per task and evaluated on 1,280 prompts; configurations keep parameter counts roughly matched and capacity is plotted against training FLOPs.Training samples function-and-input prompts and minimizes expected loss over prefixes containing prior examples plus the next input.Source: §3.1 · pp. 4–5, §3 · p. 4

3.2.1 Learning regression

• Regression benchmarks use Gaussian inputs and include linear, sparse-linear, two-layer ReLU, and decision-tree functions.Source: §3.2.1 · p. 5

3.2.2 Learning with outliers

• Outlier tasks replace prompt examples with out-of-distribution or confounding dummy vectors; evaluation uses the 50th clean example at fixed locations.Source: §3.2.2 · p. 5

3.2.3 Learning discrete functions

• Sparse parity samples inputs and subset S uniformly, trains with cross-entropy, and scores binary accuracy.Source: §3.2.3 · p. 6

3.2.4 Learning Chain-of-Thought

• The Chain-of-Thought task interleaves intermediate hidden features with inputs and labels, then evaluates final outputs and intermediate predictions.Source: §3.2.4 · p. 6

3.2.5 Learning retrieval

• Vector MQAR tests retrieving associated value vectors for queries matching keys sampled from the sequence.The task uses mean squared error loss.Source: §3.2.5 · p. 6

3.2.6 Learning synthetic formal languages

• Synthetic-language probes include GINC, generated from hidden Markov models, and RegBench, based on random finite automata and valid-next-token accuracy.Source: §3.2.6 · p. 6

4.1 Mamba can in-context learn!

• Across the ICL suite, Mamba outperforms S4-Mamba and S4, performs on par with Transformer as training FLOPs scale, and can lead in smaller equal-depth settings.The reported comparisons are at equal computation resources; Figure 2 plots each architecture’s best-performing model at the same FLOPs.Source: Figure 2 · p. 7, §4.1 · p. 8

4.2 Performance gaps in more complex ICL tasks

• On robust regression, Mamba beats Transformer with prevalent fixed outliers but loses on orthogonal outliers; with 90% identical replacements it leads on clean-test OOD accuracy.The orthogonal-outlier comparison uses sequence length 101; the many-outlier setting uses sequence length 512.Source: Figure 3 · p. 8, §4.2 · p. 9
• On Chain-of-Thought-I/O, Mamba is comparable overall, leading at smaller problem sizes while Transformer surpasses it as model size increases.Mamba’s performance is more influenced by hidden dimension; Transformer performance is relatively stable across problem sizes.Source: Figure 4 · p. 9, §4.2 · p. 9

4.3 Challenges in parity and retrieval

• In vector MQAR, Mamba struggles to retrieve vectors, with normalized-vector mean squared error above 0.1 in all reported cases.The authors suggest increasing key-value pairs can overwhelm SSM hidden-state capacity; the 16-query setting has lower test loss than the four-query setting, which they conjecture reflects more training samples per sequence.Source: Table 3 · p. 10, §4.3 · p. 10
• For sparse parity with (d, k) = (10, 2), Mamba and S4-Mamba succeed while Transformer remains at random-guessing performance despite larger models and longer training.The authors suggest convolution or gating may matter more than input-dependent selection for parity, and leave the responsible architectural component for future study.Source: §4.3 · pp. 10–11, Figure 5 · p. 11

5 The Advantage of Hybrid Architectures for In-context Learning

• MambaFormer is competitive across the ICL suite and excels at sparse parity and retrieval, tasks that Mamba and Transformer struggle with individually.The authors identify a Mamba initial layer as the key ingredient.Source: §5 · p. 12, §5.2 · p. 13

5.1 Simultaneously learning parities and retrieval

• MambaFormer learns sparse parity as quickly as Mamba in sample complexity, while Standard Hybrid learns it less efficiently.A Transformer with an initial Mamba block also solves parity almost as efficiently as Mamba.Source: §5.1 · p. 12
• MambaFormer nearly closes Transformer’s vector MQAR retrieval gap, while SSMs may be overwhelmed as key-value pairs increase substantially.Adding a Standard Hybrid attention layer significantly improves Mamba on vector MQAR.Source: §5.1 · pp. 12–13

5.2 All-in-one ICL performance

• Both hybrid models match Transformer on decision trees and outperform Mamba, including at larger parameter sizes.A small MambaFormer trained for 100k iterations (< 10^17 FLOPs) performs as well as models trained with nearly five times the FLOPs.Source: §5.2 · p. 13
• On clean test sequences, MambaFormer resembles Transformer on outlier regression, while Standard Hybrid resembles Mamba; which performs better depends on outlier structure.Mamba learns with one outlier vector, whereas Transformer learns better with a subspace of outlier vectors.Source: §5.2 · p. 13

5.3 In-context learning formal languages

• The formal-language experiments use GINC and RegBench as proxies for language ICL, but benchmark performance may not translate to real-world language ICL.The authors state that the relationship remains unclear for GINC.Source: §5.3 · pp. 13, 15
• On GINC, Mamba has the best ICL accuracy among non-LSTM models, while Transformer has lower perplexity; Standard Hybrid is comparable to both and MambaFormer performs slightly worse.These comparisons are reported for the GINC benchmark.Source: Table 4 · p. 13, §5.3 · p. 14
• On RegBench, hybrid architectures converge faster than Mamba and Transformer and achieve higher accuracy.The authors note that RegBench favors Transformers over attention-free models.Source: §5.3 · p. 15, Table 5 · p. 15

6 Discussion

• SSMs, especially Mamba, are capable in-context learners, but SSMs struggle with decision trees and retrieval while Transformers struggle with sparse parity.MambaFormer achieves best-of-both-worlds performance on the paper’s ICL suite.Source: §6 · p. 15
• The discussion frames MambaFormer’s performance as a response to the different task weaknesses of SSMs and Transformers.Source: §6 · p. 15
• Future work includes relating ICL-suite performance to language modeling, integrating architectural elements, identifying ICL-relevant features, and testing language-modeling effects.Source: §6 · p. 15

Impact Statement

• The paper studies language-modeling architectures to identify their strengths and weaknesses.Source: Impact Statement · p. 15
• The impact statement says the study provides recipes for new architectures.Source: Impact Statement · p. 15
• The authors say the outcomes may facilitate efficiency and architectural improvements for large language models.Source: Impact Statement · p. 15

A Experimental Setup

• The section describes the experimental design and configured setup.Source: A Experimental Setup · p. 20
• The authors say their code is available in the project repository.Source: A Experimental Setup · p. 20
• The authors direct readers to the repository for detailed implementations.Source: A Experimental Setup · p. 20

A.1 Model architectures

• The evaluated models include decoder-only GPT-2-family Transformers, Mamba, and hybrid variants.The hybrid variants include Standard Hybrid and MambaFormer configurations.Source: §A.1 · p. 20
• Transformer layers pair a Multi-Head Attention block with a Multilayer Perceptron, while Mamba layers contain two Mamba blocks.Source: §A.1 · p. 20
• Standard Hybrid and MambaFormer layers combine one Multi-Head Attention block with one Mamba block, and the attention blocks use eight heads.Source: §A.1 · p. 20

A.2 Model training and configurations

• Table 6 groups architectures by size and includes narrower, deeper variants; models deeper than 32 layers were excluded based on prior accuracy observations.S4 embedding dimensions were scaled by 1.75 to match Transformer parameter counts.Source: Table 6 · p. 20
• Most models train for 500,000 steps with Adam; vector-valued MQAR uses 300,000, and learning-rate choice is especially consequential as larger models can experience gradient explosions.The default learning rate is 1e−4; learning rates and gradient-norm clipping values are also swept.Source: §A.2 · p. 20
• The data setup fixes x dimension at 20 and batch size at 64, while a curriculum increases input dimension and context-point count every 2,000 steps.The passage describes the fixed dimensions as the train and test data setup, then separately specifies the curriculum.Source: §A.2 · p. 20
• The configurations treat two Mamba, attention, or feed-forward blocks as one layer to roughly align parameter counts across Transformer, Mamba, and S4-Mamba.Standard Hybrid and MambaFormer omit feed-forward networks, while S4 embeddings are enlarged to match Transformer parameter counts.Source: §A.2 · p. 20
• The linear-regression ablation varies width and depth at fixed FLOPs to assess their effects on ICL and identify strong configurations.It also tests whether optimizer choice changes the results and conclusions.Source: §A.2 · pp. 20–21

A.3 Chain-of-Thought-I/O settings

• For Chain-of-Thought-I/O, the study uses a two-layer ReLU network, varying model scale at fixed input and hidden dimensions and varying hidden dimension at small scale.The scale experiment fixes d = 10 and k = 8; the problem-scale experiment varies k among 4, 8, and 16.Source: §A.3 · p. 21, Table 7 · p. 21

A.4 Many-outlier regression settings

• Many-outlier regression uses two model-size groups and a long context of 1,024 tokens, requiring multi-GPU training.The context is specified as length 1,024 (N = 512).Source: §A.4 · p. 21, Table 8 · p. 21

A.5 Vector-valued MQAR

• The training set consists of 300,000 samples; models are trained for 64 epochs with batch size 64 and evaluated on a test set of 3,000 samples.For each setting, the best result among four learning rates from 10^-4 to 10^-2 is reported.Source: §A.5 · p. 21

B FLOPs Computation

• FLOPs are estimated by multiplying block multiplication counts by six for forward and backward multiply-accumulate costs, assuming batch size one.A Standard Hybrid block combines attention and Mamba, with multiplication count 10LD^2 + 2L^2D when linear terms are ignored.Source: B FLOPs Computation · p. 21

C Exploring effects of width vs. depth in ICL

• The ablation tests four width–depth configurations across four model-size groups on linear regression, comparing Transformer, Mamba, and MambaFormer at fixed FLOPs.It compares Adam with AdamW using linear warmup and cosine decay; for each configuration, the best of three learning rates is selected.Source: C Exploring effects of width vs. depth… · pp. 21–23
• AdamW with a scheduler improves performance, especially for larger models and for Mamba and MambaFormer, while the study’s conclusions remain optimizer-agnostic.The optimizer ablation reports that Mamba and MambaFormer may perform even better under this setup.Source: Figure 9 · p. 22, C Exploring effects of width vs. depth… · p. 24
• Two-layer Transformers struggle on linear regression, whereas MambaFormer can learn it with two layers; 128-layer models also struggle when too narrow.The study concludes that exploring 4–32 layers is sufficient for finding configurations on other tasks at fixed FLOPs.Source: Figure 10 · p. 23, C Exploring effects of width vs. depth… · pp. 23–24