Can Custom Models Learn In-Context? An Exploration of Hybrid Architecture Performance on In-Context Learning Tasks

Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 22:00:11

Abstract: In-Context Learning (ICL) is a phenomenon where task learning occurs through a prompt sequence without the necessity of parameter updates. ICL in Multi-Headed Attention (MHA) with absolute positional embedding has been the focus of more study than other sequence model varieties. We examine implications of architectural differences between GPT-2 and LLaMa as well as Llama and Mamba. We extend work done by Garg et al. (2022) and Park et al. (2024) to GPT-2/LLaMa hybrid and LLaMa/Mamba hybrid models -- examining the interplay between sequence transformation blocks and regressive performance in-context. We note that certain architectural changes cause degraded training efficiency/ICL accuracy by converging to suboptimal predictors or converging slower. We also find certain hybrids showing optimistic performance improvements, informing potential future ICL-focused architecture modifications. Additionally, we propose the "ICL regression score", a scalar metric describing a model's whole performance on a specific task. Compute limitations impose restrictions on our architecture-space, training duration, number of training runs, function class complexity, and benchmark complexity. To foster reproducible and extensible research, we provide a typed, modular, and extensible Python package on which we run all experiments. This code is available at \url{https://github.com/anonymousforneurips64/neurips2024-submission21757}.

Main paper: 
 introduction: 
 Popularized by Large Language Models such as GPT-2 [1] and GPT-3 [2] , In-Context Learning (ICL) is the ability for highly expressive generative sequence models to predict phenomena by processing demonstrations without performing traditional gradient steps. Such phenomena vary from effective control systems [3] to answering questions in natural language [4, 5] . A large body of recent work has studied this phenomenon in transformer models [6, 7, 2, 1, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25] , which derive in structure from Vaswani et al. [26] .Some recent examples of this research on ICL include Garg et al [6] , which studies ICL by providing a variety of function classes for models to learn, additionally benchmarking robustness by testing performance on out-of-distribution data. Guo et al [11] shows the validity of composing simple function classes to produce complex ones, while Liu et al [20] produced a metric for model information recall.These works give us a set of metrics with which we can use to compare model performance on ICL.ICL was initially primarily studied in attention-based models but has recently been explored in other sequence models, creating discussion on its differences across those models and why these Submitted to 38th Conference on Neural Information Processing Systems (NeurIPS 2024). Do not distribute. ) -Table 1 : Summary of tasks. Each regression target f θ (x i ) is either parametrized by a randomly sampled θ or directly computed/sampled as detailed above.occur architecturally. In our paper, we study this by substituting key modern transformer (Llama) components with Mamba blocks and GPT-2 components and richly benchmarking.Since ICL for complete natural language understanding often requires training models with over a billion parameters, the effects of architectural changes on fine-grained ICL abilities are often left unexplored. As a consequence, although language models have progressed quickly and entertained radically new architectures, there is limited extensible research that explores the effects of fine-grained architecture choices on ICL ability [8, 14] . Garg et al. established using simple function classes to evaluate ICL ability and examined solely GPT-2 as a sequence model. Lee et al. [8] expanded this analysis on a slightly different set of function classes for a variety of base models. Park et al. [14] evaluated ICL performance of 2 hybrid architectures between Mamba and GPT-2. Using unmodified Llama/Mamba/GPT-2 as a control, we analyze GPT2-Llama and Llama-Mamba hybrid architectures derived from replacing portions of GPT2 components with analogous Llama sections and LLama with Mamba blocks, respectively, in 12 total architectures (3 unmodified + 9 hybrid).We observe that the code written to analyze ICL with simple function classes -although almost unanimously extensions of Garg et al.'s -often requires substantial, structural changes to the parent codebase 1 , greatly heightening the barrier to extending each project in turn. Inspired by Donoho's ideal of Frictionless Reproducibility [27] , we provide a set of simple abstractions and interfaces to facilitate extensions and modifications to our code while promoting interoperability between forks. 
related work: 
 There are many ways to capture qualitative aspects of ICL with quantitative measures. Weber et al. [17] compare the agreement between generations of a language model under varying prompts of equal meaning to test robustness to variations. Olsson et al. [22] compute a heuristic "ICL score" to measure an accuracy increase in predictions of a model given more context. We adapt this metric to fit our experimental setup more aptly, regularizing along both the number of in-context examples and against a baseline predictor.In general, evaluating ICL ability has been approached from two primary avenues: both when the only solution at train time is to meta-learn an algorithm [6, 8, 28, 11, 19] and when optimal loss at train time can also be satisfied by memorization or otherwise leveraging previously trained-on data [10, 23] . In this work, we take the former approach through learning a regression algorithm to randomized simple function classes [6, 11, 15] .Further still, non-transformer architectures are capable of ICL [8] . Lee et al. [8] observed ICL in numerous sequence model architectures (e.g. RNNs, Mamba, S4, CNNs, GPT-2, and Llama) and found qualitative differences in each architecture's performance. Chan et al. [25] found that Transformers depend on "burstiness" and long-tail distributions of natural data to outperform RNNs and LSTMs in ICL tasks. Park et al. [14] uses simple function classes similar to Garg et al. [6] in evaluating the ICL ability of Mamba, S4, S4-Mamba, and GPT-2. They find an overlapping but inequivalent set of function classes for which each model succeeds and construct a hybrid architecture to achieve the union of these abilities. We further this work by closely examining the contributions of individual architectural changes for GPT-2 and Llama-style transformers towards ICL ability.Question: Does the paper discuss the limitations of the work performed by the authors?Answer: [Yes]Justification: We briefly mention some limitations in our analysis and experiments in Section 5.1. We acknowledge our limited training runs, inexhaustive training horizon, incomplete evaluation of ICL Regression Score, and no metrics on the usability of our codebase.Similarly here we acknowledge that this list of limitations is by no means exhaustive.Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.• The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.• The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.• The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.• The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.• If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.• While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
methods: 
 As established by Garg et al. and extended by recent work, our ICL tasks take the following form [6, 8, 14] :x 0 , f θ (x 0 ), x 1 , f θ (x 1 ), ..., query x N prompt P , f θ (x N )completion where P is a series of input-output pairs followed by a lone query. The model predicts a completion based on the prompt it received. The function parameters θ and the inputs x i are randomly sampled from a function class domain and an input domain, respectively. The tasks we regress to are summarized in Table 1 and detailed in Section 3.1We train models for ICL by minimizing the expected loss over a distribution of prompts and corresponding function outputs. This approach allows us to observe qualitative differences in model architectures by their ability to behave similarly to optimal or baseline estimators. To further simplify ICL aptitude evaluation, we introduce a proxy value summarizing a given model's ICL ability for a specific task. This metric averages the error of a model normalized by the baseline error at each context length. We detail this further in Section 3.3.To determine task-specific ICL ability, our sequence models regress onto the functions shown above [14] . We replicate the function classes Linear Regression, Sparse Linear Regression, 2-Layer MLP Regression, and Decision Tree Regression from Garg et al. [6] as they present a wide range of "difficulty" for sequence models. In addition, to capture the existence of some ICL ability, we also regress onto the two function classes examined in Park et al. [14] : parity function with induced sparsity (Sparse Parity) and parallel associative recall (Vector MQAR).Unless otherwise specified, we train all models with 12 layers, 8 attention heads, an expansion factor of 4 (in the case of models with Mamba Mixer layers), and linear layers to transform the input sequences into and from the embedding dimension of 256. We use the ADAM optimizer with a learning rate of 0.0001 for 500k steps. Our expansion factor was selected to ensure similar parameter counts across baselines and all other hyperparameters were chosen for consistency with Garg et al. [6] . Note for the four function classes from Garg et al., the same curriculum was used during training.No curriculum is used for the two new function classes from Park et al. [14] . For our compute2 , we utilized 898.90 hours on an A10, 55.74 hours on an RTX 3090, 151.90 hours on an RTX 4090, 75.48 hours on an RTX 4070 Ti, and 9.83 hours on an RTX 6000.Linear Regression and Sparse Linear Regression Each function in these tasks is parametrized as a single weight vector (w) of dimension equal to that of the x-values (i.e. 20) so that y = w T x. We sample the coordinate values from a normal distribution and (in the Sparse Linear case) zero out all values except a uniformly at random selected k coordinates. In essence, one can consider LinearRegression to be the degenerate case where the k = 20. We preserve these tasks from Garg et al. [6] to verify that none of our hybrid modifications lose the near-optimal performance that was already found with GPT-2.We fill two weight matrices W (1) ∈ R 100×20 and W (2) ∈ R 1×100 with scalar samples from a normal distribution. y values are computed as the result of a forward pass through a 2-layer multi layer perceptron with a ReLU activation. That is: y = W (2) ReLU(W (1) x).This is a more complex function class that Garg et al. [6] found that GPT-2 can perform very well at, suggesting that this task can capture some ICL ability of an architecture.We construct full decision trees of depth 4 with leaf values sampled from a normal distribution and branching conditions to be selected uniformly at random over the coordinates of the input dimension. The left branch is taken if the selected input coordinate is less than 0 and the right branch is taken otherwise. Garg et al. [6] found that GPT-2 was able to achieve much lower error for lower context lengths than XGBoost or Greedy Tree Learning, suggesting that this task can capture some ICL ability of an architecture.We select k = 2 values to consider and compute their parity, expressed as either -1 or1.That is, we uniformly sample without replacement θ ∼ {1, ..., 10} k and compute y = i∈θ x[i].Along with a higher learning rate of 0.0004, this is identical to the scheme implemented in Park et al. [14] . They [14] found that GPT-2 style transformers do not perform well on this task, suggesting that this is a discerning proxy for measuring ICL ability. Finally, as convergence was quick for this task, we only trained models up to 200k steps.Vector MQAR We sample 2N points from the d-sphere of radius √ d and group them randomly into pairs to forming N key-value pairs. For consistency with the experiments of Park et al. [14] and to reliably allow for the formation of transformer circuits highly relevant to this task [22, 14] , we reduce model complexity by using an embedding dimension of 128, 2 layers, and a higher learning rate of 0.0002. Park et al. [14] found that Mamba, our representative of SSM-type models, performed poorly, suggesting that this task can serve to ensure we don't lose capabilities provided by transformers. As detailed by Radford et al. [1] , GPT-2 is almost identical to the original decoder-only transformer, with absolute positional embedding, pre-norm layer normalization, and a GELU activation function in the feed-forward network (FFN) (which is otherwise a multi-layer perceptron). In contrast, Llama [29, 30] combines a number of modern transformer modifications, including swapping layer norm with RMS norm [31] , changing the architecture and activation function of the FFN, and using rotary Table 2 : A summary of the primary architectural differences between GPT-2, Llama, and Mamba. We examine all variations between GPT-2 and Llama and all variations between Llama and Mamba.positional embeddings instead of absolute positional embeddings [32] . We acknowledge that the larger variations of Llama2 [30] and both variations of Llama3 [33] used Grouped-Query Attention (GQA), however we surmise that at our model scales of ∼10 million parameters, GQA will not significantly affect the performance of our models. From an entirely different method of sequence modeling, Mamba forgoes positional embedding entirely, combining features of the Gated LinearUnit and state space expansion to remove the need for distinct attention and feed-forward blocks.We summarize these architectural differences in Table 2 . We examine all combinations of these different components, training 12 total architectures (listed in Figure 1a ) on our 6 tasks for a total of 72 model-task pairs. Figure 1b illustrates how each of these variations compose into a model. We provide individual diagrams of each architecture in Appendix A.In addition to the baseline metric (squared error as a function of context length) from Garg et. al.[6], we've established another metric: ICL regression score. This is a scalar expressing overall performance of a model on a task. Abstractly, the metric aims to capture the proportion of the baseline error saved by a model. The regression score is calculated by (1) computing the difference in error achieved by the model and the zero estimator at each context length, (2) computing the average of this value over the length of the sequence, (3) computing the same value for the baseline estimator, and ( 4) taking the ratio of these.In summary, ICL regression score can be calculated as follows:EQUATIONwhere ξ 2b .We replicate the baseline predictors for linear regression, sparse linear regression, and MLP regression from Garg et al. [6] due to the lack of a higher-performing baseline. However, we opted to use a pretrained GPT-2 model with identical structure to that used in Garg et al. to serve as a more calibrated baseline than Greedy Tree Learning or XGBoost. They showed superior decision tree ICL performance for a trained GPT-2 transformer compared to Greedy Tree Learning or XGBoost. For consistency with Park et al. [14] and due to the algorithmic hardness of Sparse Parity, we used our Mamba model trained on this task. Park et al. showed that Mamba can effectively learn this task, so we repeat our strategy as in Decision Tree Regression with our Mamba model (instead of GPT-2) as a baseline.For ease of experimentation and reproducibility, we have built a typed, extensible, and modular Python codebase. We achieved this by identifying isolated processes in the training regime and to the experiment archetype as proposed by Garg et al. [6] and repeated by others [8, 15, 14] . We integrate standard reporting software Weights and Biases [34] and leverage fast implementations of attention [35] and 1-D convolutions [36] . We also implement a configuration-based system for training, loading, and evaluating models to facilitate frictionless repeatability of all experiments. 
results: 
 We confirm the results from Garg et al. [6] and Park et al. [14] that GPT-2 and Mamba can learn our first four regression tasks in context. Park et al. [14] that Mamba struggles to performVector MQAR while transformers and hybrid architectures excel. We note that Llama and GPT-2 have very comparable performance in Sparse Parity and Vector MQAR. We plot all qualitatively non-optimal squared error profiles in Figure 3 and all squared error profiles in Appendix B.(a) Notable phenomena for Sparse Linear. We observe that while GPT-2 (orange) performs very similarly to our baseline, adding RMS norm without RoPE (red and green) leads to models performing notably worse than optimal.(b) Notable phenomena for Decision Tree. We note that Mamba (green) performs somewhat suboptimally while GPT-2 RMS (orange) fails to learn the task entirely.Figure 3 : Squared error profiles that do not exhibit near-optimal behavior. Shaded regions are 99% confidence intervals.Models can converge to suboptimal regression schemes. We find that some model-task pairs produce suboptimal predictions, not as a result of insufficient training. A clear example is GPT-2 RMS SwiGLU (model 1.4) on Sparse Linear. This model appears to not achieve optimal error -achieving an ICL Regression Score of only 0.754, opposed to ∼ 0.93 by other models -and yet its performance does not significantly improve with more gradient steps. We plot the squared error achieved by various checkpoints for model 1.4 in Figure 4a . We observe that this error profile appears similar to that of models trained on the Linear task and so also examine the prediction quality of the (a) GPT-2 RMS SwiGLU Checkpoints on Sparse Linear. We see that GPT-2 RMS SwiGLU converges to the least squares solution, despite Lasso being the optimal solution. This suggests that GPT-2 RMS SwiGLU fails to learn to utilize its context to its fullest extent.(b) GPT-2 RMS SwiGLU trained on Sparse Linear and evaluated on Linear.When evaluated on a similar task to which it was trained on, GPT-2 RMS SwiGLU appears to perform better than its siblings, despite the fact that it performed worse than its siblings on its original task! This suggests that it learned a different regression scheme than GPT-2 on the same training data. same model (GPT-2 RMS SwiGLU trained on Sparse Linear) on Linear in Figure 4b . We find that it indeed mimics the error profile of least squares. This result builds on Akyürek et al.'s findings [19] in what functions transformer models develop representations of. Sparse Linear on adopts a suboptimal regression scheme (least squares) partway in training, eventually unlearning its scheme in favor of the optimal regression scheme (lasso). We plot the squared error on Sparse Linear achieved by various checkpoints for Model 1.3 in Figure 5a , noting that the error of the checkpoint at 100k steps closely matches the error of least squares. Further, we examine the squared errors on Linear Regression for the various checkpoints for Model 1.3 in 5b and see that the checkpoint at 100k most closely matches least squares. This suggests that model 1.3 learned the linear regression scheme in the beginning of training, but was eventually able to learn to utilize the sparse nature of its training data.Models can fail to converge within our training horizon. We find that a number of models performed strikingly poorly in their trained task. In particular, GPT-2 with Layer norm replaced by RMS norm (model 1.1) performed very poorly on Sparse Linear Regression and Decision Tree, as indicated by the lowest ICL Regression Score achieved in those tasks (0.535 and 0.114, respectively) and in Figures 3a and 3b . We also observe that GPT-2 with RMS and SwiGLU (model 1.4) also did not converge to a regression scheme, despite apparently modelling a different regression scheme entirely. Similarly, Mamba (model 3) did not converge to a training scheme on Decision Tree as illustrated in Figure 6a . We believe this suggests a lower training efficiency for certain architectures on these tasks.Models can fail to learn the task entirely. In the case of Decision Tree, GPT-2 with RMS (model 1.1) failed to learn the task entirely as not only indicated by its final ICL Regression Score but also its consistency in achieving very high error throughout training. We plot squared error for various checkpoints in Figure 6b . 3) starting by learning a more general but sub-optimal regression scheme but eventually converging to the optimal regression scheme when trained on Sparse Linear. Shaded regions are 99% confidence intervals. Regression and 2-Layer MLP, along with their corresponding graphs of squared error as a function of context length, corroborate the claims from Garg et al. [6] that transformers can "learn" these tasks.ICLFurther, the ICL Regression Scores for Sparse Parity are consistent with Park et al. [14] , with all hybrids between GPT-2, and Llama failing to "learn" the task and all hybrids between Llama and Mamba succeeding in "learning" the task. Indeed, the ICL Regression Score achieved by Mamba captures the qualitatively sub-optimal performance detailed above on Decision Tree. 
discussion: 
 Even simple function classes leave room for local minima. We find that despite distilling down the phenomenon of In Context Learning to regression against simple function classes, there still exists room for models to adopt various regression schemes. This is supported by the apparent convergence Model L in ea r ± 0 .0 0 1 S p ar se L in ea r ± 0 .0 0 1 2 -L ay er M L P ± 0 .0 6 D ec is io n T re e ± 0 .0 0 1 S p ar se P ar it y ± 0 .0 0 1 biases on certain solution forms, resulting in extreme convergence times when these solution forms greatly vary from the optimal predictor's form.Extensible Research as Reproducible Research. In the development of this work, continuously iterating to minimize the friction of reproduction has enabled rapid extension of our Python artifacts to support even abstractly defined hybrid architectures, which are often considered inextricable from highly bespoke code or dedicated packages such as xFormers [37] . We implore the reader to seriously consider the value of making their research extensible with a minimum of friction. We hope that our attempts to maximize extensibility and reproducibility contribute to the longevity of this work as a reliable, tested, and simple framework to use for studying simple function classes in context.We have only one training run performed on each model-task pair. As a result, we have no estimation for how consistently observed phenomena appear with the given architectures. We only train each model for a maximum of 500K steps. Thus, when a model fails to converge within this window, we lose information on insightful trends that could possibly occur with further training.We do not empirically evaluate the effectiveness of ICL Regression Score or the usability of our provided code platform. We compute no verifying metrics to establish how well ICL Regression Score generalizes or is robust to qualitatively distinct ICL regression tasks. Similarly, we perform no user study on the effectiveness of our code platform, presenting only our own experience. Justification: The abstract/introduction briefly covers the limitations of the paper while introducing the main claims/findings/contributions. We reference relevant work we are building off of.Guidelines:• The answer NA means that the abstract and introduction do not include the claims made in the paper.• The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.• The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.• It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?Answer: [Yes] Justification: This paper contributes its codebase, which contains sufficient information in the ReadME for reproducibility. This paper's "Methods" section discusses data generation for the simple function classes.Guidelines:• The answer NA means that paper does not include experiments requiring code.• Please see the NeurIPS code and data submission guidelines (https://nips.cc/ public/guides/CodeSubmissionPolicy) for more details.• While we encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).• The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https://nips.cc/public/guides/CodeSubmissionPolicy) for more details.• The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.• The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.• At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).• Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted. 
experimental setting/details: 
 Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?Answer: [Yes] Justification: These details were explained under the "Methods" section are sufficient to understand our results.Guidelines:• The answer NA means that the paper does not include experiments.• The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.• The full details can be provided either with the code, in appendix, or as supplemental material. 
experiment statistical significance: 
 Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?Answer: [Yes] Justification: Every figure and number presented in this paper has confidence of intervals of 95% or 99% (specified in each case).Guidelines:• The answer NA means that the paper does not include experiments.• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.• The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).• The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)• The assumptions made should be given (e.g., Normally distributed errors).• It should be clear whether the error bar is the standard deviation or the standard error of the mean.• It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.• For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).• If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text. 
experiments compute resources: 
 Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?Answer: [Yes] Justification: The list of all GPU types utilized for the experiments were included, along with the time spent for compute on each of them. Furthermore, we provide a breakdown of the time spent for each experiment type on the GPU that we utilized the most (an A10).Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.• The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.• The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper). 
code of ethics: 
 Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines?Justification: We study ICL using simple function classes and do not use real world data.No human subjects are involved and there are no direct paths for negative societal impact.Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.• The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [NA] Justification: Since this paper studies ICL in hybrid models using simple function classes, there is no direct path to negative applications.Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.• Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.• The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.• The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.• If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: As this paper discusses ICL on simple function classes, it does not utilize real-world data or models real-world capabilities. Thus, there is no risk for misuse.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.• Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.• We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes]Justification: We provide credit to the authors of the three codebases that inspired some of the features in our own and cite that their codebases fall under the MIT License for the first two, and could not be found for the third. URLs are provided for each codebase as well.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL.• The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.• If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.• For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.• If this information is not available online, the authors are encouraged to reach out to the asset's creators. 
new assets: 
 Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?Answer: [Yes] Justification: Our codebase contains a ReadME providing thorough documentation. Our paper also explains high-level functionality of our codebase.Guidelines:• The answer NA means that the paper does not release new assets.• Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.• The paper should discuss whether and how consent was obtained from people whose asset is used.• At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file. 
crowdsourcing and research with human subjects: 
 Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?Justification: There are no crowdsourcing experiments or research with human subjects, and thus no text of instructions or compensation information is included.Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.• Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper. Justification: As there were no study participants in this study, this information was not included in this paper.Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.• Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.• We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.• For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.As mentioned, our code takes notable inspiration from the code distributed by Garg et al.[6], Park et al.[14], and Lee et al.[8], which can be found at https://github.com/dtsip/in-context-learning, https://github.com/krafton-ai/mambaformer-icl, and https://github.com/ivnle/synth-icl respectively. The first two repositories are licensed under the MIT License and we could not identify the license for the third.On an A10, the approximate training time for Linear Regression and Sparse Linear Regression was 12 hours, for 2-Layer MLP Regression and Decision Tree Regression was 2 days, and for Vector MQAR was 5 hours. 
 
 Appendix:  


Decision: The paper presents an empirical study on in-context learning (ICL) using hybrid architectures built from existing large language models, but its contributions are limited, lacking both novelty and clear motivation. While the open-source codebase is a positive aspect, the study's results are difficult to interpret, and the paper fails to offer meaningful insights or intuition behind performance variations. Additionally, the writing is unclear, with important details either missing or poorly explained.

As a summary, this paper is rejected due to the following reasons:
1. The paper offers little in terms of technical innovation or new methodologies.
2. The rationale for exploring these particular hybrid architectures is not well-articulated, diminishing the study's significance.
3. Clarity issues in both the presentation and explanation of results make it difficult to derive value from the findings.

Reviews: review: paper_summary: paper_summary: This work presents an analysis of in-context learning (ICL) for a variety of hybrid architectures (composed of different blocks from preexisting large language model architectures) on different regression tasks. The experiments are built on top of a couple of prior works [1, 2] that also explored ICL in similar contexts. This paper highlights that several prior results can be reproduced, and for novel hybrid architectures which are the main focus of this work – most of them converge to optimal solutions while some others can escape suboptimal solutions or even fail to converge in the first place. The authors also propose a new metric “ICL regression score” to evaluate ICL performance in comparison to a known baseline. The modularized code for this work is publicly available for the broader scientific community.


[1] Garg S, Tsipras D, Liang PS, Valiant G. What can transformers learn in-context? a case study of simple function classes. Advances in Neural Information Processing Systems. 2022.
[2] Park J, Park J, Xiong Z, Lee N, Cho J, Oymak S, Lee K, Papailiopoulos D. Can mamba learn how to learn? a comparative study on in-context learning tasks. arXiv preprint arXiv:2402.04248. 2024.
strength_weakness: strength_weakness: - The paper is not very well-motivated. Why are hybrid architectures (especially the two that are focused on) important to study? What intuitions or profound reasons drive the authors to make the experimental design choices that they did? 
- Additionally, use cases for ICL itself are not well-motivated. Are there any practical use cases that warrant such extensive evaluation? The writing is not easy to understand for a reader not very up-to-date with the ICL literature. 
- The technical novelty of the work is limited.
- The results are presented in a manner where the performance metrics are reported for the 12 architectures and 5 tasks but it is not very clear what the reader or the scientific community working on ICL should take away from the results. Are there patterns regarding why certain hybrid architecture + task combinations make ICL shine compared to the baselines and why some others do not? A lot of the interpretation of such results is left to the reader to figure out. A lack of a deeper understanding and intuition about the reasons behind the results makes it hard to see solid/impactful takeaways that others could build on top of. 
- The authors mention and describe a 6th task Vector MQAR but do not report or discuss any of its results in detail in the main text of the paper. One figure is present in the Appendix but it is not explained and it is too hard to read the text in the figure.
- Some typos:
     - Line 161: Mention the word “**Figure**” before 2a.
     - Line 164: “**Figure 2b**” instead of “Table 2b”
     - Line 185: Park et al. [14] **show** that ..
     - Line 202: “Sparse Linear ~~on~~ adopts a suboptimal..”
- References to result tables (Table 3 for lines 213, 230)  and model descriptions (Figure 1a for lines 194, 204, 201, 16 etc.) could further enhance readability and the user’s understanding.
questions: questions: - What is the main motivation behind studying hybrid architectures for ICL? For a reader not well-versed with the ICL literature, it is not clear from the paper. 
- Why are the specific hybrid architectures chosen (GPT2-Llama and Llama-Mamba)? Is it only because no one has studied them before or are there other profound reasons why these could be of interest to study? 
- On line 35, what is meant by “richly benchmarking”?
- Could you briefly describe what kind of structural changes were necessary to the codebase of Garg et al.? A general understanding will make it more clear why the codebase itself is a major contribution as listed currently in the draft.
- Lines 60-61: What optimal train loss is being referred to here? There is no loss associated with ICL itself, right (as that only consists of in-context examples/demonstrations)? Does this loss refer to the ground truth functions being learned or the baselines? 
- Line 80: which models are being referred to here? There is no “real training” during ICL itself. Are the functions/regression targets being referred to here or the baselines? The terminology used in the paper should be clarified.
- It is not clear what data the 12 hybrid architectures are trained on. Do the authors train them from scratch? If yes, what datasets are used? Else are pretrained weights used for the different blocks? Are the models finetuned? Something else?
- Why aren’t the baselines and detailed results for Vector MQAR discussed in the paper draft? Why does Mamba struggle with MQAR but does pretty well on all other 5 tasks?
- Why are some hybrid architectures converging to suboptimal regression schemes versus others are not? Do you have some intuition behind this finding?
- Similarly, why do some models escape suboptimal regression schemes while others fail to converge entirely? 
- For Fig 6(a), could you try training Mamba checkpoints longer than 500k checkpoints to verify if it actually converges to the baseline (GPT-2’s performance)? Were these checkpoints trained by the authors themselves or used from a pretrained model? What data was used for training?
- What do the authors mean by “context length” in most of their figures? Does it refer to the number of ICL examples or the number of prompt tokens? Can the authors share a few examples of sample prompts that were used for ICL with varying context lengths?
- In Figure 6(b), what about GPT2-RMS might explain its inability to learn the decision tree function while all other hybrid architectures used do pretty well? Why is the RMS norm as the distinguishing factor for this hybrid architecture detrimental to the decision tree and sparse linear tasks but not others?
- In lines 244-246, the authors mention that “certain hybrid architecture variations may place inductive biases on certain solution forms, resulting in convergence times when these solution forms greatly vary from the optimal predictor’s form”. Can the authors give some examples of such inductive biases for some architectures studied in this work?
- Why are some values in the last column of Table 3 not filled in?
limitations: limitations: In terms of negative societal impacts, could the potential misuse of natural language tasks via the hybrid architectures presented be a legitimate concern? It is certainly out of scope of this work for evaluation but could be listed as part of the broader impacts section of the checklist.
paper_summary: paper_summary: This work presents a codebase for benchmarking the in-context learning ability of language models, especially for hybrid models. In addition, several empirical results are presented to show that some model architectures fail entirely or have suboptimal performance on specific in-context learning tasks.
strength_weakness: strength_weakness: * I feel that it is hard to assess the contribution of this work. It seems that this work's main contribution is the implementation of the in-context learning ability benchmark codebase. While such a codebase is important and useful, I did not find what technical challenges the codebase is trying to address and the effectiveness of the codebase.
* Another contribution is the empirical findings of the relationship between architectures and per-task performance on in-context learning. However, I found the empirical results are not systematic and hard to interpret. I am unsure how these findings motivate future architecture design.
questions: questions: As stated above, what's the main contribution/innovation of the codebase? What conclusion/takeaway/intuition can we learn from the empirical results?
limitations: limitations: see above
paper_summary: paper_summary: The authors build on the ICL work of [1] and [2], wherein networks are "trained to in-context learn" several tasks of varying complexity (e.g., Linear Regression, Sparse Linear Regression, Vector MQAR, etc.).  Carrying on from [2], the main contribution of the presented work is the combination of various permutations of the attention/SSM blocks from GPT-2, LLama-2, and Mamba models.  Such permutations include swapping the Layer Norm in GPT-2 attention blocks with RMSNorm, the LLama-2 SwigGLU with a Mamba Mixer, and so on; in total, 9 different hybrid combinations are considered.  For each task and specific hybrid model, the model is trained over task samples for 500k steps, then evaluated on ICL performance on that task.  In addition to reporting squared error per task, the authors also propose a new metric, called the "ICL regression score."  The results of specific model-task pairs are plotted and several trends are discussed.

[1] Shivam Garg, Dimitris Tsipras, Percy S Liang, and Gregory Valiant. What can transformers learn in-context? a case study of simple function classes. Advances in Neural Information Processing Systems, 35:30583–30598, 2022.
[2] Jongho Park, Jaeseung Park, Zheyang Xiong, Nayoung Lee, Jaewoong Cho, Samet Oymak, Kangwook Lee, and Dimitris Papailiopoulos. Can mamba learn how to learn? a comparative study on in-context learning tasks. arXiv preprint arXiv:2402.04248, 2024.
strength_weakness: strength_weakness: ### Clarity
There is significant room to improve the clarity of the presented work.  Currently, several important details are missing (discussed below), key concepts are not fully explained, and the current paper could benefit from an editorial pass (it is currently difficult to read).  It is also important to note that, while the authors presented the results of model + task pairs and detailed where models failed, no possible explanations and follow up ablation studies were conducted; the question of "why" remains for all the presented results.  E.g.,
> Specific hybrid architectures can hesitate to learn/converge for certain function classes.

Do the authors have intuition or an explanation which can be explored to explain this?  Arguably, this is one of the most important contributions to be made in a large empirical review like the presented paper, i.e., to make sense of the experiments to gain a greater intuition for why an LLM behaves in a certain way.

Wrt important missing information:
- "We replicate the function classes Linear Regression, Sparse Linear Regression, 2-Layer MLP Regression, and Decision Tree Regression from Garg et al. [6] as they present a wide range of "difficulty" for sequence models. In addition, to capture the existence of some ICL ability, we also regress onto the two function classes examined in Park et al. [14]: parity function with induced sparsity (Sparse Parity) and parallel associative recall (Vector MQAR)." <- - How training instances are produced per task? How many test samples are produced per task?  If this follows Garg et al., then each model is trained *from scratch* on 40 samples per task. Can you please clarify and state these in the main text?
- It is not clear	what the author's mean by "zero estimator."  Is this the zero shot prediction? Correspondingly, it is not clear exactly what the presented ICL regression score represents.
- "To determine task-specific ICL ability, our sequence models regress onto the functions shown above [14]."  <- It would help to clearly state the paper trains the models "from scratch" to in-context learn, as in previous works.
questions: questions: See weaknesses.
limitations: limitations: The did well to state limitations of the presented evaluation.
paper_summary: paper_summary: The authors explore a number of different attention-based architectures along with Mamba on a series of in-context regression tasks. The architectures vary in their choice of normalization, positional encodings, activations, and hybridization with Mamba. They discover some varying capacities for the different ICL tasks.
strength_weakness: strength_weakness: As mentioned above, the present work seems to be closely related to [Park et al.](https://arxiv.org/abs/2402.04248), with little added insight. While an attempt has been made to explore minor architectural aspects like the choice of normalization, positional encodings, and activations, for the most part the changes seem to matter little (according to Table 3). I'm unsure what the key take-away is. Is there a particular architectural configuration that works best? What concrete practices can a user apply to improve their models' performance? The authors appear to have an ambition towards answering questions like these, but do not ultimately resolve them.

The results are sometimes difficult to interpret. At times, this is simply because the plots are unreadable, with text that is too small. I'm unsure how to interpret the in-context regression score. It looks like it's often <1 across models. Does this mean they all fail to outperform the baseline? Is this score comparable across different tasks?

A main objective of the paper is to compute numerics comparing models, but only one training run was executed for each experiment. If the compute budget is very limited, a more valuable approach may be to consider a smaller subset of models (e.g. keeping either GPT-2 or Llama, but not both) and simpler task parameterizations. Doing so will enable you to sweep across many more settings and increase your experiment replications, generating more convincing numerics. 

Additional minor formatting comments:
- Line 25: consider compressing citations (e.g. with sort&compress)
- Table 1: third hline from the top intersects with text
- Consider plotting figures with point-markers for each data point, to clarify where exactly your data points fall
- For related models that vary by parameter (e.g. training iteraitons), consider using the same color but with different shadings / style
- Figure 15: text is unreadable
- Figure text overall is small and hard to read
questions: questions: - Line 255 mentions that there is only a single run performed for each model-task pairs, yet many plots feature confidence intervals. Where do these confidence intervals come from?
- Why do you pick the model subsets that you do in constructing each plot? For instance, plot 3b has a title "Decision Tree: All Architectures" but only plots GPT-2 and Mamba (what about Llama, and the other variations?)
- Why apply Mamba only to Llama, and not to GPT-2? Since you mentioned compute was a bottleneck, why not keep either GPT-2 or Llama as a base model?
- How is the ICL regression score meant to be interpreted? It seems to depend strongly on the task, where some tasks will feature a score uniformly less than 1 across models, and greater than 1 for others.
limitations: limitations: The authors adequately state the limitations of their approach.

Comments: title: Acknowledgement of rebuttal
Acknowledgement of rebuttal
Author Response (part 2/2)
Author Response
Author Response
Author Response
Author Response (part 1/2)
Clarification to motivations and contributions
comment: I thank the authors for the detailed answers to all my questions. The responses presented in the rebuttal can greatly improve the readability of the draft. I agree with other reviewers that the paper has further scope for improvement. I will thus retain my score and wish the authors the best in preparing a revised version of this work.
I thank the authors for their response.  I've read the other reviews and agree the submission looks to be in an unfinished state and the presented experiments need further exploration to better distinguish from previous work (as well as articulate what any major takeaways are).  I am retaining my score and wish the authors the best in future revisions.
Thank you for your detailed rebuttals and clarifying comments. After reading the other reviews and rebuttals, I concur with the other reviewers in that 1) the impact of these experiments remains unclear and 2) the present manuscript seems to be in a rough, unfinished state. For these reasons, I retain my current score.
We respond to clarifying questions in this response.

We trained all models from scratch, using standard initializations. The training data is synthetic: randomly sampled vectors $x_i$ for each in-context example and only one set of parameters $ \theta $ for each sequence arranged as illustrated in Section 3. The distributions are summarized in Table 1. For example, a batch of $8$ sequences with 40 in-context examples each results in $8 \times 40 =  320$ sampled $x$ values and $8$ sampled sets of parameters ($ \theta $). 

We conflate “context length” with the number of $ \( x, f(x) \) $ pairs (i.e. in-context examples). This makes the number of prompt tokens simply one more than double the “context length”. This is because in-context examples must have both an input value $x$ and label value $ f(x) $ *and* the prompt must have a query token of an unlabeled input value $ x_{query} $. See the beginning of Section 3 for a (hopefully helpful) visual.

For example, consider a linear function of input dimension 3 (we use 20 in the paper) and weights [1, 1, 1], producing a prompt of “context length” 2: 

$$ \\begin{bmatrix} 1 \\\\ 2 \\\\ 1 \\end{bmatrix}, \\begin{bmatrix} 4 \\\\ 0 \\\\ 0 \\end{bmatrix}, \\begin{bmatrix} -1 \\\\ 1 \\\\ 0.5 \\end{bmatrix}, \\begin{bmatrix} 0.5 \\\\ 0 \\\\ 0 \\end{bmatrix}, \\begin{bmatrix} 0 \\\\ 0.2 \\\\ 1 \\end{bmatrix} $$

In this case, the regression target is 1.2.

Another example: consider a sparse linear function of input dimension 3 and weights [0, 1, 0], producing a prompt of “context length” 3: 

$$ \\begin{bmatrix} 0.6 \\\\ 0.4 \\\\ -0.2 \\end{bmatrix}, \\begin{bmatrix} 0.4 \\\\ 0 \\\\ 0 \\end{bmatrix}, \\begin{bmatrix} -0.3 \\\\ 0.8 \\\\ -0.1 \\end{bmatrix}, \\begin{bmatrix} 0.8 \\\\ 0 \\\\ 0 \\end{bmatrix}, \\begin{bmatrix} 1 \\\\ -1.4 \\\\ 0.4 \\end{bmatrix}, \\begin{bmatrix} -1.4 \\\\ 0 \\\\ 0 \\end{bmatrix}, \\begin{bmatrix} 0.1 \\\\ -0.8 \\\\ 1.1 \\end{bmatrix} $$

In this case, the regression target is -0.8.

---

Where we say “optimal loss,” there are closed-form predictors that achieve theoretically minimal (average) loss at each context length. For example, optimal loss at linear regression with squared-error loss is that of the least squares regression line (as opposed to 0 error at all context lengths, which is the ground truth). As far as the authors are aware, regressing to a random 2-layer MLP or depth 4 decision tree does not have an optimal scheme, so we use the more general term “baselines” to refer to regression schemes that serve as points of comparison.

Line 80 refers to transformer-/mamba-like models. We defined our training loss to be the average squared error between the model’s prediction $ \hat f_{\theta}(x_{query}) $ and the ground truth $ f_{\theta}(x_{query}) $. Thus our training scheme was to minimize expected loss over many sampled functions.

We continued training Mamba up to 1 million gradient steps and found that it continued to (very slowly) approach the regression scheme achieved by GPT-2. As of this posting, our 1.1 million step checkpoint only continues this trend, but does not yet reach the same error profile.

The last column of Table 3 contains unfilled entries because we inferred that since both GPT-2 and Llama failed entirely at sparse parity, that hybrid architectures would perform similarly. As a result, we omit those model-task pairs.

The form of the Vector MQAR we replicated did not consider querying models at multiple context lengths, so using any baseline to compute a regression score would reduce to multiplication by a scalar value. This makes formulating a baseline for Vector MQAR not meaningful in our context. We did not discuss detailed results for Vector MQAR as our results almost exactly matched direct extrapolations from the results of Park et al. We will expand discussion of these results in the paper.

We do not have a definitive answer for why Mamba only struggles with MQAR. Park et al. [1] propose that the latent representation space expressible by Mamba for this task (a vector of dimension 512 per layer) is not expressive enough for this task (which calls for storing 128 pairs of 20-dimensional vectors, or 5120 real values) and becomes “overwhelmed” with long contexts. We generally agree with this, and also affirm that any recurrent model’s latent space along the sequence length can be saturated this way. To clarify, the models trained on this task are smaller than those from other tasks to ensure that the attention-based models converged, suggesting that computing a gradient from only a single regression target per sequence limits the complexity of models that are trainable on this task.

Again, we thank you for the complete review and helpful guidance.

[1] Jongho Park, Jaeseung Park, Zheyang Xiong, Nayoung Lee, Jaewoong Cho, Samet Oymak, Kangwook Lee, and Dimitris Papailiopoulos. Can mamba learn how to learn? a comparative study on in-context learning tasks. arXiv preprint arXiv:2402.04248, 2024.
We thank you for your thoughtful feedback and helpful comments. We briefly respond to some of the above questions and apologize for the otherwise low quality submission.

The reviewer above highlights the limited scientific findings of the paper and proposes some avenues for uncovering some by building on the existing analysis.

We indeed only perform a single training run for each model-task pair. We constructed confidence intervals from a distribution of errors achieved on 2048 batches of 64 functions and corresponding input points sampled from the task distributions (detailed in Table 1). This resulted in approximately 131k error values for each model at each context length, on which we assumed a normal distribution (due to similarity by visual inspection) and computed confidence intervals.

In constructing each plot, we picked specific model subsets to increase readability. For instance, in plot 3b titled "Decision Tree: All Architectures," we only plot GPT-2, GPT-2 RMS, and Mamba. In all plotted tasks, we refrained from plotting more than one architecture that achieved near-optimal or otherwise identical error distributions.

Mamba was applied only to Llama and not to GPT-2 to avoid repeating the contributions of Park et al. [1] addressing GPT-2/Mamba hybrid architectures. We analyzed hybrids between Llama and Mamba as representatives of attention-based and attention-free architectures because (to our knowledge) this was novel. We chose to also examine architectural “steps” between GPT-2 and Llama to verify that modern advancements in NLP are captured by our proposed regression score. We confirm this (admittedly loosely) by observing that native Llama obtains statistically similar or superior regression scores across all tasks. 

The proposed ICL regression score is intended to summarize the overall “performance” of a specific architecture on a certain task. It compares the predictive error of an architecture against that of a baseline. The ICL regression score depends strongly on the accuracy of our chosen baseline for each task and cannot be compared between tasks. For example, with a sub-optimal/poor baseline as used in the decision tree task, ICL regression scores will typically be higher. Conversely, for an information-optimal baseline such as in the linear regression or sparse linear regression tasks, regression scores will be strictly no greater than 1. Intuitively, the score should allow practitioners to quickly compare candidate architectures (such as GPT-2 RoPE SwiGLU vs Llama SwiGLU-less), but will hold little arithmetic value.

We believe that the primary contributions of this paper are (1) the development of a codebase amenable to open-source contributions and (2) the creation of a score that can compare models within a synthetic regressive task. We thank the reviewer for clarifying that this is an obscured claim in the current draft.

We once again thank you for the detailed review and informative suggestions.

[1] Jongho Park, Jaeseung Park, Zheyang Xiong, Nayoung Lee, Jaewoong Cho, Samet Oymak, Kangwook Lee, and Dimitris Papailiopoulos. Can mamba learn how to learn? a comparative study on in-context learning tasks. arXiv preprint arXiv:2402.04248, 2024.
We appreciate your diligent feedback and practical critique. We attempt to clarify some details missed by the submission below.

We trained all checkpoints used in our analysis from scratch on synthetic data generated by sampling parametrized functions and x values from the distributions specified in Table 1. Those sampled x values and functions are then collected into prompts that follow the structure described in Section 3. We will add this with more clarity to the main text.

The “zero estimator” predicts 0 for all regression targets, which in all tasks is the estimate with minimal expected squared error conditioned on *no* context. 

The ICL regression score aims to summarize the overall “performance” of a specific architecture on a certain task. It compares the predictive error of an architecture against that of a baseline. ICL regression score strongly depends on the accuracy of our chosen baseline for each task and consequently cannot be compared between tasks. For example, with a poor “baseline” as we use in the decision tree task, ICL regression scores will typically be higher. Conversely, for an information-optimal baseline such as in the linear regression task, regression scores will be strictly no greater than 1. Intuitively, we suggest treating it as a “score” achieved by an architecture on a task – one that holds little arithmetic value, but can allow practitioners to quickly compare candidate architectures.

We would like to reiterate that we are grateful for the reviewer’s genuine and constructive feedback.
We thank you for the earnest review and encouraging feedback. We otherwise apologize for the unpolished writing.

The reviewer above identifies the primary contribution of this paper to be the associated code but finds that the empirical analysis is lacking systematism, motivation, and a digestible upshot.

We cede that the paper does not explicitly present a digestible result for practitioners, but we claim that the ICL regression score we detailed can serve as a quick-and-dirty tool for comparing the ICL ability of architectural variations. 

As for our code, we improved the modularity and implementation of abstractions over existing codebases such as in Garg et al. [1], Park et al. [2], and Lee et al. [3]. We allowed for extensions to architecture, benchmarking, and even training schemes to occur often independently and simultaneously without causing conflicts in standard version control systems like Git. A particularly illustrative example of this is our yaml-only hybrid architecture specification, which was developed concurrently with our evaluation tooling. We also used best practices in open-source repository management such as continuous integration, behavioral testing, and reproducible development and deployment environments.

Again, we thank you for the sincere review.

[1] Garg S, Tsipras D, Liang PS, Valiant G. What can transformers learn in-context? a case study of simple function classes. Advances in Neural Information Processing Systems. 2022.

[2] Jongho Park, Jaeseung Park, Zheyang Xiong, Nayoung Lee, Jaewoong Cho, Samet Oymak, Kangwook Lee, and Dimitris Papailiopoulos. Can mamba learn how to learn? a comparative study on in-context learning tasks. arXiv preprint arXiv:2402.04248, 2024. 

[3] Ivan Lee, Nan Jiang, and Taylor Berg-Kirkpatrick. Is attention required for ICL? Exploring the Relationship Between Model Architecture and In-Context Learning Ability. arXiv preprint arXiv:2310.08049, 2023.
We are grateful for your detailed feedback and thorough critique of the submission. We provide a clarification to questions about motivating the submission in this response and reserve responses to specific questions for another response. We apologize for the poor submission.

The reviewer above details limited motivations, novelty, and clarity of the paper, while mentioning the potential usefulness of an improved codebase.

Leading language models have ballooning inference costs, and finding a deeper correlation between architectural sub-blocks and ICL capacities can help inform designing for the cost/performance trade-off at inference time. Since in-context learning is the primary avenue by which LLMs are useful to a wide range of tasks, we propose that it is then useful to gauge the ICL ability conferred by small architectural choices. To this end, we contend that since there exists some synthetic ICL regression/classification tasks with varying performance across architectures, it is reasonable to assume that there is a set of benchmarks that can capture different qualitative capacities that we as humans intuitively aggregate into “ICL ability”. As a result, performance on such a set of benchmarks should be indicative of ICL capacities in language models. We referred to a suite of benchmarks that approach this comprehensiveness as “richly benchmarking” on line 35.

Existing work on synthetic ICL tasks such as those in the paper examine either GPT-2 or attention-free architectures (such as Mamba in Park et al. [1], a wide array of SSMs and recurrent models in Lee et al. [2]). The only exception we are aware of to this is the inclusion of Llama in the survey performed by Lee et al. [2]. As a result, we selected Llama and Mamba as architectures representative of current work in NLP for both attention-based and attention-free architectures and chose mixes between them to examine the change (if any) in the ICL capacities of both attention-based architectures produced by recent advances and hybrid attention architectures.

Our changes to the codebase were primarily logistical, but constitute an overhaul from repositories that build directly from Garg et al. [3]. We improved distinctions between abstractions and their implementations, enabling (1) extensions to model architecture, evaluation techniques, and even training schemes to occur often independently and simultaneously without causing version control conflicts and (2) best practices in open-source repository management such as continuous integration, behavioral testing, and reproducible development/deployment environments. We contend that this is a necessary contribution to facilitate examination of such small architectural variations in this paper and future work.

[1] Jongho Park, Jaeseung Park, Zheyang Xiong, Nayoung Lee, Jaewoong Cho, Samet Oymak, Kangwook Lee, and Dimitris Papailiopoulos. Can mamba learn how to learn? a comparative study on in-context learning tasks. arXiv preprint arXiv:2402.04248, 2024. 

[2] Ivan Lee, Nan Jiang, and Taylor Berg-Kirkpatrick. Is attention required for ICL? Exploring the Relationship Between Model Architecture and In-Context Learning Ability. arXiv preprint arXiv:2310.08049, 2023.

[3] Garg S, Tsipras D, Liang PS, Valiant G. What can transformers learn in-context? a case study of simple function classes. Advances in Neural Information Processing Systems. 2022.
We deeply thank the reviewers for their thorough and thoughtful responses. We apologize for the rough writing and late rebuttal. We understand that this may not be read, but we would like to have this collected with this submission for future reference.

Leading language models have ballooning inference costs, and finding a deeper correlation between architectural sub-blocks and ICL capacities can help inform designing for the cost/performance trade-off at inference time. Since in-context learning is the primary avenue by which LLMs are useful to a wide range of tasks, we propose that it is then useful to gauge the ICL ability conferred by small architectural choices. To this end, we contend that since there exists some synthetic ICL regression/classification tasks with varying performance across architectures, it is reasonable to assume that there is a set of benchmarks that can capture different qualitative capacities that we as humans intuitively aggregate into “ICL ability”. As a result, performance on such a set of benchmarks should be indicative of ICL capacities in language models. We referred to a suite of benchmarks that approach this comprehensiveness as “richly benchmarking” on line 35.

Existing work on synthetic ICL tasks such as those in the paper examine either GPT-2 or attention-free architectures (such as Mamba in Park et al. [1], a wide array of SSMs and recurrent models in Lee et al. [2]). The only exception we are aware of to this is the inclusion of Llama in the survey performed by Lee et al. [2]. As a result, we selected Llama and Mamba as architectures representative of current work in NLP for both attention-based and attention-free architectures and chose mixes between them to examine the change (if any) in the ICL capacities of both attention-based architectures produced by recent advances and hybrid attention architectures. 

We propose an ICL regression score to summarize the overall “performance” of a specific architecture on a certain task. It compares the predictive error of an architecture against that of a baseline. The ICL regression score depends strongly on the accuracy of our chosen baseline for each task and cannot be compared between tasks. For example, with a sub-optimal/poor baseline as used in the decision tree task, ICL regression scores will typically be higher. Conversely, for an information-optimal baseline such as in the linear regression or sparse linear regression tasks, regression scores will be strictly no greater than 1. This score can serve as a quick-and-dirty tool for comparing the ICL capacities of candidate architectures. 

Our changes to the codebase were primarily logistical, but constitute an overhaul from repositories that build directly from Garg et al. [3]. We improved distinctions between abstractions and their implementations, enabling (1) extensions to model architecture, evaluation techniques, and even training schemes to occur often independently and simultaneously without causing version control conflicts and (2) best practices in open-source repository management such as continuous integration, behavioral testing, and reproducible development/deployment environments. We contend that this is a necessary contribution to facilitate examination of such small architectural variations in this paper and future work.

We would like to again thank the reviewers for their time and consideration.

[1] Jongho Park, Jaeseung Park, Zheyang Xiong, Nayoung Lee, Jaewoong Cho, Samet Oymak, Kangwook Lee, and Dimitris Papailiopoulos. Can mamba learn how to learn? a comparative study on in-context learning tasks. arXiv preprint arXiv:2402.04248, 2024. 

[2] Ivan Lee, Nan Jiang, and Taylor Berg-Kirkpatrick. Is attention required for ICL? Exploring the Relationship Between Model Architecture and In-Context Learning Ability. arXiv preprint arXiv:2310.08049, 2023.

[3] Garg S, Tsipras D, Liang PS, Valiant G. What can transformers learn in-context? a case study of simple function classes. Advances in Neural Information Processing Systems. 2022.