A Simple Model of Inference Scaling Laws

Dear Reviewer i2X4,

Thank you for your careful and positive review of our paper, finding that our proposed model is simple and elegant. We hope that our responses will alleviate your concerns.

Below, we address the issues you raised.

Comments and Suggestions

1. Thank you for pointing this out, we will complete L87: For $n$ training samples, the expected single feature error $E_i$ is.
2. Thank you for finding this typo, it will be fixed in the revised version.

Weaknesses and Questions

1. Data generation process - While it is true that the data generation process of LLMs may seem very different from the simple memorizing model, the claim of the paper is not that the simple model perfectly captures every detail of the LLM inference process, but rather the opposite – that these details are not crucial to address inference scaling, as it is often true with pre-training scaling laws [1,2]. The only requirements from the LLM in order to fit our assumptions are that the model is sufficiently large, or has sufficient capacity to (approximately) store its training data, and that inference is done by an imperfect sampling process. The specific task we chose – which is recovering exactly the training data itself, is not what LLMs do, but it is a proxy for the task of “recovering the right answer from the training data”. Meaning that if the answer, or the path to the answer (a sequence of retrievals from its memory) existed in the training data, this model is able, with some error probability, to reach this answer. This approach explains why this simple model captures very different scenarios (LLM and VAE). In order to understand the precise connection between our model and LLM performance, it would be wise to study the internal representations of LLMs, and look for “memorizing modules”, possibly in specialized attention heads, as was suggested by reviewer kFgD, and will be part of the focus of our future works. We will better explain this point in the main text of the revised version, if the reviewer finds this answer acceptable and useful in understanding our goal.

2. Other functional fits to the data - We agree with the reviewer that ablation with simple functions could be useful, we will include a short appendix considering several functions of the suggested type to show that they cannot be used to capture the full behavior of the model, typically sigmoid type functions are somewhat able to capture the small $k$ regime, but never both limits. We provide a preliminary figure in https://imgur.com/kTsDM0R .

3. Similar parameters for different descriptions - The parameters should not be the same, but there exists a simple mapping between them that can be derived from the large and small $k$ limits.

4. Beta fit to practical data - Thank you for this comment. The two figures are conceptually different. Regarding Fig 4(c), we will include the empirical curves, which match very well the asymptotic behavior predicted by our model (https://imgur.com/a/8Ftb76R). Regarding Fig 1(b), this is a prediction made by our model that requires further study to properly interpret. In essence, this figure shows the "difficulty" distribution for the different models, on the specific set of maths questions. It does not fit any empirical data, but should be taken as the interpretation of the $\alpha,\beta$ parameters. The goal of the figure is to show that different models can perceive the same data at different difficulty levels, and a detailed study of this point could lead to new inference sampling techniques, based on the ratios of easy and difficult questions. We hope this explanation clarifies this point and would be happy to discuss further.

5. Figures 6(b) and 7(a) fitting the theoretical curve - Figs 6(b) and 7(a) are meant to serve as qualitative evidence for the effective correlation between trials, rather than an exact characterization. A priori, the inference attempts themselves need not be correlated, but we see that there are effective long range correlations between different trials from both the reconstruction error 6(b) and its eigenvalues 7(a). We see that it is sufficient that only a bulk of the eigenvalues conform to a power-law decay to have good predictions for the pass@k metric. We will explain this point more clearly in the revised text.

6. emergent behavior at specific threshold value - The logic is that for larger threshold values the model must have smaller reconstruction error, and so require fewer inference attempts to succeed. The "emergence" is just a feature of the functional form of Eq. 7. It would be interesting to analyze this model in more detail and consider whether one can define "emergence" as the point where the function changes from concave to convex perhaps, and study the internal structure of the model at these points.

References:

[1] Maloney et al. https://arxiv.org/abs/2210.16859

[2] Bahri et al. https://www.pnas.org/doi/abs/10.1073/pnas.2311878121

Dear Reviewer FFcT,

Thank you for your positive appraisal of our submission. We’re glad that you found valuable insights and potential practical applications, which were our goals.

We address the weaknesses raised, as well as questions below.

Weaknesses

1. strong assumptions - We completely agree that both the independence assumption and the perfect verification limit are strong assumptions. However, given that our paper is the first to propose a model for inference scaling (as far as we know), it is natural to begin with a limiting, ideal setting and extend it in future works. It is rather surprising that even under these strong assumptions, the simple memorizing model captures the inference scaling behavior of real models as well as it does.

We would like to kindly bring to the reviewer’s attention the fact that the independence assumption itself is also discussed in the main text, and the “effectively correlated trials” section (Sec. 4.2) is meant to describe this effect, if we understand the reviewer correctly. We apologize if this point was not sufficiently clear, and will try to highlight it further in the revised version. Furthermore, we are sure that relaxing the perfect verification assumption will be a very interesting avenue for future works, since verification methods is a very broad topic all in itself.

2. More diverse models - We agree in principle that adding more diverse controlled experiments could extend the scope of the paper, but we believe that the current evidence on LLMs and VAE reconstruction supports our predictions. The VAE experiments are meant as intermediate step, that shares some of the complexity of LLMs without the full pipeline. If the reviewer has a particular experiment in mind that would make sense within the “memorization-inference” setup, which is neither VAE nor the LLMs, we would be happy to test our predictions on it.

Questions

1. Scaling beyond tasks in the main text - The experiments given in the main text which consider LLMs on reasoning tasks and the VAE reconstruction are different, but share some common features. In particular, in both tasks, the model is asked to “learn” the training data distribution and perform sampling (as opposed to regression for instance). In that sense, the exact details of the architecture, model and task are not particularly important, as long as the model can be thought of as a two component “memorization” module and “inference” module. Therefore, our predictions are quite universal, as long as the data is not uniformly “easy”, and so we believe our results should extend to other generative settings, for instance to diffusion models.

2. Sensitivity to $\alpha,\beta$ - The question of sensitivity to the parameter choices is a bit unclear to us, since the conclusions (namely the inference scaling laws) are analytically given in terms of $\alpha,\beta$, so the exact dependence on the parameters can be characterized. Could the reviewer please clarify their meaning? If the question is of interpretation, then different choices of $\alpha$ correspond to a different small $k$ behavior, since as we explain in the main text, the average “difficulty” of the samples is $\frac{\alpha}{\alpha+\beta}$, and so larger $\alpha$ values would imply a slower improvement with $k$ for a smaller number of trials, while $\beta$ dictates the large $k$ behavior of the inference loss/pass@k, meaning how difficult it is to improve performance by increasing $k$ at a large number of inference attempts, where larger $\beta$ means greater scaling improvement.

We hope that our replies are sufficient to raise the reviewer’s confidence in our work, and potentially accept the paper. We welcome any further questions and comments.

Dear Reviewer kFgD,
Thank you for carefully reading our manuscript, we are glad that you found our theoretical analysis sound and our empirical results well backed.

We would like to address the various points you raised below.

Regarding Evidence

We agree that the VAE results are interesting, and are happy to see that they were appreciated. To be honest, having discussed with other researchers and due to the broader interests of the community, we preferred to highlight the functional matching to the LLM performance rather than the more controlled VAE setting. If you believe this is sufficiently interesting to the community, we will try to highlight these results a bit more in the revised manuscript.

Regarding Broader Scientific Literature

As far as we know, this is indeed the first theoretical model for inference scaling. From a brief reading of the Park et al. paper, we agree that there are interesting connections there, though their work is for in-context learning and ours was tested on inference without further context. It is not unreasonable to think that one way to interpret the in-context update to the representation as somewhat equivalent to resampling from the memory of the pre-trained model during inference. Perhaps in the ICLR setting our "error probabilities" are not fixed, but updated according to the context provided. We thank the reviewer for bringing this to our attention, and will add it to the related works. If the reviewer has more insights to share along these lines during the discussion period, possibly regarding something we've missed, we'd be happy to discuss it further.

Comments/Suggestions

This is an extremely interesting question, and we agree that developing a mechanistic understanding of the memorization module ansatz would be a valuable complement to our simple model. Our goal in this work was to introduce a first model that provides a clear explanation of the key inference observation. Understanding how this final result emerges from training would be a natural extension.

As a next step—currently a work in progress—we aim to explore the connection between pre-training scaling performance and inference scaling in solvable models, which we believe exist. This direction may align with the reviewer's intuition, as memorization must occur during training, and we may observe different inference scaling behaviors depending on whether the "memorization assumption" holds or breaks down. Moreover, the simple model presented here seems to give quite universal results, but it still might be that there are several "universality classes", where different models might fall into.

Finally, while we have not yet considered RNNs, their study should, in principle, be feasible, and would very likely lead to different scaling behaviors, at least in some cases (if not considering SSMs perhaps).

We appreciate this valuable suggestion and will certainly investigate it in future work.

We hope that our replies have affirmed your confidence in our work, and potentially lead to accepting the paper. We welcome any further questions and comments.

Dear Reviewer 4E5M,
Thank you for your positive reading of our manuscript.

Below, we address your comments/questions:

Weakness:

The caption for Figure 1 was quite dense – In the revised version, we will shorten the caption, and move some of its content to the main text, namely L131 – 133 could be moved. Do you believe this will make the caption easier to understand?

Comments/Questions:

1. The authors should include some motivational real-world examples where the memorization assumption holds (in approximation) - As per your suggestion, we will include an appendix which contains a very simple neural network classifier which can memorize (approximately) its training data, combined with an inference model which is allowed to make mistakes based on the Beta distribution. This way we separate the memorizing assumption from the inference assumption.
2. Thank you for the correction, we will fix the mistake.

Questions:

1. Regarding the Beta distribution - The Beta distribution is a particular choice made in this paper, but it is certainly not unique, as we point out in the main text: One way to model…. The reason for this choice is the fact that the inference behavior differs for small $k$ and large $k$, and so at least a two-parameter distribution is required. For instance, using the classical Zipf law type distribution with a single decay parameter $\alpha$ would only capture one of these limits. Other two parameter distributions would also be acceptable, but the interpretation will remain the same: the model perceives some inference tasks as “difficult” and others as “easier” depending on the two parameters of the distribution.

We hope that our replies are sufficient to raise the reviewer’s confidence in our work, and potentially accept the paper. We welcome any further questions and comments.