MLPs Learn In-Context on Regression and Classification Tasks

Thanks for the comments! We’re glad you found the paper to be well-written and clear. Answering your questions:

The contribution of this work is hard to identify, and the results are not novel nor inspiring. First, it has been widely studied that MLP can learn relational data, i.e., can predict for set/sequence inputs, i.e., can learn in-context. Thus the studied topic is not novel.

We’re sorry to hear your thoughts, and will strive to clarify our contributions in our manuscript. We are unaware of prior work that has demonstrated MLPs can learn in-context, but would be very interested to know what references you had in mind? There is certainly related work we cite that rely on architectural modifications like relational bottlenecks or graph neural networks purpose-built for relational tasks, but none as far as we are aware that examine unmodified MLPs or Mixer models. Indeed, there is active literature arguing that vanilla MLPs cannot learn in-context on relational tasks (Boix-Adsera et al make an especially strong claim in this regard. See also Webb et al, Santoro et al, and Marcus), which we refute in our manuscript.

Second, it is also known that naive MLP for relational prediction (concatenating input) has disadvantages, e.g. lack of capturing permutation-invariance or variant input number. But this paper does not point out this keys, which would be helpful to explain why the MLP-Mixer and RB MLP can learn in-context better.

Thanks for the note! It is indeed surprising that the disadvantaged, naive MLP outperforms the Transformer on our relational tasks, despite the weaknesses you mention. Documenting surprising performance gaps like these is a core contribution of our work. We will emphasize these points further in our manuscript.

Third, the result that MLP competitively with Transformers given the same compute budget is obtained with a synthetic setting, considering above nature of MLPs, it is not convincing enough to be generalizable.

Great question. We were interested in cleanly evaluating ICL performance alone, and sought to avoid confounds introduced in naturalistic settings like developing fluency in natural language or learning complex visual features (which are certainly important topics of study in their own right, but outside the scope of our investigation). As a result, we drew from simple, synthetic ICL tasks common in the literature. We do cite other work demonstrating that MLP Mixers in particular continue to perform competitively with Transformers in real-world NLP and computer vision settings (for example Liu et al, Tolstikhin et al), suggesting that our conclusions may continue to hold in these cases as well.

In playful but intuitive words, while RB MLP and MLP-Mixer can be viewed as MLP Pro for relational data (including ICL), Transformer can be viewed as MLP Pro Max. So while it is mainstream to study ICL with transformer, it is not meaningful enough to study ICL with MLP.

We appreciate the fun and quirky analogy! To clarify, we are not proposing that Transformers should be replaced wholesale by MLPs. Rather, MLPs are an important class of neural networks that are still intensely studied in the present day, with hitherto untested capabilities on in-context learning. That ICL can be accomplished by such a simple model was very surprising to us, and runs against the conventional wisdom that ICL is uniquely suited for attention-based architectures. To improve our understanding of this phenomenon, it is very worthwhile documenting the properties of these models on commonly studied in-context learning tasks.

Could the authors provide empirical results on real-world few-shot learning data?

Thanks for the suggestion! As mentioned above, our focus has deliberately been on a synthetic setting where the models’ specific competency for ICL can be clearly interrogated, without the confounds introduced by real-world data. Naturalistic data falls outside the scope of this manuscript, but we look forward to investigating it further in the future.

Thank you for the equally fascinating and thought-provoking comments! We’re glad you enjoyed reading the paper, and found it well-written and compelling.

The character limitations in openreview prevent us from fully engaging in your discussion points within a single comment box, so we've split them in two. Apologies for the inconvenience! Part I below:

The paper compares transformers applied auto-regressively against MLP and MLP mixer models that access the data all at once. This seems like a strange choice for a fair comparison especially when trying to compare training FLOPs. I strongly suggest including a bi-directional transformer as a fairer baseline.

Our apologies, we do indeed run experiments in an auto-regressive setting (Fig. 5), where we demonstrate that the exact same results continue to hold. In this case, the MLP and Mixer models are also prompted auto-regressively, with zero-padding up to max length to handle variable-length inputs. The reference to this figure was easy to miss, so we’ll clarify our manuscript to emphasize its existence. Note: in the non-autoregressive setting, the Transformer is not prompted auto-regressively, so it too sees the entire input before being asked for an output, just like the MLP and Mixer. Incorporating bi-directionality in this case may not produce a difference.

The paper does not really offer a concrete definition of what ICL is. “In-context learning (ICL) refers to a task paradigm where exemplars from a novel task are presented during inference time rather than during training. ” One could argue working out a linear projection for a “new point” is a new task and hence standard linear regression is “in-context learning”. For clarity I think promoting this sort of discussion about what is and isn’t ICL is a strength of the paper, but I would like to the see the authors add more discussion on this.

Interesting take! We implicitly follow the definition of ICL commonly used in the literature in these synthetic settings (for example, in Garg et al, Zhang et al), where a “task” is defined to be a member of a function class $\mathcal{F}$ such that for every task/function $f \in \mathcal{F}$, the model can correctly predict the response to an input $x_1, f(x_1), x_2, f(x_2), \ldots, x_q$. Under this definition, an experiment consisting of a single linear regression example would constitute a function class $\mathcal{F}$ that contains a single member. While one could conceivably label this as “in-context learning,” it is a bit pathological as in-context learning typically implies that the model is generalizing to unseen, novel functions in $\mathcal{F}$, which we stipulate for our experiments. We will certainly add more discussion on this topic.

The paper only consider small toy problems which makes it hard to know if these observation generalise to more real world problems.

Great point. We were interested in cleanly evaluating ICL performance, and sought to avoid confounds introduced in naturalistic settings like developing fluency in natural language or learning complex visual features (which are certainly important topics of study in their own right, but outside the scope of our investigation). Hence, we drew from simple, synthetic ICL tasks common in the literature. It is an important direction of future research to see how these conclusions may break down as more complexity is introduced to the task.

Thank you for answering the computational comparison questions your choices do indeed sound the most reasonable given the clarifications. In light of your response I will raise my score from a 5 to a 6, and the soundness to a 4. While the work is of undoubtedly of high quality and thought provoking, I believe the overall contribution is still limited given the scope of the toy problems considered, and thus I feel unable to raise further.

Thank you for the comments and suggestions! We’re glad you found the study to be thorough, impactful, and convincing. Answering your questions:

The paper could have included real regression data. Most existing literature focuses on synthetic tasks, and exploring real data (even simple regression datasets) with somewhat complex underlying distributions would have added valuable insights.

Great suggestion. Our focus was primarily on measuring in-context learning in a controlled setting, leading us to synthetic ICL tasks. Exploring naturalistic regression data sounds like an excellent next step to examining more complex settings.

Can you provide more insights on why transformers are failing in relational tasks?

Great question. To clarify this result, Transformers in fact succeed quite beautifully on relational tasks. However, comparatively smaller MLPs also succeed beautifully, and (surprisingly) generalize better on out-of-distribution perturbations. It remains unclear where this compute gap originates, but is a topic of great interest under study in a separate, theory-oriented follow-up we are preparing. We will update our manuscript to clarify this nuance.

Thanks for the comments and suggestions! We’re glad you found our paper to be well-written and convincing. Answering your questions:

The paper in my opinion overclaims the signifiance of the work, of how surprising the findings are. MLPs are universal function approximators, and ofc, can to some extend approximate self-attention layers. Its nevertheless somewhat interesting that gradient descent can install such solutions into architectures purely consisting of MLPs. It is, especially on tractable problems such as linear regression / classification, clear that, if optimized well, neural networks will find / approximate the (known) Byaes optimal solution of these problems. I therefore not find the results very surprising.

Thanks for raising an interesting point of discussion. MLPs are indeed universal function approximators, though it remains unknown at what cost – solving an in-context problem might require an intractably wide MLP. It is also unclear, as you note, whether an appropriate solution can be discovered by gradient descent. Given the field’s focus on attention-based architectures, it was very surprising to us that an MLP can indeed learn in-context with compute comparable to a Transformer’s, especially on tasks that are most heavily studied using attention-based constructions.

I would benefit the authors to highlight that Transformers architectures dynamically allocating compute / memory based on its in-context length. This is a unique feature, when comparing to RNNs or MLPs. Even if they might be performing similarlry to MLPs for a given sequence length given a fixed memory, compute budget, the flexiblity of Transformers is their strength.

This is a great point, and we will add it to our manuscript. Thanks!

The authors, afaiu, do not run experiments in an autoregressive model as e.g. Garg et al., 2022 (What Can Transformers Learn In-Context? A Case Study of Simple Function Classes) or von Oswald et al. 2023 (Uncovering mesa-optimization algorithms in Transformers). I find this setting very important, see Questions below.

Our apologies, we do indeed run experiments in an autoregressive setting (Fig. 5), and show that the exact same qualitative results hold. We will clarify this point further in our manuscript, and consider moving this figure out of the appendix.

Can you please provide additional experiments and provide analyses of these results when training autoregressively as Garg et al., 2022 or von Oswald et al. 2023.

Gladly! The results of these experiments were discussed in Fig. 5, and show that the exact same results continue to hold in an autoregressive setting. We will clarify our manuscript to make these results easier to find.

If the MLPs / MLP mixers networks are chaning from in-weights to in-ciontext learning, do they approximate e.g. gradient descent in their archicture or mimic functionally of the self-attention layers. For the MLP mixer variants, I would find an analyses of functional similarity between architectures interesting. Can you e.g. train read-out layers at different depths of the trained network which approximate the solution / optimal prediction similarly when going down the network. It would be interesting to see if the networks are comptuing similar things at the same depth of the network (and gradually approximating the final solution). One could even try to see if the output of the self-attention and MLP mixers are similar.

These are really interesting suggestions to examine, thank you! You raise a great question in asking mechanistically what an MLP may be doing when it performs a task in-context. Mechanistic interpretability falls outside the scope of this particular study, but we are currently working on a mechanistic follow-up and will consider incorporating your proposed methods. Thanks!

Thanks for the comments! We’re glad you found our results compelling.

The analysis is mostly in stylized and restricted settings. It is not entirely clear what this means for kinds of ICL that is observed in realistic settings. Even within simplistic settings, some more complex problems can be considered

Great point. We were interested in cleanly evaluating ICL performance, and sought to avoid confounds introduced in naturalistic settings like developing fluency in natural language or learning complex visual features (which are certainly important topics of study in their own right, but outside the scope of our investigation). Hence, we adopted synthetic tasks common in the literature. The tasks you reference are indeed quite interesting, and there are an ever-growing collection of fantastic tasks to try. In our case, we found that regression, classification, match-to-sample, sphere oddball, line oddball, and same-different were a sufficiently diverse and complex collection to make our point. In subsequent follow-ups, we will be delighted to examine other settings further.

Some useful description of the experimental setup were missing or in the appendix. There is some theoretical discussion in the appendix that are not referred to in the main paper.

Great catch. We apologize for the gaps, and will correct them in our manuscript.

The paper mostly shows empirical evidence that MLPs can be competitive, however there is not much discussion about why this might be

Great point. We comment further below.

What is the distribution of inputs $x_i$ for Section 2.1? What could happen if the tasks were even more diverse by changing the data covariance?

Inputs $x$ are sampled iid from a standard Gaussian. Changing the data covariance sounds like a fascinating direction for future study, thanks for the suggestion!

Any understanding of why Transformers fair poorly in the relational tasks and why MLPs might be better? These seem right up the alley for Transformers

Great question. To clarify, the Transformer does succeed beautifully at match-to-sample, as well as the relational tasks generally. However, comparatively smaller MLPs succeed beautifully as well, and (surprisingly) generalize better out-of-distribution. This study was intended to be primarily empirical, documenting surprising performance gaps. We are currently preparing a separate, related follow-up with theoretical arguments.

Are there previous papers that argue attention is required for ICL?

There are many. Akyurek et al, Von Oswald et al and Zhang et al are a few. As Von Oswald and colleagues put it, “Transformers show remarkable in-context learning behavior. Mechanisms based on attention, associative memory and copying by induction heads are currently the leading explanations for this remarkable feature of learning.” Many theories about ICL explicitly ground their constructions in attention-based architectures. One of our core motives is to encourage the field to move past this restriction, and consider ICL as a broader phenomenon.

These experiments mix the role of expressivity (how large the model needs to be) and optimization (can standard algorithms learn a good solution)

That’s correct, and this is a deliberate design choice. We measure models based on floating point operations, sweeping over a wide range of sizes and training time. By budgeting based on FLOPs, a model cannot have both unbounded parameter count and arbitrarily large training time, forcing an optimal tradeoff. Doing so allows us to compare diverse architectures fairly and reflects the real-world cost of training these models.

Any reason to look at linear regression/classification and not other ICL tasks (like in Garg et al.)?

Great question. The additional tasks in Garg et al. are certainly interesting, as are the ever-growing collection in the literature, and regrettably we must select a finite subset to try. In our case, we found that regression, classification, match-to-sample, sphere oddball, line oddball, and same-different were sufficiently diverse and complex to make our point. Additional perturbations like varying data diversity, increasing context length (Fig. 1) or the various out-of-distribution settings (Fig. 2) probe whether our models have learned a shortcut that leverages the simplicity of the tasks, or are truly performing in-context. We have found the latter to be true, and look forward to confirming this finding with additional tasks in the future.

One strength of Transformers is the ability to handle different context lengths n. Is this property also true for MLPs?

Yes, this is also true for MLPs. We consider variable-length inputs in Fig. 5, where the models are prompted auto-regressively. For the MLPs, inputs are zero-padded up to a maximum context length. We show that the exact same results continue to hold in this setting.