Is Mamba Capable of In-Context Learning?

Authors: riccardo  grazzi
julien  siems
simon  schrodi
thomas  brox
frank  hutter
Venue: automl_cc_AutoML_2024_Conference
Published: 2024-02-26 16:01:46

Abstract: The surprising generalization capabilities of foundation models have been enabled by in-context learning (ICL), a new variant of meta-learning that denotes the learned ability to solve tasks during a neural network  forward pass, exploiting contextual information provided as input to the model. This useful ability emerges as a side product of the foundation model's massive pretraining. While transformer models are currently the state of the art in ICL, this work provides empirical evidence that Mamba, a newly proposed state space model which scales better  than transformers w.r.t.\@ the input sequence length, has similar ICL capabilities. We evaluated Mamba on tasks involving simple function approximation as well as more complex natural language processing problems. Our results demonstrate that, across both categories of tasks, Mamba closely matches the performance of transformer models for ICL. Further analysis reveals that, like transformers, Mamba appears to solve ICL problems by incrementally optimizing its internal representations. Overall, our work suggests that Mamba can be an efficient alternative to transformers for ICL tasks involving long input sequences. This is an exciting finding in meta-learning and may also enable generalizations of in-context learned AutoML algorithms (like TabPFN or Optformer) to long input sequences. The anonymous code to reproduce our experiments is available [here](https://anon-github.automl.cc/r/is_mamba_capable_of_in_context_learning-7C49/README.md).

Main paper: 
 introduction: 
 Recent advancements in large-scale neural language modeling (Brown et al., 2020) have demonstrated that Transformer models (Vaswani et al., 2017) exhibit in-context learning (ICL) capabilities: after (self-supervised) pre-training, they can infer how to perform tasks only from input examples without the need for explicit training nor fine-tuning. This ability represents a departure from established in-weights learning of traditional machine learning and has sparked considerable academic interest as a new type of meta-learning. In contrast to standard meta-learning approaches (Hospedales et al., 2021) , ICL emerges in transformer models from pre-training: without explicit training on a distribution of tasks, without bi-level optimization, and without any specific inductive bias. Recent studies advanced the understanding of how transformers can implement and learn in-context gradient-based methods when trained on distributions of simple supervised learning tasks, e.g., on linear regression tasks (Von Oswald et al., 2023; Ahn et al., 2023; Bai et al., 2023) . Despite such results, whether pre-trained transformers perform in-context gradient methods on more complex tasks remains an ongoing discussion (Shen et al., 2023) .Orthogonal to these investigations into the transformer architecture, recent work proposed deep state space models to overcome limitations of transformers in processing long sequences (Tay et al., 2021) , such as S4 (Gu et al., 2021a) or H3 (Gu et al., 2021b) . These models merge elements from recurrent and convolutional networks with state space approaches (Kalman, 1960) . However, their success on NLP tasks was limited due to problems handling dense information tasks. A key feature of state space models is that they can run the forward pass in two modes with different complexity w.r.t. the input sequence length: a parallel mode, ideal for training but with superlinear time complexity and a recurrent mode with linear time complexity and more suited for inference.In contrast, the forward pass of transformer models has quadratic time complexity and hence it is less efficient.This work conducts an investigation into the ICL capabilities of the recently proposed Mamba architecture (Gu and Dao, 2023) , a successor to S4 and H3. Mamba has already shown its potential in different applications beyond NLP, such as visual representation learning (Zhu et al., 2024) or image segmentation (Ma et al., 2024) . Concurrent to our work, the ICL capabilities of Mamba were investigated on synthetic language learning tasks by Akyürek et al. (2024) and on other tasks, including simple function classes, by Park et al. (2024) . Differently from those works, we study the performance of the pre-trained Mamba model on NLP tasks, conduct a probing analysis on the representations at intermediate layers to better understand Mamba's ICL solution mechanisms, and measuring the models' ability to extrapolate beyond the context length used for training on simple function classes.Models capable of ICL have shown good AutoML performance. A remarkable example are Prior-data Fitted Networks (PFNs) (Müller et al., 2022) , which are (usually transformer) models pretrained on a range of synthetic tasks generated from a prior distribution, to approximate the posterior predictive distribution. Interestingly, by harnessing ICL and with a rich enough prior, PFNs, specifically TabPFN (Hollmann et al., 2023) , exhibit strong generalization performance on real tabular data, enabling predictions on new data without the need for explicit specification and optimization of a table-specific model. Other applications of ICL relevant to AutoML are in Black-Box Optimization (Müller et al., 2023; Yang et al., 2023a; Chen et al., 2022) , Learning-Curve Extrapolation (Adriaensen et al., 2024) , and Time-Series Forecasting (Dooley et al., 2024) . Our experimental setup for simple function classes, which closely follows the one in (Garg et al., 2022) , is similar to that of PFNs, but the priors we use concern only simple regression tasks and the model only outputs the mean of the posterior predictive distribution.Contributions. After introducing state-space models and Mamba more formally (Section 2), we make the following contributions:• We demonstrate that Mamba is capable of ICL and performs on-par with transformers on simple function classes (Section 3.1) and more complex NLP tasks (Section 4). This highlights Mamba as an efficient alternative to transformers for ICL tasks entailing long sequences. In addition, we find that Mamba outperforms its predecessor S4, and RWKV (Peng et al., 2023) , a recent parallel/recurrent architecture.• Using a simple probing approach, we provide preliminary insights into the mechanism by which Mamba incrementally solves ICL tasks (Section 3.2). We find that the optimization processes exhibited by Mamba are similar to those of transformer models. 
state space models and the mamba architecture: 
 State Space Models (SSMs) are sequence-to-sequence models (when discretized) with learnable parameters which are inspired by a continuous-time model. SSMs map an input sequence (𝑥 1 , . . . , 𝑥 𝐿 ) to the output sequence (𝑦 1 , . . . , 𝑦 𝐿 ) by computing latent states (ℎ 1 , . . . , ℎ 𝐿 ). The map of linear time invariant SSMs such as S4 (Gu et al., 2021a) or H3 (Fu et al., 2023) can be equivalently written asℎ 𝑡 = Āℎ 𝑡 -1 + B𝑥 𝑡 𝑦 𝑡 = 𝐶ℎ 𝑡(1)EQUATIONwhere ℎ 0 = 0, Ā = 𝑓 1 (𝐴, 𝐵, Δ), B = 𝑓 2 (𝐴, 𝐵, Δ), 𝑓 1 , 𝑓 2 specify the type of discretization and (𝐴, 𝐵, 𝐶, Δ) are learnable parameters. In particular, if 𝑥 𝑖 ∈ R 𝐷 and ℎ 𝑖 ∈ R 𝑁 𝐷 , then 𝐴, 𝐵, 𝐶 are matrices of dimensions 𝑁 𝐷 × 𝑁 𝐷, 𝑁 𝐷 × 𝐷 and 𝐷 × 𝑁 𝐷 respectively, while Δ is a scalar discretization parameter.To have a sufficiently large memory of the past, the dimension of ℎ 𝑡 is usually much larger than that of 𝑥 𝑡 . The outputs and states in equation 1 can be computed recurrently with 𝑂 (𝐿) time complexity and 𝑂 (1) space complexity. However, since the state update is linear in ℎ 𝑡 -1 and 𝑥 𝑡 and does not change with 𝑡 (linear time invariant) the outputs can also be obtained (see equation 2) by first computing 𝐾 𝑖 in parallel for 𝑖 = 0, . . . , 𝐿 -1 and then computing (𝑦 1 , . . . , 𝑦 𝐿 ) through a convolution implemented via Fast Fourier Transform (FFT), which can be easily parallelized and has time complexity of 𝑂 (𝐿 log(𝐿)). The recurrent mode is ideal for auto-regressive inference, since in that case parallelization is not possible, while for training the convolutional mode can fully take advantage of the parallelism of modern specialized hardware. A fundamental limit of linear time invariant SSMs is that they do not have a mechanism to select which information to retain in the latent state based on the input sequence, and this makes them perform badly in tasks involving content-aware reasoning like selective copying, in which transformer models excel. The Mamba architecture (Gu and Dao, 2023 ) is a linear time varying SSM which overcomes this limit via a selection mechanism allowing the latent state dynamics to change with the current input, thus enabling selective retention of information. In particular, state and outputs of Mamba follow the recursion (starting from ℎ 0 = 0)ℎ 𝑡 = Ā𝑡 ℎ 𝑡 -1 + B𝑡 𝑥 𝑡 𝑦 𝑡 = 𝐶 𝑡 ℎ 𝑡 where Ā𝑡 = exp(𝐴Δ 𝑡 ), B𝑡 = exp(𝐴Δ 𝑡 ) -1 exp(Δ 𝑡 𝐴 -𝐼 )Δ 𝑡 𝐵 𝑡 (zero order hold discretization), 𝐵 𝑡 = 𝑊 1 𝑥 𝑡 + 𝑏 1 , 𝐶 𝑡 = 𝑊 2 𝑥 𝑡 + 𝑏 2 , Δ 𝑡 = softplus(𝑊 3 𝑥 𝑡 + 𝑏 3) and (𝐴,𝑊 1 ,𝑊 2 ,𝑊 3 , 𝑏 1 , 𝑏 2 , 𝑏 3 ) are learnable parameters and 𝐴 is diagonal. The now time-varying discretization parameter Δ 𝑡 ∈ R enables a gating mechanism which can selectively ignore the current input or reset the state. However, the selection mechanism hinders the use of convolutions for fast and parallel training. Despite this, Mamba has a similar training time as linear time invariant SSMs thanks to a hardware-aware algorithm that uses a parallel scan (𝑂 (𝐿) time complexity with only 𝑂 (log(𝐿)) sequential steps) in place of FFT and stores and updates the large latent states only in the fast SRAM, rather than in GPU HBM memory. We refer to Gu and Dao (2023) for further details on Mamba's architecture. 
investigation of simple function classes: 
 In this section, we assess Mamba's ability to learn task distributions involving simple function classes. We followed the experimental protocol of Garg et al. (2022) : each model is trained on a task distribution and then tested on the same distribution. This process was repeated for 4 regression task distributions, each falling into a specific function class: linear functions, sparse linear functions, 2-layer ReLU neural networks, and Decision trees. Models trained on linear functions were also tested on out-of-distribution (OOD) tasks. Differently from Garg et al. (2022) , we also tested the models on tasks with more input examples than seen during training, to measure whether they can extrapolate to longer inputs. Details on each task distribution are provided in Appendix A.1.3.We compare Mamba to a causal transformer model using the GPT2 architecture (Radford et al., 2019) and to some baselines specific to each function class (see Appendix A.1.4). We also compare with S4, a linear time invariant model, to measure the impact of the selection mechanism of Mamba. To ensure a fair comparison, the architectures of Mamba and S4 are adjusted to have a comparable number of parameters to the transformer (9.5 M). For further training details, we refer to Appendix A.1.1.We removed the positional encoding used in the causal transformer by Garg et al. (2022) , since we observed that it hinders the transformer's ability to extrapolate to longer inputs, as also shown by Press et al. (2021) . Similar to Müller et al. (2022) , we argue this to be more natural in the present setup as the input is actually a set of data points, not a sequence.To sample the inputs of each training task for linear regression, Garg et al. (2022) used a normal distribution, while we used a skewed Gaussian distribution. Our findings show that this improves robustness w.r.t. OOD tasks of both Mamba and transformers. Formally, we trained on linear functions in the set 𝐹 = {𝑓 : 𝑓 (𝑥) = 𝑤 𝑇 𝑥, 𝑤 ∈ R 𝑑 } with input dimension 𝑑 = 20. For each training task, we sampled 𝑤 from an isotropic Gaussian N (0, 𝐼 𝑑 ) and 𝑥 1 , . . . , 𝑥 𝑘 , 𝑥 𝑘+1 from N (0, Σ), where Σ is a skewed covariance matrix with its eigenbasis chosen at random and the 𝑖th eigenvalue proportional to 1/𝑖 2 . Following this, we set 𝑦 𝑖 = 𝑤 𝑇 𝑥 𝑖 and assembled the input prompt 𝑃 = (𝑥 1 , 𝑦 1 , 𝑥 2 , 𝑦 2 , ..., 𝑥 𝑘 , 𝑦 𝑘 , 𝑥 𝑘+1 ). We used the mean squared error (MSE) 1 𝑘+1 𝑘+1 𝑖=1 ( ŷ𝑖 -𝑦 𝑖 ) 2 as loss function, where ŷ𝑖 is the output of the model corresponding to 𝑥 𝑖 . For all results in this section, we report the MSE/𝑑 (𝑑 = 20), which is what is reported by Garg et al. (2022) for linear regression.Results are shown in Figures 1 and 2 . In skewed linear regression, both Mamba and the Transformer model closely match the least squares baseline both in-distribution and out-of-distribution as long as the context length is shorter than the one used for training. Increasing the context length generally decreases performance, by an amount depending on the training run: two runs of the transfomer model present very little degradation while Mamba degrades substantially for all three runs. In contrast, S4 performs much worse than the least squares baseline in all setups. We hypothesize that S4's poor ICL performance is due to its linear time invariance. A similar hypothesis was also drawn by Gu and Dao (2023) for the task of selective copying. Similar results for Mamba and the Transformer also hold for sparse linear regression, while for neural networks and Decision trees, Mamba and the transformers are comparable and even show a promising input length extrapolation: the performance improves after the dotted vertical line indicating the number of in-context examples used during training. Interestingly, in the case of Decision trees, two out of three runs of the Transformer perform substantially better than Mamba, while Mamba's error degrades less than the one of the Transformer for ReLU neural networks. We also note that when trained on non-skewed linear regression tasks, Mamba and Transformers are less robust to OOD tasks (Figure 9 ), even if they both perform well in-distribution (Figure 8 ). In the appendix we also report the performance when increasing the number of in-context examples during training for skewed linear regression (Figure 10 ).To better understand how the Mamba and transformer models perform ICL, we aim to test if they both employ a solution strategy akin to iterative optimization, i.e., we study if they incrementally improve their solutions layer after layer (Von Oswald et al., 2023; Ahn et al., 2023; Bai et al., 2023) . We adopted a probing strategy similar to the one by Geva et al. (2021) for transformer language models. Differently from other probing strategies like the one by Akyürek et al. (2023) , who learn a (non-linear) probe on the high-dimensional intermediate representations, this strategy learns a linear probe on the output of the decoder applied to the intermediate representations after each layer. We argue this to be a less biased probing strategy, since in our setup it reduces the degrees of freedom of the probe to just 2 parameters per task (scale and shift), since the models have scalar output.More specifically, our probing strategy works as follows. Let {(𝑥 𝑖 , 𝑦 𝑖 )} 𝑖 be i.i. examples for the train, validation, and test set, respectively. First, we separately feed the prompts ) . . . ( ỹ𝑙,𝑘+𝑚 , 𝑦 𝑘+𝑚 ) (i.e., using the validation tokens). We can then measure the accuracy of these intermediate predictions with the normalized mean squared error on the test tokens (𝑛𝑑) -1 𝑛 𝑗=1 ( ŷ𝑙,𝑘+𝑚+𝑗 -𝑦 𝑘+𝑚+𝑗 ) 2 (𝑑 = 20). We ran our analysis over 128 tasks (e.g. different linear regression weight vectors) with 𝑘 as specified in Figures 3 and 4 , 𝑚=10 and 𝑛=40. We conducted the probing analysis on Mamba and transformer models trained on skewed linear regression, sparse linear regression, ReLU neural networks, and Decision trees. For skewed linear regression, we additionally compared to Gradient Descent (GD) and GD ++ as done in Von Oswald et al. (2023) . GD ++ is a version of preconditioned gradient descent in which the data samples undergo a transformation (𝑥 𝑖 ← (𝐼 -𝛾𝑋𝑋 𝑇 )𝑥 𝑖 ) and we tuned 𝛾 for optimal performance via a grid search (c.f., Figure 11 in Appendix), while the step-size was set for each task to the theoretically optimal value, i.e., as 2/(𝐿 + 𝜇), where 𝐿, 𝜇 are the largest and smallest eigenvalues of the empirical covariance matrix of the task. We ran GD and GD ++ for 24 iterations to match the 24 layers of our Mamba model. Figure 3 provides strong evidence to the hypothesis that both Mamba and the transformer employ an iterative solution scheme on the skewed and sparse linear regression task (Figure 3d ), since the log-MSE decreases (almost) linearly. In ReLU neural networks, the error also decreases somewhat gradually after some layers. Interestingly, for Decision trees the error stays high for both models for more than half of the initial layers. The comparison between GD, GD ++ , Mamba, and the transformer in Figure 3d reveals that both GD and GD ++ are outperformed by Mamba and the Transformer, while GD converges more slowly due to the tasks having skewed covariance.(𝑥 1 , 𝑦 1 , 𝑥 2 ,In Figure 4 we show how, similarly to Figure 1a , having more in-context examples than the ones used during training negatively affects the learning curves for linear regression. Both models never seem to take advantage of the additional number of examples in this task: the performance with 40 in context examples (the number used for training) is the best (or close) for all layers. Moreover, we note that the learning curve for Mamba with 160 in-context examples exhibits a clear U-shape: the model actually found a good solution in intermediate layers but ultimately degrades.In Figure 5 , we compare the scale and shift parameters estimated by the linear probing model for Mamba and the transformer in the context of Skewed Linear Regression. Additionally, for sparse linear regression, ReLU neural networks, and Decision trees, we present the findings in Figure 12 in the Appendix. Interestingly, for linear regression and ReLU neural networks, the linear probing model applies only minor modifications, with scale values near one and shift values approaching zero after the first layers. We also observe that for all tasks, the variance across tasks of the estimated scale and shift is higher in the early layers of the model and quickly approaches zero as we move towards the output layer. In this section, we evaluated the ICL performance of various pre-trained Mamba language models, with parameter counts from 130 million to 2.8 billion. The pre-training for all variants was done on the Pile dataset (Gao et al., 2020) , while the Mamba 2.8B (SP) checkpoint underwent further fine-tuning on ca. 600 billion tokens from the SlimPajama dataset.We compared the Mamba variants to another RNN model with linear state dynamics named RWKV (Peng et al., 2023) , which was also pre-trained on the Pile1 , and popular transformer-based language models, such as LLama (Touvron et al., 2023) , Pythia (Biderman et al., 2023) , and GPT-J 6B (Wang and Komatsuzaki, 2021) . We did not compare to S4 because we are not aware of S4 models pretrained on the Pile.We followed the experimental protocol of Hendel et al. (2023) , which tested 27 NLP tasks spanning a wide range of categories, including algorithmic tasks (e.g., list element extraction), translation (e.g., English to Spanish), linguistic tasks (e.g., singular to plural conversion), and knowledge-based tasks (e.g., identifying country-capital pairs). For evaluation, we used the same datasets as Hendel et al. (2023) except for the algorithmic tasks, which were randomly generated (we use the same generation parameters). We report the mean accuracy over 400 generated test sets per task, each having five in-context examples. Figure 6 shows that the ICL performance improves for all models with increasing number of parameters. Notably, Mamba 2.8B achieves an ICL accuracy close to LLama 7B, and on par with GPT-J and Pythia models. In addition, we find that Mamba consistently outperforms the similarly scalable RWKV at comparable parameter sizes. We provide a detailed table of per-task accuracies in Appendix A.2.Finally, we find that Mamba scales well with the number of in-context examples; see Figure 7 . Particularly, Mamba 0.37B and 2.8B maintain a considerable performance edge over RWKV 0.43B and 3B, respectively. 
related work: 
 In-context learning capabilities in transformer-based models. Brown et al. (2020) introduced GPT-3 and were the first to show the remarkable in-context learning (ICL) abilities of large language models. They showed that larger models tend to utilize more in-context information compared to smaller models. Other work by Garg et al. (2022) showed that even small transformers exhibit ICL abilities on linear modeling tasks. These small transformer-based models have been subject to in-depth analysis by recent work (Dai et al., 2022; Von Oswald et al., 2023; Ahn et al., 2023; Bai et al., 2023) that showed that they can implement mechanisms akin to gradient descent or more complex optimization algorithms in the forward pass.On a different tangent, Olsson et al. (2022) found that a two-layer attention-only network can develop the so-called "induction heads" mechanism, which outputs the token succeeding a previous instance of the current token, precisely when its ICL performance increases. Chan et al. (2022) investigated properties of the data-distribution which lead to the emergence of ICL abilities, while Reddy (2024) identified factors for the abrupt emergence of the induction heads. Another line of work (Hendel et al., 2023; Todd et al., 2024) showed that some intermediate output of the forward pass of transformer models, named "task vectors", encodes most of the information for the in-context task. In particular, merging a task vector with another from a different in-context task allows the model to solve a combination of the two tasks.In-context learning benchmarks. Garg et al. (2022) tested transformers abilities to learn task distributions from simple function classes, like linear regression or Decision trees. Hendel et al. (2023) tested pre-trained LLMs on a suite of NLP tasks such as English to French translation. Lastly, Akyürek et al. (2024) proposed to test NLP models' ability to learn formal languages.Sub-quadratic transformer alternatives. Transformer-based models have a forward pass with time complexity scaling quadratically with respect to the input sequence length. To reduce the amount of computation, previous work proposed to introduce sparsity into the attention layers (Child et al., 2019; Qiu et al., 2019; Beltagy et al., 2020) , where the tokens only attend a subset of the other tokens. Other works by Wang et al. (2020) proposed a low-rank factorization of the attention mechanism. Others proposed linear attention, where the attention weights are computed using possibly normalized dot-products without the softmax (Katharopoulos et al., 2020; Choromanski et al., 2020; Kasai et al., 2021; Peng et al., 2021; Yang et al., 2023b) . Intriguingly, causal linear attention can be re-formulated either as a linear RNN or as the ratio between two linear RNNs, and hence its time complexity scales linearly with the length of the input sequence.Recently, state-space sequence models gained significant interest due their strong performance. They are broadly related to recurrent neural networks with ideas from classical state space models, such as the Kalman filter (Kalman, 1960) , particle filters (Gordon et al., 1993) , or hidden Markov models (Baum et al., 1970; Rabiner, 1989) . There exist various state-space models: S4 (Gu et al., 2021a) , DSS (Gupta et al., 2022) , S5 (Smith et al., 2023) , GSS (Mehta et al., 2023) , H3 (Fu et al., 2023) , Selective S4 (Wang et al., 2023) , or RWKV (Peng et al., 2023) . While they have been successfully applied to audio and vision (Goel et al., 2022; Nguyen et al., 2022) , they lacked performance for text. The recently proposed Mamba (Gu and Dao, 2023) exhibits great performance in large language modeling and similar scaling properties as state of the art transformer models. 
broader impact statement: 
 In-context learning (ICL) is an emergent phenomenon that drives the generalization skills of large language models (LLMs). While the widespread use of a poorly-understood technology like ICL may have many potential negative societal impacts, our work contributes to mitigate this problem by helping to build a mechanistic understanding of ICL. This understanding is crucial to build better, more explainable and more trustworthy AI systems in the future. Since the favourable scaling of Mamba (Gu and Dao, 2023) makes it a strong candidate for dealing with long context sizes, we do believe it to be particularly important to understand the mechanisms behind ICL in this architecture, and we thus focus on this aspect in our paper. 
limitations: 
 Our investigation into the Mamba architecture's in-context learning (ICL) capabilities shows promise as an efficient alternative to transformer models for processing long input sequences, yet it is not without limitations. The study's focus on simple function approximation and natural language processing (NLP) tasks raises questions about the generalizability of our findings to other domains like image or audio analysis, suggesting the need for broader research. The linear probing method used to understand Mamba's approach may oversimplify the model's complex optimization processes, indicating that more sophisticated probing techniques could yield deeper insights. Additionally, our evaluation covered a limited range of model sizes and configurations, and a more comprehensive comparison is necessary to fully assess Mamba's performance relative to extensively optimized transformer models. The scalability of Mamba with increasing in-context examples and its computational efficiency also remain underexplored, highlighting areas for future investigation to better understand its practical applicability and potential across a wider range of tasks and settings. 
conclusions and future directions: 
 In this work, we have demonstrated that the recently proposed Mamba architecture is capable of effective in-context-learning (ICL) across tasks involving simple function approximation as well as more complex natural language processing problems. Our analysis showed that Mamba performs on par with transformer models, while also outperforming the S4 and RWKV baselines. We provide initial evidence that Mamba appears to solve ICL problems by incrementally refining its internal representations in a manner akin to an iterative optimization strategy, as transformers do. Overall, our findings suggest that Mamba can be an efficient and performant alternative to transformers for ICL involving longer input sequences.In future work, it would be interesting to replace the attention-based backbone of existing incontext learned AutoML algorithms, such as TabPFN (Hollmann et al., 2023) , ForecastPFN (Dooley et al., 2024) or Optformer (Chen et al., 2022) , enabling them to handle longer sequences. those of the European Union or the European Research Council. Neither the European Union nor the granting authority can be held responsible for them. 
 
 Appendix: 1. For all authors. . . (h) Did you include the total amount of compute and the type of resources used (e.g., type of gpus, internal cluster, or cloud provider)? Yes.(i) Did you run ablation studies to assess the impact of different components of your approach? No.3. With respect to the code used to obtain your results. . .(a) Did you include the code, data, and instructions needed to reproduce the main experimental results, including all requirements (e.g., requirements.txt with explicit versions), random seeds, an instructive README with installation, and execution commands (either in the supplemental material or as a url)? Yes.(b) Did you include a minimal example to replicate results on a small subset of the experiments or on toy data? Yes.(c) Did you ensure sufficient code quality and documentation so that someone else can execute and understand your code? Yes.(d) Did you include the raw results of running your experiments with the given code, data, and instructions? Not for the section 3, but for the simple NLP tasks.(e) Did you include the code, additional data, and instructions needed to generate the figures and tables in your paper based on the raw results? Yes.4. If you used existing assets (e.g., code, data, models). . .(a) Did you cite the creators of used assets? Yes.(b) Did you discuss whether and how consent was obtained from people whose data you're using/curating if the license requires it? No consent required.(c) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content? No personally identifiable information was collected.5. If you created/released new assets (e.g., code, data, models). . . In Figure 10 , we observe how the context length extrapolation performance is affected by using a higher number of examples during training for skewed linear regression tasks. In particular, the window where both models obtain good performance is wider the higher the number of incontext examples during training, with the transformer generally exhibiting a much lower degree of degradation.A.1.6 ICL learning curves. The scale and shift were estimated for each layer and each task and are reported in Figure 12 . We note that for skewed linear regression and ReLU neural networks the value for scale approaches 1, and the one for shift approaches 0 relatively quickly.We added Mamba and RKWV to the code provided by (Hendel et al., 2023) at github.com/roeehendel/icl_task_vectors. We leveraged their results for GPT-J, Llama and Pythia and provided our experimental results for Mamba and RWKV. A detailed break down of the per task performance of each model is given in Table 1 . We used RWKV v4 checkpoints trained on the Pile dataset by the authors available on huggingface. 


Decision: This paper investigates the in-context learning capabilities of Mamba. Experiments demonstrate that Mamba is able to successfully perform in-context learning and shed light on the optimization process by which it does so. Reviewers agreed that the paper is well-written and the experimental methodology is sound. Some doubt was expressed about the extent to which the results are relevant to AutoML. However, the authors clarified the connections to meta-learning and potential applications of the work towards improving AutoML algorithms such as Optformer and TabPFN. I therefore recommend acceptance.

Reviews: review: paper_summary: paper_summary: The paper "Is Mamba Capable of In-Context Learning?" explores the in-context learning (ICL) capabilities of the Mamba architecture, a state space model proposed as an alternative to transformer models for processing long input sequences. The study investigates Mamba's performance across tasks involving simple function approximation and complex natural language processing (NLP) challenges, comparing its ICL abilities to those of transformer models, its predecessor S4, and a recurrent neural network architecture, RWKV. The results demonstrate that Mamba matches or even outperforms these models in ICL tasks, particularly in handling long sequences efficiently. The analysis suggests that Mamba, like transformers, employs an iterative optimization process to refine its internal representations for task solving.
main_review: main_review: It would be valuable to look into recent advancements and discussions around in-context learning capabilities in large language models, state space models versus transformers for long sequence handling, and novel methodologies for analyzing model behaviors. This additional information could offer alternative perspectives on Mamba's strengths and limitations, as well as potential directions for future work in enhancing its performance and understanding its mechanisms of action.
review_summary: review_summary: The authors investigate Mamba, a popular alternative architecture to Transformers with a focus on In-Context Learning.The experiments on simple tasks raises questions about comprehensiveness of the study. 
I am a little unsure if this paper has direct relevance to AutoML
paper_summary: paper_summary: The paper explores Mamba's in-context learning (ICL) capabilities. The authors demonstrate that Mamba is capable of ICL and outperforms other state space models (SSMs) while performing similarly to Transformers. The evaluation is conducted using simple function classes and NLP tasks.
main_review: main_review: Strengths:  The paper provides empirical evidence and demonstrates Mamba's ICL capabilities with simple tasks. The methodology and analysis are sound.  
Weaknesses: The experiments on simple tasks might raise concerns about the claims' strength.  The related work section has room for improvement. ICL capabilities in transformer-based models can be discussed in the Appendix.
review_summary: review_summary: The paper reasonably investigates Mamba, a popular alternative architecture to Transformers. The focus is on its in-context learning (ICL) capabilities, and the authors do a good job of exploring Mamba's ICL strengths and limitations. Although I have some concerns regarding the complexity of the tasks used for their experiments, the paper presents a good initial investigation of Mamba's ICL capabilities, and I recommend a weak acceptance.
paper_summary: paper_summary: &nbsp;

The authors investigate the question as to whether Mamba is a competitive approach to transformers for in-context learning (ICL). The authors evaluate Mamba's ICL capabilities on function approximation and NLP tasks. Additionally, the authors provide diagnostic experiments to examine the mechanism of ICL on a layer-by-layer basis in both Mamba and transformers.

&nbsp;
main_review: main_review: &nbsp;

While shedding light on the mechanism of ICL in Mamba is interesting, the link to AutoML is tenuous and it would be nice to see explicit experiments where the findings are directly used in an AutoML context. 

&nbsp;
review_summary: review_summary: &nbsp;

Given the lack of clarity regarding the experimental setup, I would encourage the authors to revise the exposition of the paper and potentially add experiments in an AutoML setting to make the work relevant for the conference. I would be open to hearing the authors' opinions on where the current work might be a good fit (without a connection to AutoML) with reference to the call for papers.

Update: Now able to edit score, have upgraded!

&nbsp;
paper_summary: paper_summary: This paper investigates the In-Context Learning (ICL) capabilities of the Mamba architecture, a state space model that scales better than transformers with respect to input sequence length.
main_review: main_review: Paper Strength
1. This paper provides empirical evidence of Mamba's comparable performance to transformer models in ICL tasks.
2. It offers a novel analysis of Mamba's mechanism for solving ICL problems through probing studies.
3. The research opens up possibilities for generalizing in-context learned AutoML algorithms to long input sequences.

Paper Weakness
Unclear Impact on AutoML: The paper suggests that Mamba could be used for generalizing in-context learned AutoML algorithms to long input sequences. However, the relevance of the paper's focus to the AutoML community is not clear to me.
review_summary: review_summary: This paper presents interesting findings on Mamba's capabilities for ICL tasks. The authors might consider addressing the limitations above.