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

Jongho Park1, Jaeseung Park2, Zheyang Xiong3, Nayoung Lee3, Jaewoong Cho1,
Samet Oymak4, Kangwook Lee1,3, Dimitris Papailiopoulos1,3

1 KRAFTON, 2 Seoul National University,
3 University of Wisconsin-Madison, 4 University of Michigan, Ann Arbor This
work was done during an internship at KRAFTON.

###### Abstract

State-space models (SSMs), such as Mamba Gu & Dao (2023), have been proposed
as alternatives to Transformer networks in language modeling, by incorporating
gating, convolutions, and input-dependent token selection to mitigate the
quadratic cost of multi-head attention. Although SSMs exhibit competitive
performance, their in-context learning (ICL) capabilities, a remarkable
emergent property of modern language models that enables task execution
without parameter optimization, remain underexplored compared to Transformers.
In this study, we evaluate the ICL performance of SSMs, focusing on Mamba,
against Transformer models across various tasks. Our results show that SSMs
perform comparably to Transformers in standard regression ICL tasks, while
outperforming them in tasks like sparse parity learning. However, SSMs fall
short in tasks involving non-standard retrieval functionality. To address
these limitations, we introduce a hybrid model, MambaFormer, that combines
Mamba with attention blocks, surpassing individual models in tasks where they
struggle independently. Our findings suggest that hybrid architectures offer
promising avenues for enhancing ICL in language models.

00footnotetext: Email<EMAIL_ADDRESS>Correspondence:
<EMAIL_ADDRESS>
##  1 Introduction

Modern large language models (LLMs) exhibit remarkable in-context learning
(ICL) capabilities, enabling them to learn new tasks with a few demonstrations
and without further weight fine-tuning. Although the exact emergence mechanism
of these capabilities warrants further theoretical and empirical investigation
(Chan et al., 2022; Wei et al., 2022; Min et al., 2022b; Schaeffer et al.,
2023), experiments on larger Transformer-based models consistently demonstrate
that their ICL capabilities improve as training loss reduces (Brown et al.,
2020; Kaplan et al., 2020; Muennighoff et al., 2023).

Meta-learning, or “learning to learn,” has been extensively studied
(Schmidhuber et al., 1997; Ravi & Larochelle, 2016) and recently regained
interest in the context of ICL, particularly concerning Transformer models
(Vaswani et al., 2017). Garg et al. (2022), for example, proposed various ICL
tasks, such as learning linear regression, and evaluated the ability of
transformers to perform them when specifically trained to do so. On the other
hand, Min et al. (2022a) studied fine-tuning language models to explicitly
learn and perform ICL. Following these footsteps, numerous research studies
have been dedicated to understanding the mechanics of Attention that enable
such meta-learning capabilities, either through constructive arguments or
extensive experimental investigation (Akyürek et al., 2022; Li et al., 2023b;
von Oswald et al., 2023b; Bai et al., 2023; Yang et al., 2023a; Li et al.,
2023a; von Oswald et al., 2023a).

As Transformer language models are currently the only large models that have
been reported to be capable of ICL in practice, this raises the question:

Can attention-free models perform ICL?

| Transformer | Mamba | MambaFormer
---|---|---|---
Linear regression | ✓ | ✓ | ✓
Sparse linear regression | ✓ | ✓ | ✓
2NN regression | ✓ | ✓ | ✓
Decision Tree | ✓ | $\blacktriangle$ | ✓
Orthogonal-outlier regression | ✓ | $\blacktriangle$ | ✓
Many-outlier regression | $\blacktriangle$ | ✓ | ✓
Sparse parity | ✗ | ✓ | ✓
Chain-of-Thought I/O | ✓ | ✓ | ✓
Vector-valued MQAR | ✓ | ✗ | ✓
Table 1: Model performances on various ICL tasks. We label the model’s
performance with ✓ if the model performs on par with other baseline models, ✗
if the model fails to learn the task, and $\blacktriangle$ if there exists
performance gap compared to other models. Transformer fails in learning sparse
parity Mamba fails in vector-valued MQAR. Our proposed MambaFormer succeeds in
all tasks.

This question holds merit, especially considering that several recent studies
have attempted to move beyond attention-based networks due to their quadratic
cost (Gu et al., 2022b; Dao et al., 2022; Gu & Dao, 2023; Poli et al., 2023;
Peng et al., 2023; Sun et al., 2023; Yang et al., 2023b). In this work, we
focus specifically on state-space models (SSMs), and particularly Mamba (Gu &
Dao, 2023). Mamba was recently demonstrated to be highly efficient while
achieving near state-of-the-art performance in standard pretraining language
data sets, such as the Pile (Gao et al., 2020), but at smaller model scales
(_e.g._ , up to 3 billion parameters), surpassing transformers and other
attention-free architectures across various language and non-language tasks.
However, ICL capabilities usually emerge at scales beyond 3 billion
parameters. As a result, the potential of these attention-free models to
perform ICL remains underexplored, as testing such hypotheses usually requires
scaling beyond the 7 billion parameter level. Nonetheless, we can still
investigate small-scale ICL capabilities by specifically training a model to
perform in-context learning, following the approach of Garg et al. (2022).

##### Contributions.

In this study, we introduce a diverse set of ICL tasks to evaluate the
performance of Transformer and various SSMs, including state-of-the-art models
like Mamba and S4 (Gu et al., 2022b). Our findings reveal that most of these
SSMs can effectively perform ICL, matching the performance of Transformers
across multiple tasks. However, Mamba demonstrates some limitations in
learning decision trees and retrieval tasks (as also noted by (Arora et al.,
2023)), but can outperform Transformers in other complex ICL tasks, such as
sparse parity, where Transformer models struggle. Performance of different
models on each task is summarized in Table 1.

Since there seem to be tasks where either family of models is better, we
explore the impact of interleaving SSM blocks with multi-head attention
blocks, similar to (Gu & Dao, 2023). We introduce MambaFormer, a novel hybrid
architecture that integrates Mamba and Attention layers, while eliminating the
need for positional encodings, as shown in Figure 1. MambaFormer seems to
leverage the strengths of both Mamba and Transformers, exhibiting good
performance across all evaluated ICL tasks and simultaneously learning sparse
parity and retrieval.

We believe that our findings underscore the importance of broadening the
understanding of ICL beyond Transformers, as significant progress has been
made in the context of attention-free architectures.

We acknowledge that a limitation of our study lies in the focus on non-
language ICL tasks and smaller models. It is possible that an architectural
comparison between SSMs and transformers for more general ICL tasks in actual
language settings at higher parameter counts might not be yield the same
observations as we offer here. Nevertheless, our results indicate that, apart
from its difficulty in some retrieval tasks, similar to those noted by (Arora
et al., 2023), there seems to be no fundamental obstacle for Mamba to perform
in-context learning.

Figure 1: MambaFormer is a hybrid architecture that replaces MLP blocks within
the transformer with Mamba blocks. Importantly, the architecture also starts
with a Mamba block and does not use positional encoding. In our ICL
evaluations, we find that MambaFormer consistently achieves a best-of-both-
worlds performance compared to Transformer and Mamba.

##  2 Related Work

##### Transformer-based in-context Learning.

The role of attention in ICL has been the focus of both theoretical and
empirical research. Studies have primarily focused on meta-learning (Ravi &
Larochelle, 2016; Min et al., 2022a), where one explicitly trains for ICL.
Notably, Garg et al. (2022) have examined transformers in in-context
regression tasks, from learning linear regression to learning decision trees.
Subsequent works have suggested that attention may mimic various optimization
algorithms (Akyürek et al., 2022; von Oswald et al., 2023b; Dai et al., 2023).
In fact, Ahn et al. (2023); Mahankali et al. (2023) have provably shown that
gradient descent is optimal in linear regression ICL for linear attention.

While these settings might appear simplistic and detached from language
models, Bhattamishra et al. (2023) showed that a frozen GPT-2 can implement
the nearest neighbor algorithm, drawing connections between the ICL in
existing language models and the stylized setting of training for ICL from
random initialization. Furthermore, Olsson et al. (2022) also empirically
demonstrate that “induction heads”, which are attention heads that solve a
simple retrieval problem, correlate with ICL behavior, providing a strong
connection between retrieval and ICL.

##### Sub-quadratic architectures.

The number of effective floating point operations in an attention layer scales
quadratically with respect to the input sequence length. Numerous
approximations or alternative model architectures have been proposed to
overcome the quadratic dependence. These range from approximating attention
mechanisms (Beltagy et al., 2020; Wang et al., 2020) to the development of
novel recurrent convolutional models such as structured state-space models (Gu
et al., 2022b).

S4 (Gu et al., 2022a) is a family of sequence models characterized by a
discretized state-space model

$\mathbf{h}_{t}=\overline{\mathbf{A}}\mathbf{h}_{t-1}+\overline{\mathbf{B}}\mathbf{x}_{t},~{}y_{t}=\mathbf{C}\mathbf{h}_{t},$
(1)

where $\mathbf{h}_{t}$ represents the hidden state and
$(\overline{\mathbf{A}},\overline{\mathbf{B}},\mathbf{C})$ are input-
independent (transformed) parameters. The recurrence is expressible as a
convolution, enabling near-linear complexity using Fast Fourier Transform.
Viewed in this framework, Linear Transformers (Katharopoulos et al., 2020),
which employ linear attention without softmax, can be seen as a variant of
linear SSM.

Building upon this concept, H3 (Dao et al., 2022), which integrates an S4 with
dual gated connections. The recent Mamba (Gu & Dao, 2023) departs from the
standard SSM by introducing a selection mechanism that makes
$(\overline{\mathbf{A}},\overline{\mathbf{B}},\mathbf{C})$ in Equation 1
dependent on $\mathbf{x}_{t}$, which allows for input-dependent sequence
mixing.

There are other notable attention-free models such as Hyena (Poli et al.,
2023), RWKV (Peng et al., 2023), RetNet (Sun et al., 2023), and GLA (Yang et
al., 2023b). Despite of state-of-the-art performance for models like Mamba,
Arora et al. (2023) have demonstrated that subquadratic models still lag
behind attention on multi-query recall tasks, which is a generalization of the
induction head task (Olsson et al., 2022).

In Xie et al. (2021), the authors proposed a synthetic language-based in-
context learning dataset and show that transformers and LSTMs are capable of
ICL. Moreover, Akyürek et al. (2024) suggested a langauge based ICL benchmark,
by training on regular languages generated by random finite automata, and also
underscored the gap between Transformers and subquadratic complexity models.

##  3 Experimental Setup

We evaluate the ICL capabilities of SSMs and Transformers by training each
model from scratch on each specific task, detailed in Section 3.1. Section 3.2
outlines the ICL and related tasks investigated in our study. We provide a
brief summary of our tasks in the following Table 2.

Task | dim $(d)$ | points ($N$) | Example/Function Sampling | Task-specific
---|---|---|---|---
Linear regression | 20 | 41 | $\mathbf{x},\mathbf{w}\sim\mathcal{N}(0,{\mathbf{I}}_{d})$ | –
Sparse Linear regression | 20 | 101 |  $\mathbf{x},\mathbf{w}\sim\mathcal{N}(0,{\mathbf{I}}_{d})$, $\texttt{sparsity}(\mathbf{w})\leftarrow k$ | $k=3$
2NN regression | 20 | 101 | $\mathbf{W}^{(1)}_{ij},\mathbf{W}^{(2)}_{ij}\sim\mathcal{N}(0,1)$ | –
Decision Tree | 20 | 101 | $\mathbf{x},\mathrm{Leaf}\sim\mathcal{N}(0,1),\mathrm{non\\_leaf}\sim\\{1,...,d\\}$ | $\mathrm{depth}=4$
Orthogonal-outlier regression | 20 | 101 |  $\mathbf{x},\mathbf{w}\sim\mathcal{N}(0,{\mathbf{I}}_{d})$, $\mathbf{u},\mathbf{v}\sim\mathbf{w}^{\perp}$ | $p=0.5$
Many-outlier regression | 20 | 512 |  $\mathbf{x}\sim\mathcal{N}(0,I)$ w.p. $1-p$, else $(\mathbf{x},y)=(\mathbf{1},1)$ | $p=0.9$
Sparse Parity | 10 | 140 | $\mathbf{x}\sim\\{-1,1\\}^{d},y=\prod_{j\in I}\mathbf{x}[j]$ | $k=2$
Chain-of-Thought I/O | 10 | 101 |  $\mathbf{x}\sim\mathcal{N}(0,{\mathbf{I}}_{d})$, ${\mathbf{W}}_{ij}\sim\mathcal{N}(0,2/k)$, $\mathbf{v}\sim\mathcal{N}(0,{\mathbf{I}}_{k})$ | $h=8$
Vector-valued MQAR | 20 | 128 | $\mathbf{k},\mathbf{v}\sim\mathrm{Unif}(\mathcal{S}^{d-1})$ |  $32$ k-v pairs

Table 2: Summary of Tasks. All models are trained for 500,000 iterations
(except for the vector-valued MQAR; see Section B.2).

###  3.1 Model Training for In-context Learning

We train models to learn specific function classes $\mathcal{F}$ in-context.
Training begins by generating random prompts: selecting a function
$f\in\mathcal{F}$ from distribution $\mathcal{D}_{\mathcal{F}}$ and sampling a
sequence of random inputs
$\mathbf{x}_{1},\ldots,\mathbf{x}_{N}\in\mathbb{R}^{d}$ i.i.d. from
$\mathcal{D}_{\mathcal{X}}$. Here, $N$ and $d$ represent the number of in-
context examples and the dimension of $\mathbf{x}_{i}$, respectively. These
inputs create the prompt
$P=(\mathbf{x}_{1},f(\mathbf{x}_{1}),\ldots,\mathbf{x}_{N},f(\mathbf{x}_{N}))$.
We train the model ${f}_{\theta}$, parameterized by $\theta$, by minimizing
the expected loss over all prompts:

$\min_{\theta}\mathbb{E}_{P}\left[\frac{1}{N}\sum_{i=1}^{N-1}\ell(f_{\theta}(P^{i}),f(\mathbf{x}_{i}))\right],$
(2)

where
$P^{i}:=(\mathbf{x}_{1},f(\mathbf{x}_{1}),\ldots,\mathbf{x}_{i},f(\mathbf{x}_{i}),\mathbf{x}_{i+1})$
and $\ell(\cdot,\cdot)$ is a loss function.  For
$f:\mathbb{R}^{d}\rightarrow\mathbb{R}$, we append $d-1$ zeros to
$f(\mathbf{x})$. We use appropriate loss functions for each ICL task.

##### Model architecture.

We primarily focus on SSMs, including (1) Mamba (Gu & Dao, 2023), a state-of-
the-art SSM model with selection mechanism; (2) S4 (Gu et al., 2022a), a
linear time-invariant counterpart to Mamba; and (3) S4-Mamba, a variant where
Mamba’s input-dependent S6 is replaced with input-independent S4. The primary
differences between the S4 models lie in the application of multiplicative
gating and the module order.111https://github.com/state-
spaces/s4/blob/main/models/s4

##### Training.

We train each model by sampling a batch of random prompts at each training
step and updating the model parameters using Adam optimizer (Kingma & Ba,
2014). We use a batch size of 64 and trained for 500,000 iterations (except
for the vector-valued MQAR task; see Section B.2).

##### Evaluation.

We evaluate model performance on in-context learning using task and data
distributions $\mathcal{D}_{\mathcal{F}}$ and $\mathcal{D}_{\mathcal{X}}$
consistent with training. A function and a sequence of $N$ inputs are sampled
from $\mathcal{D}_{\mathcal{F}}$ and $\mathcal{D}_{\mathcal{X}}$,
respectively, to generate a test prompt $P_{\sf
test}=(\mathbf{x}_{1},f(\mathbf{x}_{1}),\ldots,\mathbf{x}_{N},f(\mathbf{x}_{N}))$.
We create 1,280 prompts and measure the empirical mean of Eq. (2) across them
for in-context learning performance.

To plot performance as model capacity grows, we calculate the total floating
point operations (FLOPs) used for training. The calculation for Transformer
and Mamba can be found in Appendix C, which are based on (Kaplan et al., 2020;
Gu & Dao, 2023). Model configurations and training implementation details are
provided in Appendix A.

###  3.2 In-context learning tasks

We provide an overview of the ICL and related tasks investigated in this
study. Some tasks are adapted from (Garg et al., 2022), and we follow the
settings outlined in their work. The tasks are summarized in Table 2.

####  3.2.1 Learning regression

For all regression tasks, in-context examples $\mathbf{x}_{i}$ are sampled
from the Gaussian distribution $\mathcal{N}(0,\mathbf{I}_{d})$, where
$\mathbf{I}_{d}$ is the $d\times d$ identity matrix. We use the squared error
loss for model training.

##### Linear regression.

We examine the class of linear functions
$\mathcal{F}=\\{f|f(\mathbf{x})=\mathbf{w}^{\top}\mathbf{x},\mathbf{w}\in\mathbb{R}^{d}\\}$
where $\mathbf{w}$ is sampled from the Gaussian distribution
$\mathcal{N}(0,\mathbf{I}_{d})$. We set $d=20$.

##### Sparse linear regression.

The setting is identical to linear regression, except that $\mathbf{w}$ is
sampled from $\mathcal{N}(0,\mathbf{I}_{d})$, after which $k$ coordinates are
randomly retained in $\mathbf{w}$, and the rest are set to zero. We set $k=3$.

##### Two-layer neural network.

We consider the class of two-layer ReLU neural networks
$\mathcal{F}=\\{f|f(\mathbf{x})=\mathbf{W}^{(2)}\sigma\left(\mathbf{W}^{(1)}\mathbf{x}\right)\\}$,
where $\mathbf{W}^{(2)}\in\mathbb{R}^{1\times
h},\mathbf{W}^{(1)}\in\mathbb{R}^{h\times d}$, and
$\sigma(\cdot)=\max(0,\cdot)$ is the ReLU function. Each element of the weight
matrices is independently drawn from $\mathcal{N}(0,1)$. We use $d=20$ and
$h=100$.

##### Decision Tree

We consider a full binary tree with a fixed depth and input
$\mathbf{x}\in\mathbb{R}^{d}$. Leaf node values are sampled from
$\mathcal{N}(0,1)$, and the rest are sampled uniformly from $\\{1,...,d\\}$,
functioning as indices of $\mathbf{x}$.  At a given non-leaf node, we move to
the right if $x[i]>0$, where i is the sampled index, and otherwise move to the
left. $y$ is the leaf node value when the traversal terminates.

####  3.2.2 Learning with outliers

The problems that belong to this family adopt the basic setting of the
standard linear regression task. With a fixed probability $p$, each pair of
$({\mathbf{x}}_{i},f({\mathbf{x}}_{i}))$ in the prompt is replaced with
“dummy” vectors which are either out of the training distribution, or
confounders designed to increase the complexity of the task. We test
$p\in\\{0.5,0.9\\}$ as replacement probabilities for tasks described below.
During training, we do not compute the loss for the replaced outliers.

##### Orthogonal-outlier regression.

Each pair of $(\mathbf{x}_{i},f(\mathbf{x}_{i}))$ is randomly replaced with
$((a_{x}\mathbf{u}+b_{x}\mathbf{v})/(a_{x}^{2}+b_{x}^{2}),(a_{y}\mathbf{u}+b_{y}\mathbf{v})/(a_{x}^{2}+b_{x}^{2}))$,
where $\mathbf{u},\mathbf{v}\in\mathbf{w}^{\perp}$.
$(\mathbf{u},\mathbf{v}):=(\mathbf{w}_{1}-\text{proj}_{\mathbf{w}}(\mathbf{w}_{1}),\mathbf{w}_{2}-\text{proj}_{\mathbf{w}}(\mathbf{w}_{2}))$
and $\mathbf{w}_{1}$ and $\mathbf{w}_{2}$ are sampled from
$\mathcal{N}(0,\mathbf{I}_{d})$ and the coefficients $a_{x},b_{x},a_{y},b_{y}$
are independently sampled from $\mathcal{N}(0,1)$.

##### Many-outlier regression.

In this setting, ${\mathbf{x}}_{i}$ and $f({\mathbf{x}}_{i})$ are randomly
replaced with a $d$-dimensional vector of ones $\\{1\\}^{d}$ and an one-hot
vector $[1,0,\ldots,0]$, respectively, with probability 90%.  Here, we test
longer sequences of $N=512$.

####  3.2.3 Learning discrete functions

##### Sparse parity.

Following the setting from Bhattamishra et al. (2023), we consider the class
of functions
$\mathcal{F}=\\{f|f(\mathbf{x})=\prod_{j\in\mathcal{S}}\mathbf{x}_{i}[j]\\}$,
where $\mathbf{x}_{i}[j]$ denotes the $j$-th element of the vector
$\mathbf{x}_{i}$ and $\mathcal{S}$ is a subset of $\\{1,\ldots,d\\}$ with the
size $k$. Each $\mathbf{x}_{i}$ is sampled uniformly at random from
$\\{-1,1\\}^{d}$, and $\mathcal{S}$ of size $k$ is randomly sampled from the
set $\\{1,\ldots,d\\}$. For this task, we train a model using the cross-
entropy loss and evaluate the model using a binary indicator for accuracy,
which assigns 1 to correct predictions and 0 to incorrect ones.

####  3.2.4 Learning Chain-of-Thought

##### Chain-of-Thought-I/O.

Following the setting from Li et al. (2023b), we consider the class of two-
layer ReLU neural networks
$\mathcal{F}=\\{f|f(\mathbf{x})=\mathbf{W}^{(2)}\sigma\left(\mathbf{W}^{(1)}\mathbf{x}\right)\\}$,
where $\mathbf{W}^{(2)}\in\mathbb{R}^{1\times
h},\mathbf{W}^{(1)}\in\mathbb{R}^{h\times d}$, and $\sigma(\cdot)$ is the
$\mathrm{ReLU}$ function. We set $d=10$, and $h=8$. We additionally interleave
the intermediate hidden feature
$\mathbf{s}_{i}=\sigma\left(\mathbf{W}^{(1)}\mathbf{x}_{i}\right)$ in our
input training sequence in a Chain-of-Thought style. Given the input sequence
$(\mathbf{x}_{1},\mathbf{s}_{1},f(\mathbf{x}_{1}),\cdots,\mathbf{x}_{N},\mathbf{s}_{N},f(\mathbf{x}_{N}),\mathbf{x}_{\text{test}})$,
the model is evaluated on the final output prediction $\hat{\mathbf{y}}$ based
on the input sequence and the intermediate layer prediction
$\hat{\mathbf{s}}_{\text{test}}$.

####  3.2.5 Learning retrieval

Vector-valued multi-query associative recall

We test the model’s ability to do multi-query associative recall (MQAR) (Arora
et al., 2023). While MQAR is not an ICL task, model’s ability to do
associative recall (AR) is highly related to model’s ability to learn in-
context (Olsson et al., 2022). To better measure the model’s ability to
retrieve information from context, we consider a variant of MQAR such that
keys and values are vector-valued so each vector can be seen as a “unique
token” and the retrieval accuracy can be measured by the mean squared error
between retrieved vectors and target vectors. Specifically, in this task, the
model is given a sequence of key-value pairs of vectors
$\\{\mathbf{k}_{1},\mathbf{v}_{1},...,\mathbf{k}_{n},\mathbf{v}_{n}\\}$, where
$\mathbf{k}_{i},\mathbf{v}_{i}\in\mathcal{S}^{d-1}$ are sampled uniformly from
the unit $d$-sphere. The query consists of sequence of vectors
$\\{\mathbf{q}_{1},...,\mathbf{q}_{m}\\}$. For each query $\mathbf{q}_{j}$,
there exists some $1\leq l\leq n$ such that $\mathbf{q}_{j}=\mathbf{k}_{l}$.
The model must learn to output $\mathbf{v}_{l}$ associated with the query
$\mathbf{q}_{j}$ for each of the queries, producing $m$ outputs total. We
train a model using the squared error error.

##  4 Experiment results

Figure 2: A suite of ICL tasks ran for Transformer, Mamba, S4, and hybrid
architectures where each color represents a different architecture.  More
transparent points indicate earlier stages of training; plotted models are
trained in between [100k, 500k] iterations.  We represent each model in terms
of its number of floating point operations (FLOPs) used for training.

In this section, we demonstrate that Mamba can be trained from scratch to
perform various ICL tasks. Furthermore, we identify specific tasks in which
one model performs better than the other and vice versa.

###  4.1 Mamba can in-context learn!

As shown in Figure 2, Mamba consistently outperforms its more simple
counterparts S4-Mamba and S4. In simple tasks such as linear regression, the
gap between Mamba and S4-Mamba is much smaller than that of S4-Mamba and S4.
Given that the main difference between Mamba and S4-Mamba is the input-
dependent selection mechanism, appropriate gating and stacking of MLPs (_i.e._
, the difference between S4-Mamba and S4) seem to be more significant for such
tasks. However, in comparison, input-dependent selection makes meaningful
progress for more complex tasks such as 2NN regression and learning decision
trees.

Mamba can also perform on par with Transformer even as the total FLOPs scale
up. This is surprising given that Transformer and attention have been the
focus of many previous works for its unique ICL capability. Moreover, Mamba
tends to perform better in smaller parameter settings when controlling for
equal depth, _i.e._ , keeping the number of attention, MLP, and Mamba blocks
equivalent.

###  4.2 Performance gaps in more complex ICL tasks

We also consider a family of more complex ICL tasks, namely learning decision
tree, sparse parity, and Chain-of-Thought (Figures 2 and 4). The figure shows
that Transformers can solve Decision Tree and Vector-valued MQAR, while Mamba
cannot. In Sparse Parity task of Figure 5, however, Transformer is unable to
learn the function family while Mamba can.

Figure 3: Scaling curves with respect to the total number of FLOPs on
Orthogonal-outlier Regression Task.

##### Filtering outliers in regression.

Orthogonal-outlier regression and many-outlier regression, like other outlier
tasks, focus on the model’s ability to learn to ignore dummy vectors, either
by the fact that the $\mathbf{x}_{i}\in\mathbf{w}^{\perp}$, or by the fact
that $\mathbf{y}_{i}$ is a vector instead of a zero-padded scalar value. This
explicitly requires the models to look at the previous input sequences, and
discover the properties that distinguish the dummy vectors from training
examples while learning the class of functions the training prompt represents.

For orthogonal-outlier regression task with a relatively short sequence length
of 101 (see Table 2 for task descriptions), Mamba performs on par with
Transformer, as seen in Figure 3. Interestingly, for many-outlier regression
where we test on a sequence length of 512 and 90% all-ones replacement, Mamba
significantly outperforms Transformers. This is also in line with what Gu &
Dao (2023) report, in which Mamba fares better for the induction task for long
sequence lengths. These two results indicate that Mamba has no significant
issue with filtering out unnecessary information, while retaining the ability
to learn linear regression in-context.

##### Chain-of-Thought I/O.

Figure 4: Performance of Transformer and Mamba models on the Chain-of-
Thought-I/O task. Experiments on varying the model size (left) and varying the
hidden dimension (right).

Figure 4 shows that Mamba models are capable of in-context learning in a
chain-of-thought manner, performing comparably to Transformer models across
the tested configurations. In smaller model configurations, Mamba models
exhibit superior performance compared to Transformer models. However, as model
size increases, Transformer models begin to surpass Mamba models. The
performance of Transformer models remains relatively stable across different
problem sizes, while Mamba models’ performance is significantly influenced by
the size of the hidden layer. Specifically, Mamba models excel over
Transformer models at smaller problem sizes (i.e., smaller hidden dimensions),
but their advantage diminishes as the problem size expands.

###  4.3 Challenges in parity and retrieval

Embedding dimension ($d$) | 64 | 128
---|---|---
Mamba | 7.23e-1 | 1.50e-1
6 MambaBlocks + 1 Standard Hybrid | 1.54e-3 | 5.86e-5
Transformer w/o PE | 7.61e-5 | 5.55e-5
Transformer w/ PE | 3.99e-5 | 2.46e-7
MambaFormer | 1.03e-5 | 3.79e-7
(a) $32$ key-value pairs with $16$ queries.

Embedding dimension ($d$) | 64 | 128
---|---|---
Mamba | 8.64e-1 | 1.64e-1
6 MambaBlocks + 1 Standard Hybrid | 1.99e-2 | 1.37e-2
Transformer w/o PE | 1.14e-3 | 8.66e-5
Transformer w/ PE | 5.17e-6 | 8.76e-7
MambaFormer | 7.30e-6 | 3.37e-6
(b) $32$ key-value pairs with $4$ queries.

Table 3: Test loss (mean squared error) on vector-valued MQAR and respective
model configurations. We test for both Transformer with Positional Encoding
(PE) and without. All models have $4$ layers with roughly the same number of
parameters. An exception is “$6$ MambaBlocks + $1$ Standard Hybrid” model, but
we still consider it as (equivalently) having $4$ layers since $6$ MambaBlocks
are equivalent to $3$ Mamba layers as described in Figure 6.

We run vector-valued MQAR on two settings: (1) $32$ key-value pairs with $16$
queries and (2) $32$ key-value pairs with $4$ queries. From Table 3, we can
see that Mamba fails to retrieve vectors accurately as the mean squared error
for retrieving normed vectors are greater than $0.1$ in all cases.

As a sidenote, all models trained with $16$ queries have lower test loss than
models trained with $4$ queries. A possible explanation is that, for a single
sequence of data that represents an MQAR task, we can think of each
$(\mathbf{q},\mathbf{v})$ pair as a “training sample”, so a sequence with $16$
queries contains more “training samples” than that of a sequence with $4$
queries. This also shows that having more queries does not necessarily make
the task harder.

Figure 5: Although Transformer fails to converge, Mamba and S4-Mamba can learn
sparse parity of $d=10$ and $k=2$. Each model is trained with  an embedding
dimension of 256 and depth of 12 layers (approximately 10 million parameters)
up to 500k iterations. Transformer failed to learn even up to an embedding
dimension of 768 and 24 layers.

While Mamba fails on simple retrieval tasks such as MQAR, the tables turn for
the task of learning sparse parity (Figure 5). Transformer fails to do better
than random guessing, in line with the empirical evidence of Bhattamishra et
al. (2023). We confirm this is the case for Transformer sizes of embedding
dimensions up to 768 and up to 24 layers when trained for at most 1 million
iterations. However, Mamba succeeds in this task with ease, solving sparse
parity for $(d,k)=(10,2)$ with a network as small as 2 layers. Even more
surprisingly, S4-Mamba is able to solve parity as well; this may mean that
proper convolution or gating may be more important than input-dependent
selection.  Our result hints at that the initial (causal) convolution that
Mamba provides before the attention layer may be crucial to solving parities,
a similar phenomenon observed for Vision Transformers in computer vision tasks
(Yu et al., 2022).

It is known that any algorithm for learning parities requires either a super-
linear memory of $\omega(d)$ or a super-polynomial number of samples in $d$
(Raz, 2016; Kol et al., 2017). While Transformer is known to have better
memory due to its quadratic attention mechanism, our results on learning
sparse parities brings forth the question on how different architectures may
utilize its memory differently in terms of function approximation. We leave
the theoretical and empirical question of which architectural component allows
for learning parities as an avenue for further study.

##  5 The Advantage of Hybrid Architectures for In-context Learning

Figure 6: Model Architectures. (a) and (b) denote the standard Transformer and
Mamba architectures. (c) denotes the hybrid architecture of Mamba and
Attention blocks, following the design proposed in Gu & Dao (2023). (d)
demonstrates the proposed architecture, namely MambaFormer, which replaces the
Positional Encoding with a Mamba block. For convenience, we denote 2 blocks of
either Mamba, Multi-head Attention, or a Feed Forward Network as 1 layer.

In the previous section, we have observed that Transformers perform better
than SSMs in some tasks, such as learning decision trees or retrieval, while
SSMs excel in others, such as learning sparse parities or learning heavy-
outlier linear regression, possibly due to its recurrent nature. However, can
we achieve the best of both worlds without sacrificing performance in our
suite of ICL tasks?

We answer this in the affirmative; that we can indeed reach competitive
performance in our suite of ICL tasks, achieving performance comparable to
that of Transformers and Mamba, while simultaneously excelling in specific
tasks that either fail in. We can achieve strong performance by interleaving
Attention and Mamba, where a key ingredient is having Mamba as the first
layer.

In this section, we investigate two hybrid architectures that combine
Transformer and Mamba, namely Standard Hybrid and MambaFormer as illustrated
in Figure 6. Standard Hybrid is the architecture of interleaving MHA and Mamba
by replacing the MLP block with Mamba. MambaFormer is nearly identical to
Standard Hybrid but with an additional Mamba block as its initial layer and no
particular positional encoding. Although many works have found that
interleaving multi-head attention and LTI SSMs beneficial (Zuo et al., 2022;
Mehta et al., 2022; Pilault et al., 2023), interestingly Gu & Dao (2023) have
not found significant benefits of interleaving. In the following results, we
show that interleaving with Mamba as its initial layer can help solve both
sparse parity and retrieval, each task unsolvable by Mamba and Transformer.

###  5.1 Simultaneously learning parities and retrieval

Figure 7: Median convergence time of learning parity over 5 random seeds for
max. 500k iterations, where 500k convergence time signifies failed learning.
Having the initial layer as Mamba is essential for efficiently learning
parities. Tested model configurations are specified in Appendix A.

As highlighted in Bhattamishra et al. (2023); Barak et al. (2022), learning
sparse parity in-context seems to be difficult for Transformer and some SSMs
like Hyena. Yet interestingly, as seen in Figure 7, MambaFormer successfully
learns parity as quickly as Mamba in terms of sample complexity. While the
Standard Hybrid model is also capable, it exhibits much worse sample
efficiency.

We perform an ablation study by equipping Transformer with an initial Mamba
block without any positional encoding. Although this variant Transformer only
has fewer Mamba blocks than Standard Hybrid, it solves parity almost as
efficiently as Mamba. Not only does this show us that order of layers in
interleaving matter, as shown in Press et al. (2022), but also that Mamba can
complement Transformer without hurting performance in ICL. This result brings
up intriguing difference between the function learning capabilities of
Attention and Mamba; we leave this question up for further study.

##### Closing the gap in retrieval.

The gap between Mamba and Transformer in vector-valued MQAR task is largely
due to the fact that Mamba (as an SSM) compresses context into smaller states
when generating output, while the Attention mechanism in Transformer does not
compress the context. The amount of information about the context Mamba has at
each state depends on the dimension of hidden state (as the hidden states
capture the important information in the context) and it is challenging if the
task is to accurately retrieve a specific part of the context by a query that
is placed after the context.

To close the gap in the vector-valued MQAR task between Mamba and Transformer,
and without sacrificing too much of the efficiency, we add one attention layer
within layers of Mamba blocks. In particular, in a Mamba model of $4$ layers
($8$ Mamba blocks stacked homogeneously), we replace the middle two blocks
with Standard Hybrid (w/o positional embedding). As shown in Table 3, Mamba
model gains a significant improvement in vector-valued MQAR by having one
Standard Hybrid. We further test MambaFormer on the same task and find that
MambaFormer almost entirely closes the gap to transformer in vector-valued
MQAR task.

###  5.2 All-in-one ICL performance

While MambaFormer succeeds in two tasks that were either deemed difficult for
Mamba for Transformer, it also performs equally as well as Transformer and
Mamba in the rest of our suite of ICL tasks. In Figure 2, we see that
MambaFormer and Standard Hybrid both learn decision trees as well as
Transformer, even at larger parameter sizes. Even more surprisingly,
MambaFormer efficiently learns linear regression better than both models even
in the presence of 90% noisy data in Many-outlier regression, as a MambaFormer
trained on 100k iterations ($<10^{17}$ FLOPs) performs as well as models
trained with 10 times the number of FLOPs.

In conclusion, we find the best of both worlds within our diverse array of ICL
tasks; a hybrid architecture that can solve as difficult problems as retrieval
and parity, while performing on par with Transformer and Mamba in other ICL
tasks. Given our results, it will be interesting to see how hybrid
architectures perform in other kinds of ICL tasks, such as those discussed in
(Xie et al., 2021; Akyürek et al., 2024).

##  6 Discussion

In this work, we have provided a comprehensive investigation of in-context
learning with state-space models (SSMs) and contrasted them with the
transformer architecture. Our study has revealed that SSMs, especially Mamba,
are capable in-context learners. On the other hand, our evaluations revealed
that neither SSMs nor transformers are great at all tasks, specifically, SSMs
struggle with decision tree learning and retrieval tasks whereas transformers
struggle with sparse parity. This has led us to the hybrid architecture
MambaFormer which achieves a best-of-both-worlds performance on our ICL suite.

Future research directions include exploring (1) how performance on our ICL
suite correlates with general language modeling capabilities, such as
perplexity on standard NLP benchmarks, (2) the potential for developing more
effective architectures by integrating elements from transformers, SSMs, and
gating mechanisms, (3) identifying architectural features that contribute to
effective in-context learning, and (4) assessing the impact of MambaFormer and
other innovative architectures on language modeling performance.

## References

  * Ahn et al. (2023) Ahn, K., Cheng, X., Daneshmand, H., and Sra, S.  Transformers learn to implement preconditioned gradient descent for in-context learning.  _arXiv preprint arXiv:2306.00297_ , 2023.
  * Akyürek et al. (2022) Akyürek, E., Schuurmans, D., Andreas, J., Ma, T., and Zhou, D.  What learning algorithm is in-context learning? investigations with linear models.  In _The Eleventh International Conference on Learning Representations_ , 2022.
  * Akyürek et al. (2024) Akyürek, E., Wang, B., Kim, Y., and Andreas, J.  In-context language learning: Architectures and algorithms.  _arXiv preprint arXiv:2401.12973_ , 2024.  URL https://arxiv.org/abs/2401.12973.
  * Arora et al. (2023) Arora, S., Eyuboglu, S., Timalsina, A., Johnson, I., Poli, M., Zou, J., Rudra, A., and Ré, C.  Zoology: Measuring and improving recall in efficient language models.  _arXiv preprint arXiv:2312.04927_ , 2023.  URL https://arxiv.org/abs/2312.04927.
  * Bai et al. (2023) Bai, Y., Chen, F., Wang, H., Xiong, C., and Mei, S.  Transformers as statisticians: Provable in-context learning with in-context algorithm selection.  _arXiv preprint arXiv:2306.04637_ , 2023.
  * Barak et al. (2022) Barak, B., Edelman, B., Goel, S., Kakade, S., Malach, E., and Zhang, C.  Hidden progress in deep learning: Sgd learns parities near the computational limit.  _Advances in Neural Information Processing Systems_ , 35:21750–21764, 2022.
  * Beltagy et al. (2020) Beltagy, I., Peters, M. E., and Cohan, A.  Longformer: The long-document transformer.  _arXiv preprint arXiv:2004.05150_ , 2020.
  * Bhattamishra et al. (2023) Bhattamishra, S., Patel, A., Blunsom, P., and Kanade, V.  Understanding in-context learning in transformers and llms by learning to learn discrete functions.  _arXiv preprint arXiv:2310.03016_ , 2023.  URL https://arxiv.org/abs/2310.03016.
  * Brown et al. (2020) Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J. D., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., et al.  Language models are few-shot learners.  _Advances in neural information processing systems_ , 33:1877–1901, 2020.
  * Chan et al. (2022) Chan, S., Santoro, A., Lampinen, A., Wang, J., Singh, A., Richemond, P., McClelland, J., and Hill, F.  Data distributional properties drive emergent in-context learning in transformers.  _Advances in Neural Information Processing Systems_ , 35:18878–18891, 2022.
  * Dai et al. (2023) Dai, D., Sun, Y., Dong, L., Hao, Y., Ma, S., Sui, Z., and Wei, F.  Why can gpt learn in-context? language models secretly perform gradient descent as meta-optimizers.  In _Findings of the Association for Computational Linguistics: ACL 2023_ , pp. 4005–4019, 2023.
  * Dao et al. (2022) Dao, T., Fu, D. Y., Saab, K. K., Thomas, A. W., Rudra, A., and Ré, C.  Hungry hungry hippos: Towards language modeling with state space models.  _arXiv preprint arXiv:2212.14052_ , 2022.  URL https://arxiv.org/abs/2212.14052.
  * Gao et al. (2020) Gao, L., Biderman, S., Black, S., Golding, L., Hoppe, T., Foster, C., Phang, J., He, H., Thite, A., Nabeshima, N., Presser, S., and Leahy, C.  The pile: An 800gb dataset of diverse text for language modeling, 2020\.
  * Garg et al. (2022) Garg, S., Tsipras, D., Liang, P. S., and Valiant, G.  What can transformers learn in-context? a case study of simple function classes.  _Advances in Neural Information Processing Systems_ , 35:30583–30598, 2022.
  * Gu & Dao (2023) Gu, A. and Dao, T.  Mamba: Linear-time sequence modeling with selective state spaces.  _arXiv preprint arXiv:2312.00752_ , 2023.  URL https://arxiv.org/abs/2312.00752.
  * Gu et al. (2022a) Gu, A., Goel, K., Gupta, A., and Ré, C.  On the parameterization and initialization of diagonal state space models.  In _Advances in Neural Information Processing Systems_ , volume 35, pp. 35971–35983, 2022a.
  * Gu et al. (2022b) Gu, A., Goel, K., and Ré, C.  Efficiently modeling long sequences with structured state spaces.  In _The Tenth International Conference on Learning Representations, ICLR 2022_ , 2022b.  URL https://openreview.net/forum?id=uYLFoz1vlAC.
  * Kaplan et al. (2020) Kaplan, J., McCandlish, S., Henighan, T., Brown, T. B., Chess, B., Child, R., Gray, S., Radford, A., Wu, J., and Amodei, D.  Scaling laws for neural language models.  _arXiv preprint arXiv:2001.08361_ , 2020.
  * Katharopoulos et al. (2020) Katharopoulos, A., Vyas, A., Pappas, N., and Fleuret, F.  Transformers are rnns: Fast autoregressive transformers with linear attention.  In _Proceedings of the 37th International Conference on Machine Learning, ICML 2020_ , volume 119 of _Proceedings of Machine Learning Research_ , pp. 5156–5165, 2020.  URL http://proceedings.mlr.press/v119/katharopoulos20a.html.
  * Kingma & Ba (2014) Kingma, D. P. and Ba, J.  Adam: A method for stochastic optimization.  _arXiv preprint arXiv:1412.6980_ , 2014.
  * Kol et al. (2017) Kol, G., Raz, R., and Tal, A.  Time-space hardness of learning sparse parities.  In _Proceedings of the 49th Annual ACM SIGACT Symposium on Theory of Computing_ , pp. 1067–1080, 2017.
  * Li et al. (2023a) Li, Y., Ildiz, M. E., Papailiopoulos, D., and Oymak, S.  Transformers as algorithms: Generalization and stability in in-context learning, 2023a.
  * Li et al. (2023b) Li, Y., Sreenivasan, K., Giannou, A., Papailiopoulos, D., and Oymak, S.  Dissecting chain-of-thought: Compositionality through in-context filtering and learning.  In _Thirty-seventh Conference on Neural Information Processing Systems_ , 2023b.
  * Mahankali et al. (2023) Mahankali, A., Hashimoto, T. B., and Ma, T.  One step of gradient descent is provably the optimal in-context learner with one layer of linear self-attention.  _arXiv preprint arXiv:2307.03576_ , 2023.
  * Mehta et al. (2022) Mehta, H., Gupta, A., Cutkosky, A., and Neyshabur, B.  Long range language modeling via gated state spaces.  _arXiv preprint arXiv:2206.13947_ , 2022.
  * Min et al. (2022a) Min, S., Lewis, M., Zettlemoyer, L., and Hajishirzi, H.  Metaicl: Learning to learn in context.  In _Proceedings of the 2022 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies_ , pp. 2791–2809, 2022a.
  * Min et al. (2022b) Min, S., Lyu, X., Holtzman, A., Artetxe, M., Lewis, M., Hajishirzi, H., and Zettlemoyer, L.  Rethinking the role of demonstrations: What makes in-context learning work?  _arXiv preprint arXiv:2202.12837_ , 2022b.  URL https://arxiv.org/abs/2202.12837.
  * Muennighoff et al. (2023) Muennighoff, N., Rush, A. M., Barak, B., Scao, T. L., Tazi, N., Piktus, A., Pyysalo, S., Wolf, T., and Raffel, C.  Scaling data-constrained language models.  In _Thirty-seventh Conference on Neural Information Processing Systems_ , 2023.  URL https://openreview.net/forum?id=j5BuTrEj35.
  * Olsson et al. (2022) Olsson, C., Elhage, N., Nanda, N., Joseph, N., DasSarma, N., Henighan, T., Mann, B., Askell, A., Bai, Y., and et al., A. C.  In-context learning and induction heads.  _arXiv preprint arXiv:2209.11895_ , 2022.  URL https://arxiv.org/abs/2209.11895.
  * Peng et al. (2023) Peng, B., Alcaide, E., Anthony, Q., Albalak, A., Arcadinho, S., Cao, H., Cheng, X., Chung, M., Grella, M., Kiran, K. G., et al.  Rwkv: Reinventing rnns for the transformer era.  _arXiv preprint arXiv:2305.13048_ , 2023.  URL https://arxiv.org/abs/2305.13048.
  * Pilault et al. (2023) Pilault, J., Fathi, M., Firat, O., Pal, C., Bacon, P.-L., and Goroshin, R.  Block-state transformers.  In _Thirty-seventh Conference on Neural Information Processing Systems_ , 2023.
  * Poli et al. (2023) Poli, M., Massaroli, S., Nguyen, E., Fu, D. Y., Dao, T., Baccus, S., Bengio, Y., Ermon, S., and Ré, C.  Hyena hierarchy: Towards larger convolutional language models.  _arXiv preprint arXiv:2302.10866_ , 2023.  URL https://arxiv.org/abs/2302.10866.
  * Press et al. (2022) Press, O., Zhang, M., Min, S., Schmidt, L., Smith, N. A., and Lewis, M.  Measuring and narrowing the compositionality gap in language models.  _arXiv preprint arXiv:2210.03350_ , 2022.
  * Radford et al. (2019) Radford, A., Wu, J., Child, R., Luan, D., Amodei, D., Sutskever, I., et al.  Language models are unsupervised multitask learners.  _OpenAI blog_ , 1(8):9, 2019.
  * Ravi & Larochelle (2016) Ravi, S. and Larochelle, H.  Optimization as a model for few-shot learning.  In _International conference on learning representations_ , 2016.
  * Raz (2016) Raz, R.  Fast learning requires good memory: A time-space lower bound for parity learning.  In _2016 IEEE 57th Annual Symposium on Foundations of Computer Science (FOCS)_ , pp. 266–275. IEEE, 2016.
  * Schaeffer et al. (2023) Schaeffer, R., Miranda, B., and Koyejo, S.  Are emergent abilities of large language models a mirage?  In _Thirty-seventh Conference on Neural Information Processing Systems_ , 2023.  URL https://openreview.net/forum?id=ITw9edRDlD.
  * Schmidhuber et al. (1997) Schmidhuber, J., Zhao, J., and Wiering, M.  Shifting inductive bias with success-story algorithm, adaptive levin search, and incremental self-improvement.  _Machine Learning_ , 28:105–130, 1997.
  * Sun et al. (2023) Sun, Y., Dong, L., Huang, S., Ma, S., Xia, Y., Xue, J., Wang, J., and Wei, F.  Retentive network: A successor to transformer for large language models.  _arXiv preprint arXiv:2307.08621_ , 2023.  URL https://arxiv.org/abs/2307.08621.
  * Vaswani et al. (2017) Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L., and Polosukhin, I.  Attention is all you need.  In _Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017_ , pp. 5998–6008, 2017\.  URL https://proceedings.neurips.cc/paper/2017/hash/3f5ee243547dee91fbd053c1c4a845aa-Abstract.html.
  * von Oswald et al. (2023a) von Oswald, J., Niklasson, E., Randazzo, E., Sacramento, J., Mordvintsev, A., Zhmoginov, A., and Vladymyrov, M.  Transformers learn in-context by gradient descent.  In _International Conference on Machine Learning_ , pp. 35151–35174. PMLR, 2023a.
  * von Oswald et al. (2023b) von Oswald, J., Niklasson, E., Schlegel, M., Kobayashi, S., Zucchet, N., Scherrer, N., Miller, N., Sandler, M., Vladymyrov, M., Pascanu, R., et al.  Uncovering mesa-optimization algorithms in transformers.  _arXiv preprint arXiv:2309.05858_ , 2023b.  URL https://arxiv.org/abs/2309.05858.
  * Wang et al. (2020) Wang, S., Li, B. Z., Khabsa, M., Fang, H., and Ma, H.  Linformer: Self-attention with linear complexity.  _arXiv preprint arXiv:2006.04768_ , 2020.
  * Wei et al. (2022) Wei, J., Tay, Y., Bommasani, R., Raffel, C., Zoph, B., Borgeaud, S., Yogatama, D., Bosma, M., Zhou, D., Metzler, D., et al.  Emergent abilities of large language models.  _arXiv preprint arXiv:2206.07682_ , 2022.
  * Xie et al. (2021) Xie, S. M., Raghunathan, A., Liang, P., and Ma, T.  An explanation of in-context learning as implicit bayesian inference.  In _International Conference on Learning Representations_ , 2021.
  * Yang et al. (2023a) Yang, L., Lee, K., Nowak, R., and Papailiopoulos, D.  Looped transformers are better at learning learning algorithms, 2023a.
  * Yang et al. (2023b) Yang, S., Wang, B., Shen, Y., Panda, R., and Kim, Y.  Gated linear attention transformers with hardware-efficient training.  _arXiv preprint arXiv:2312.06635_ , 2023b.
  * Yu et al. (2022) Yu, W., Luo, M., Zhou, P., Si, C., Zhou, Y., Wang, X., Feng, J., and Yan, S.  Metaformer is actually what you need for vision.  In _Proceedings of the IEEE/CVF conference on computer vision and pattern recognition_ , pp. 10819–10829, 2022.
  * Zuo et al. (2022) Zuo, S., Liu, X., Jiao, J., Charles, D., Manavoglu, E., Zhao, T., and Gao, J.  Efficient long sequence modeling via state space augmented transformer.  _arXiv preprint arXiv:2212.08136_ , 2022.

##  Appendix A Experimental Setup

###  A.1 Model Architectures

We focus on decoder-only Transformer models, particularly those from the GPT-2
family (Radford et al., 2019), Mamba (Gu & Dao, 2023), and their Hybrid
variants, including Standard and MambaFormer configurations. These models are
evaluated across a range of sizes, as detailed in Table 4. Transformer layers
consist of a Multi-Head Attention (MHA) block followed by a Multilayer
Perceptron (MLP) block. Mamba models consist of two Mamba blocks per layer.
The Hybrid variants merge these approaches, combining a single MHA block with
a Mamba block. For MHA blocks, we use 8 number of heads. Refer to Figure 6 for
a visualization of the architectures considered.

###  A.2 Model Training

We train all of our models on A100-SXM4-40GB GPUs for 500,000 training steps
on all tasks. We use Adam optimizer Kingma & Ba (2014) with a fixed learning
rate. The default value is set to $0.0001$, following the default learning
rate in Garg et al. (2022), and search various learning rates in
$\\{5e-5,1e-4,2e-4,4e-4\\}$. We observe that the training procedure is the
most sensitive to choosing the right learning rate. In particular, as the
number of parameters of the models increases, the training procedure is prone
to gradient explosions, especially in Mamba and hybrid architecutres. Hence,
we clip the gradient norm, with values in $\\{5.0,10.0,50.0\\}$.

As for the train and test data, we fix the dimension of $\mathbf{x}$ to be
$20$, and fix the batch size to be $64$. As suggested in Garg et al. (2022),
we also observe that curriculum is crucial in certain ICL tasks. We adopt a
curriculum of 15 steps every 2000 steps both on the dimension of $\mathbf{x}$
and the number of points (half the length of the training prompt).

Table 4:  Different configurations of Models.  | # layers | embed dim
---|---|---
Standard | 12 | 768
Small | 8 | 512
X-small | 4 | 256
XX-small | 2 | 128

##  Appendix B Implementation Details

This section further elaborates on the task descriptions from Section 3.

###  B.1 Chain-of-Thought-I/O results

Table 5 presents the configurations for the Chain-of-Thought-I/O task using a
2-layer ReLU neural network, following the setup described by Li et al.
(2023b). In the model scale experiment, the input dimension $d=10$ and hidden
layer dimension $k=8$ are held constant while varying the model scale.
Additionally, the hidden dimension $k$ is varied among ${4,8,16}$ while fixing
the model scale to small to identify the effect of problem scale.

Table 5:  Model configurations for Chain-of-Thought-I/O experiments.  Model | # layers | embed dim | # heads (MHA)
---|---|---|---
Standard | 12 | 256 | 8
Small | 6 | 128 | 4
Tiny | 3 | 64 | 2

###  B.2 Vector-valued MQAR

The training set consists of $300,000$ training samples. We train for $64$
epochs with batch size of $64$ and evaluate on a test set of $3,000$ samples.
For each setting, we sweep with learning rates in np.logspace(-4, -2, 4) and
report the best result among all learning rates.

###  B.3 Orthogonal-outlier Regression

We run Orthogonal-outlier regression on all four model architectures, with
varying number of parameters. The curves above generally capture the trend
that Mamba and Transformer perform on par with each other, and Standard Hybrid
and MambaFormer show better curves compared to the vanilla models.

One curve that stands out would be Mamba’s loss curve in the smaller regime.
As a future direction, we have tested trading off the depth and width while
fixing the total number of parameters with the smallest Mamba model
configuration in Table 6.

Table 6:  Various Mamba model configurations with fixed number of parameters.  Model | # layers | embed dim
---|---|---
Standard | 4 | 128
| 2 | 192
| 8 | 96
| 16 | 64

The best performing model configuration was with 16 layers, and embedding
dimension of 64. This suggests that with fixed number of parameters,
increasing the depth may boost the downstream task performance on ICL tasks.

##  Appendix C FLOPs computation

We count the number of multiplications in a Mamba block and a Transformer
block in Table 7 and Table 8. We assume batch size $B=1$. To calculate FLOPs,
we multiply the number of multiplications by $6$ to account for the multiply-
accumulate cost in both forward and backward pass. Note that a Standard Hybrid
block is an attention block stacked with a Mamba block, so the number of
multiplications in a Standard Hybrid block is $10LD^{2}+2L^{2}D$, ignoring the
linear terms.

Table 7:  Number of multiplications in a Transformer block.  | Number of multiplications
---|---
QKV projection | $3LD^{2}$
Outer product and multiply $V$ | $2L^{2}D$
Outer projection | $LD^{2}$
FFN with ffw_width=4 | $8LD^{2}$
Table 8:  Number of multiplications in a Mamba block. Unless specified otherwise, we assume $E=2$ and $R=2$.  | Number of multiplications
---|---
Input projection | $2LED^{2}$
SSM | $7LEDN+2RLED$
Output projection | $LED^{2}$