Bayesian scaling laws for in-context learning

Thanks for your review! We try to address your questions below:

1. Thank you for this feedback on the textual exposition. Eq. (3) adds the ICL efficiency coefficient ($K$) and weight-tying to reduce latent parameter count, as explained in section 3.2. We agree that our explanation of why we did this could be better. We will update the text to explain both, based on the explanations below:

   • Given some ICL curves for several tasks on a single model, we can easily deduce the diagonal values in $P$ ($P_{x, x}$) since those are the values the ICL curve for each task asymptotically converges to. However, the non-diagonal values are latent and unobserved. Having significantly more latent variables than observed ones for such a small model makes the model liable to overfitting (since there is nothing grounding the latent vars), so we come up with weight-tying schemes to reduce the number of latent vars to fit. The weight tying also makes the comparison with power laws and other baselines fairer since the number of parameters is much reduced.
   • The efficiency coefficient is explained in section 3.2 to allow for the possibility of multiple and/or noisy updates. Figure 3b shows the empirical utility of K --- it neatly captures the informativeness of the ICL examples in the GINC setting, with longer examples giving higher K. Another interesting experiment could be to compare K on the same task between a base and instruct model.

2. Our Appendix C.1 might not be entirely clear here; for all laws, we take the negative log of the whole expression, so we are computing negative-log likelihood in all cases. For the non-Bayesian laws, we can use logarithm rules to simplify the expression and make it more numerically stable, which is why we have ln n in those expressions; we aren't taking the log any differently. To make this clearer, we can provide snippets from our code in that section for the rebuttal version, and we will of course release our code with the final version of the paper.

   • We empirically tested the clamping and found no significant difference when clamping in (-20, 20) vs. (-10, 10), and picked the looser range. For our tasks, it seems the optimal values are well within the clamped range.
   • Also note that none of the clamping and log-space transformation is necessary when using Adam as an optimiser, but Adam performs worse across the board than L-BFGS; L-BFGS is just more sensitive to numerical instability.

3. $K$ is a learned parameter and so not relevant for this analysis. We do not expect error to blow up with $n$ because the Bayesian laws always converge to some value as $n \to \infty$ and the output probability is constrained to [0, 1]. As for $M$, our weight-tying scheme, which reduces the number of latent variables from O(M^2) to O(M), significantly combats error by reducing the possibility of overfitting. We can mention some of these points in the text but we believe we have already addressed these, and cannot find analogous analyses in prior ICL scaling law work such as Anil et al. (2024).

We really appreciate your comprehensive review! It raises a lot of important questions that we will try our very best to address below. First though, we do want to discuss the framing of our results. We absolutely agree that a difference in NRMSE of 0.02 is negligible (even if statistically significant) and we are not in the business of hill-climbing on error rates anyways.

The point of this work is to take inspiration from all the papers claiming ICL is Bayesian in transformers using purely theoretical proofs, and see whether it holds up in real life and is practically useful. At least on paper, a Bayesian view of post-training explains failure modes like many-shot jailbreaking (if only task priors are adjusted and not task knowledge). The quantitative results are to show that Bayesian approach is competitive and to simply justify its use as an analysis tool for ICL behaviour, and are not the end goal of the paper. We agree that the qualitative findings of how DPO/SFT work are the actual phenomenon we want to understand. We can be clearer about this in the introduction and discussion/conclusion.

Now, to address the weaknesses:

1. We are happy to relax the claim in the abstract to "our scaling laws match existing scaling laws" since the margin is so negligible.

   • In table 1, we felt that every variant was justified to include: Bayesian (original) is the exact functional form we derived, while the other two use different weight-tying schemes which seem equally justified given that the original form overparameterised and has too many latent variables to fit (off-diagonals in $P$). Figure 2b additionally shows that all methods are almost the same (and very correlated) on NRMSE, and is additional evidence for our claim. In our discussion (section 6), we are very careful to note that we haven't proven Bayesian methods are better. We can be clearer about all of this in our intro; we just want to point out we are already trying to be very careful not to overclaim.
   • The insignificance in the DPO comparison likely comes from the very very large NRMSE numbers (>1e6 in some cases!) when the DPO training set size is large and the ICL curve collapses. We don't think the insignificance will change given how noisy the fits are in that setting; figure 5a shows how severe the ICL curve collapse is for all the methods.

2. This is a great point. We have many examples of such curves and can include them all in the appendix for the various model depths and SFT/DPO settings we tested.

3. Yes, what we want to highlight is that DPO suppressed the negative subdistribution to 0 unlike SFT (such that ICL no longer works on that distribution in some cases), and that DPO training data requirements are less sensitive to model size than SFT. We can reword this sentence to be more clear -- we don't mean it uses less training data in general, just that it needs roughly the same amount for both an 8-layer and 3-layer model in our experiments.

4. Did you mean that the base model ICL curves look unusual? The instruct models seems to have Bayesian-seeming curves here (we can include the learned fit from one of the laws in this plot). We believe the unusual ICL curve for the base model stems from the base models not being trained on the instruct chat template we used, which has some special tokens that were probably not seen in base-model training ([INST] [/INST] etc.). We can report additional ICL curves for Llama 8B with a simpler input template format (e.g. "Question: ... Answer: ...") in an appendix, which we believe should show a more monotonic ICL curve.

   • Also note that figure 6b is the posterior of the Bayesian scoring-wise fit, it's not the actual computed ICL curve. The task posterior seems to be the right way to compare different models since it takes out the noisiness of the conditional probabilities for the task, which is why we show it here. We can plot the actual ICL curve on figure 6a as well to help see this.
   • In general, we believe that our findings provide additional evidence for the claim that LLMs do ICL in a Bayesian way. We are one of the few works that empirically try to test this (as opposed to just theoretically prove it) and we believe this kind of an experiment is necessary in the literature to really assess the claim. If someone finds a non-Bayesian functional form that performs significantly better than ours, that would be empirical evidence against the claim, but so far none of our baselines succeed in this (including the logistic baseline, which is newly introduced by us).

[continued in the next comment]

[continued from previous comment]

5. We can be clearer about this in section 3.2. Essentially, given some ICL curves for several tasks on a single model, we can easily deduce the diagonal values in $P$ ($P_{x, x}$) since those are the values the ICL curve for each task asymptotically converges to. The non-diagonal values are latent and unobserved. Having significantly more latent variables than observed ones for such a small model makes the model liable to overfitting (since there is nothing grounding the latent vars), so we come up with weight-tying schemes to reduce the number of latent vars to fit. The weight tying also makes the comparison with power laws and other baselines fairer since the number of parameters is much reduced.

   • The efficiency coefficient is explained in section 3.2 to allow for the possibility of multiple and/or noisy updates. Figure 3b shows the empirical utility of K --- it neatly captures the informativeness of the ICL examples in the GINC setting, with longer examples giving higher K. Another interesting experiment could be to compare K on the same task between a base and instruct model.

6. This is a very good point. For the current experiments, the models are gradient-step matched, not FLOPs-matched. They trained to convergence given the data amount we used (we can add loss curves to the appendix to show this). The FLOPs-matched setting does seem interesting though; since the models are quite small, we are happy to run some FLOPs-matched experiments and see if the loss and ICL behaviour is substantially different; if they are, we could replicate our SFT/DPO experiments on those. We would keep these in an appendix.

   • For SFT/DPO we believe gradient-step matching is more realistic since the bottleneck in these settings is usually data amount, not compute; a practitioner could ask, given some fixed dataset, what scale model should I finetune to avoid many-shot jailbreaking?

As for the additional questions:

1. You are correct, that should be $\rho$. We will carefully go over the math in the text again to check our notation.

2. This is a typo, as above (it should be included and we will fix this).

3. Yes. We will standardise this to "task" to be clear.

4. For ICL, the practically useful variable is the probability of the correct answer. Additionally, the Bayesian derivation directly operates on probabilities, and previous work (such as Anil. et al, 2024) also came up with functional forms to fit to probability or log probability. We did consider KL divergence, but unlike e.g. for the language-modelling objective, it isn't the main metric of interest when we are studying ICL. KL divergence is also unbounded and thus less numerically stable to fit than the probability (which is in [0, 1]). If we simply modify our functional form and the baselines to express KL divergence from ground truth, we don't expect to see any difference, but keeping the fits stable could be much harder.

5. Yes, thanks for catching this, we will specify that.

6. Oh this is a very good point that we had not considered. We can add texture to the boxes to make this clearer for color-blind readers.

7. We agree! We can supply many of these for the GINC experiments in the Appendix, and add them to figure 6a as well.

Thanks again for this very insightful and constructive review!

Thank you all for the reply! In general, the updated manuscript will help a lot!

1. The updated manuscript will help, thank you!

2. Same as 1.

3. I see, this explanation makes sense. I think it would be better if the following is highlighted in the updated manuscript: "DPO suppresses the negative subdistribution to 0" and some quantification of how close to zero it is will help! Irrespective of the author's statement that "DPO training data requirements are less sensitive to model size than SFT." However, does the claim "Unlike SFT, DPO can use very little training data to successfully suppress the disfavoured HMMs beyond the ability of ICL to recover" in l. 418,419, stand?

4.1 Yes the manuscript update showing the goodness of fit will help me understand this better, currently its hard to see how good the fit is. I am also a bit confused why the authors decided to (or had to) use special tokens to evaluate the base model.

4.2 I totally get the argument: "We are one of the few works that empirically try to test this (as opposed to just theoretically prove it) and we believe this kind of an experiment is necessary in the literature to really assess the claim." and agree that this empirical validation is exactly what we need beyond theoretical possibility claims. This by itself is an important contribution of the paper. My main concern is whether Table 2. actually supports "our findings provide additional evidence for the claim that LLMs do ICL in a Bayesian way.", as it seems like the logistic baseline does at least as good? Shouldn't the claim be "our findings suggest that Bayesian assumptions are as consistent as conventional scaling laws for ICL in LLMs do ICL"? Please tell me if I am misunderstanding something!

5. Sorry that the question was vague. I'm wondering what justifies this specific scheme of weight tying. At the same time, I can understand why this is the most reasonable thing to do. If you can motivate this better it would be great, but I don't think this is a weakness!

5.1 I'm a bit afraid if K will absorb any non-Bayesianness and thus make this scaling law fit a non-Bayesian ICL process too well.

6. Yes! Even showing a single FLOPs matched run (without re-running everything) would be great. But also, perhaps you can just take an earlier checkpoint of your bigger model so that you don't need to re-run? But if the authors are pretty sure the loss is saturated, I think this is uneeded.

6.1 Yes, totally, the question was only about the pre-trained models.

Summary In general, I think I am willing to raise my score if I can see some of these: (more important on top)

• The updated manuscript addressing 1,2,3,4
• Clarification of why Table 2 supports: "our findings provide additional evidence for the claim that LLMs do ICL in a Bayesian way."
• Explanation on the claim: "Unlike SFT, DPO can use very little training data to successfully suppress the disfavoured HMMs beyond the ability of ICL to recover"
• A single FLOPs matched experiment (if the loss is not saturated)
• Does the presence of the K parameter affect if we can draw conclusions about Bayesianness from these scaling laws?

Thank you!!!

My concern from these interpretations is as follows:

In general, I think conclusions about "model scaling" should be drawn in 1) equal flops or compute optimal setting 2) in a regime where scaling is smooth in loss without qualitative changes. This is because without these conditions one can easily confound the effect of scaling the model and just training more. e.g. see Figure 2 in Kirsch et al [3] where the model size is scaled keeping the # of steps equal and there seem to be a qualitative change.

The main text in general and Figure 2b, 3b, 4b, 5b seem to suggest the plots show an effect of model scaling in the sense of LLM scaling laws (Kaplan et al, Hoffmann et al), where we have a smooth and predictable scaling curve of loss without qualitative changes depending on architecture. (In fact, Kaplan et al further strengthens this latter argument by showing that "adjusting the width or architectures while holding the # of params. constant don't have a significant effect".) However, interpretation 1 and 2 suggests that in these experiments 1) models are not trained at the compute optimal frontier (or even at the equal FLOPs threshold) 2) there is a qualitative change depending on the number of layers. Thus, drawing conclusions about "model scaling" from Figure 2b, 3b, 4b, 5b seems like its confounding scaling and more training and qualitative changes. The authors suggest that the 4~12 layer models are saturated. If all conclusions are drawn from saturated models without qualitative changes, this will indeed refute the concern. However the scaling findings in Figure 2b, 3b, 4b, 5b, seem to heavily rely on the 1,2,4 layer models. Perhaps the only pure scaling result seen here is the change between the 4-layer and 8-layer model.

I think there are two options (but feel free to suggest something else):

• 1. Clarification in the manuscript that these are equal steps: I think this is the cleanest way forward. There is nothing wrong about studying model scaling with equal steps, as long as this is clear. The problem is that sometimes this causes the audience to interpret results as "bigger models are fundamentally different" even when the difference can be explained by them having received simply more training. Thus I think the paper should clearly recognize that bigger models just train faster per steps, but equal number of steps are given in this work, and clearly show the loss curve.

• 2. Reproducing results with equal loss or equal FLOPs. I think 1 is much easier, but I'm adding this mostly to clarify what I think a "model scaling" result should look like. Models which do not show qualitative changes (4layers to 12layers) should be assembled and checkpoints should be collected at equal FLOPs or compute frontier (which will require a pre-training scaling law experiment, and thus is probably out of scope). Then the results should be drawn from these checkpoints. Alternatively, one can study different checkpoints of a single model and show that the effect of scaling compute and model scale are qualitatively different, refuting the hypothesis that "model scaling effects can be simply explained by more compute given to bigger models".

I really like the idea of the paper, and again, I am ready to increase the score as soon as W6 is addressed. The fact that we are even able to have this discussion about "are bigger models more Bayesian?" is only possible because of the framework in this paper. But I really think this issue about model scaling vs simply more FLOPs is given should be clarified.

Please tell me if any argument I made is unclear, I am happy to explain it further!

Furthermore, given that the manuscript cannot be updated, I will trust any result reported from the authors without plot evidence, especially since the authors has already demonstrated their honesty by showing the loss curves very clearly.

[1] Kaplan et al [2] Hoffmann et al [3]https://arxiv.org/pdf/2212.04458

I think most concerns(1~5) are very well addressed. However, I am unsure how to handle the weakness 6. It would be great if we can have a discussion on this, I am really ready to raise the score as soon as W6 is addressed!

1. Thank you!
2. My main concern was that the scaling law fit wasn't visible like in Fig 6. Quoting my original comment: "The scaling curves and the fits are not shown as plots for the GINC experiments". The DPO's ICL collapse is indeed interesting!
3. Thank you this really makes it much easier to understand with respect to the plot.
4. Thank you, the discussion is especially very clear.
5. Thank you!

6. My interpretation of these loss curves is as follows:

• Interpretation 1: FLOPs is indeed a confounding factor: It seems like it is precisely the case that FLOPs is being a confounding factor. In fact, if you can see that the loss of the 1layer model at epoch 4 is almost precisely the loss of the 4layer model at epoch 1. Same for the 4layer model at epoch 4 and the 8 layer model at epoch 2. Of course, the simple linear relation I'm pointing out wouldn't hold either robustly across scales because of 1) early training dynamics and 2) scaling laws [1,2] themselves don't predict equal flops=equal loss.
• Interpretation 2: Going from 1-layer to 4-layer gives a qualitative change: I see that there is a qualitative change between 1,2 layer models and 4~12 layer models, as one can see from the crossover of the 1,2 layer models losses.
• Interpretation 3: Saturation in 8~12 layer models: I agree to the authors that model 8~12 are likely saturated, maybe the 4-layer model is as well.

[Continued below...]

Thank you for the review! We appreciate that you took the time to go through the derivation in the appendix and found your review very helpful.

In (2) we mean to say that the symbols in document $D$ are sampled from $\lambda$; we will change the notation to say $D \sim \lambda$ and clarify this.

We agree that the current version has an error in (17) since we E[A/B] only = E[A]/E[B] for particular distributions of A and B (cf. https://link.springer.com/article/10.1007/BF02927114), which we have not proven (or claimed) for our A and B. We have fixed this in the rebuttal response paper; we also outline here how we did so. Our fix requires no rerunning of experiments and is purely formal.

In our theorem, we want to model the effect of repeated Bayesian updates when each piece of evidence is sampled from $\lambda$. When fitting the functional form, we set $\lambda$ to one of the subdistributions of interest, but that makes the expectation $\mathbb{E}_{\sigma \sim \lambda}[A/B]$ unsimplifiable (without further assumptions about the distribution of A and B).

Instead of being stuck here, we can assume that $\lambda$ is one-hot, so that each ICL example is identical and thus interchangeable. Since $\lambda$ now collapses to having a single outcome (repeated ICL symbols A,A,A,...), we can simply get rid of the expectation. The idea here is that we have noise in the ICL process modelled entirely by $\delta$ (the probability of an ICL symbol conditional on a task) and not by the document sampled from $\lambda$. Intuitively, this simplification means that we are now finding the average Bayesian update needed to get the ICL curve we are trying to fit.

There are some possible next steps that we have left out of this work: one could introduce some assumptions about the values in $\delta$ to somehow make (17) simplifiable without collapsing the expectation to one outcome. We are aware of some unpublished work trying to do this for inference-time scaling laws. We believe that for the purposes of our work, the simplification of $\lambda$ is sufficient, and leave further modification to future work.

We would also like to point out that previous work on ICL curve fitting has not used this kind of theoretical justification to obtain their functional forms (see our discussion in appendix B), and so we are facing a novel problem when trying to derive a functional form from theoretical claims that ICL is Bayesian---we believe that doing so is the only way to really test that claim empirically.

Thank you for the positive review! To address your points:

1. It would be very interesting to see if these results hold for alternative architectures. Previous theoretical work on the learning algorithm underlying ICL, particularly those claiming that it is Bayesian, has largely focused on transformers (see section 2 of our paper) and so we prioritised studying that architecture. Future empirical work on ICL in other architectures would be great once we have some theoretical proofs to start with.
2. Yes, in our experiments we did not study the effect of model width; all model hyperparameters were held constant besides depth. We don’t expect width to have much of an effect since (1) the synthetic task we train on is very simple and is learned well by our small models (when depth is at least 3), and (2) ICL has been tied to model depth in previous work via induction heads (e.g. https://transformer-circuits.pub/2022/in-context-learning-and-induction-heads/index.html) which require model depth of at least 2 to develop.