Pre-Memorization Train Accuracy Reliably Predicts Generalization in LLM Reasoning

We thank the reviewer for the feedback! In our rebuttal, we provide more detailed explanations of the motivation of our work, potential reasons behind our observations, and new results comparing our metric to sample accuracy. We hope that with these clarifications, we are able to better convey the value of our findings.

Overall, the motivation behind the paper is not clearly articulated.

The main motivation of our work is to better understand the train/test performance gap in LLM reasoning tasks, in order to design more principled strategies for improving generalization. When finetuning LLMs, models frequently achieve perfect accuracy by the end of training, yet they often achieve vastly different test performances. This suggests that, beyond training optimization (the extent to which training loss can be reduced), the generalization behavior of a finetuned model is governed by the training dynamics (the manner in which the model changes through training). Thus, to more strategically improve generalization, we aim to first better understand how a model’s training dynamics shapes generalization.

Our findings with pre-memorization train accuracy (a property of a model’s learning dynamics) allows us to bridge a model’s train and test behaviors, which can enable us to make more targeted training improvements. Our experiments with data curation shows that collecting data based on pre-memorization train accuracy can indeed lead to significant improvements to sample efficiency.

We have updated our paper to more clearly articulate this motivation.

In my view, this metric is almost equivalent to the difficulty level of the problems. When the model samples multiple times and fails to obtain correct answers, it indicates that the problem is difficult.

Indeed our metric can be viewed as a model-specific notion of difficulty. However, measuring the model’s accuracy alone does not provide a reliable metric of difficulty that correlates with downstream generalization performance. This is because models can produce correct answers to “hard” training examples by simply memorizing the target trace, without developing generalizable reasoning abilities. In contrast, our pre-memorization train accuracy metric is able to discriminate between purely memorized examples vs ones for which the model learns in a more robust manner.

To further illustrate this point, let us consider using model accuracy to predict example difficulty at different points throughout training. At the end of training, models are often able to achieve nearly perfect accuracy on the training set, which does not allow us to discriminate between the difficulties of different examples. We can also use a model checkpoint in the middle of training to evaluate sample accuracy. To evaluate its effectiveness at predicting generalization, we compare the test accuracies of models at the end of training to the train accuracies after 1 epoch of training (https://imgur.com/a/UqQbzJe). We can see that approach achieves a R^2 around 0.5, whereas pre-memorization train accuracy has a R^2 of around or above 0.9. The accuracy at the first epoch is a worse predictor, because models both memorize and robustly learn examples at different points of training, so evaluating the accuracy at any particular epoch does not produce a reliable predictor of generalization/difficulty.

We thank the reviewer for the feedback and questions! In our rebuttal, we provide an explanation of how we do the selection of p, and additional experiments demonstrating the robustness of our observations to p, as well as the generalizability of p to heldout examples and training runs. Please let us know if you have remaining questions or concerns.

There are not enough details on how p is chosen.

We find the threshold p by sweeping across a range of values, calculating the pre-memorization train accuracy across different training runs, and selecting the value which yields the strongest predictor of test accuracy. We have added this explanation, as well as several experiments measuring the sensitivity of our observations to p, to Appendix A.1 in our submission.

What is the impact of varying p on the metric value and correlation.

We conducted a sensitivity analysis of R^2 to values of p (https://imgur.com/a/yGQrQ9C , Fig 7 in our paper), and found that the R^2 value exhibits a smooth degradation to the value of p, which makes it relatively easy to find a good value of p by sweeping a range of values.

Finding such p, we have created a list of examples which correlate high to the test set. This is a form of leakage…. Show that we can calibrate p on a hold-out set of the training-set and the experiments still hold.

We present a new experiment splitting the test set of the original dataset into halves, calibrating p on one half of the test set (calibration test set), and evaluating the prediction results on the other half (holdout test set) (https://imgur.com/a/8X9WNEm , Fig. 9 in our submission). We can see that the value of p generalizes very well, achieving a R^2 of 0.928 on the holdout test set (only 0.01 below the calibration test set). Thus, this experiment shows that the high predictive power of pre-memorization train accuracy does not stem from any sort of data leakage from calibrating p.

To further illustrate the robustness of p, we also experimented with splitting our training runs into calibration training runs and heldout training runs. We calibrate p only on the calibration training runs, and evaluate on the heldout training runs (https://imgur.com/a/HBDtzx3 , Fig. 8 in our paper). We can see that using only 1-3 training runs for calibration, we are able to arrive at a value of p which achieves a R^2 of >0.86 on heldout training runs. These experiments show that the high predictive power of pre-memorization train accuracy does not stem from “overfitting” p to specific training runs during calibration.

Thank you for your feedback! To address the concerns, we have provided an explanation of the selection of p, experiments showing the sensitivity of p, a comparison with train accuracy after 1 epoch, and a reproduction of Fig. 5 using a model trained with 3 epochs. Please let us know if this addresses the issues, or if there are further changes or clarifications that you would like.

Many details are not adequately presented, such as the selection for the threshold p in pre-memorization train accuracy

We find the threshold p by sweeping across a range of values, calculating the pre-memorization train accuracy across different training runs, and selecting the value which yields the strongest predictor of test accuracy. We have added this explanation, as well as several experiments measuring the sensitivity of our observations to p, to Appendix A.1 in our submission.

Is the correlation coefficient for predicting test set accuracy sensitive to this parameter?

We conducted a sensitivity analysis of R^2 to values of p (https://imgur.com/a/yGQrQ9C , Fig 7 in our paper). We can see that the R^2 value exhibits a smooth degradation to the value of p, which makes it relatively easy to find a good value of p by sweeping a range of values.

The authors claim that this training set accuracy metric based on the training process is superior, but they lack necessary ablation studies.

To the best of our knowledge, there are no prior works which have proposed generalization metrics with the exact same assumptions as ours. Instead, we compared our metric against a range of existing approaches which makes use of slightly different inputs (gradient variance, distance from initialization, and ATC). We have added a comparison to the accuracy at the first epoch (see below). Please let us know if there are any other ablations that you would like to see.

How would the prediction performance change if only the training set accuracy after the first epoch(rather than first m epochs) is used?

We find the average train accuracy after the first epoch to predict the final test accuracy, with R^2 around 0.5 (https://imgur.com/a/UqQbzJe ). This is because models often do not learn a subset of examples until later epochs. In contrast, pre-memorization train accuracy has a R^2 of around or above 0.9, making it a more reliable predictor of generalization.

In the experiments predicting test set accuracy, the comparisons with other methods are not entirely fair. The gradient variance and distance from initialization methods rely on significantly less input information during prediction, while the ATC method primarily focuses on OOD scenarios.

We acknowledge that the comparisons with prior methods make use of different input information, though we do not believe they use significantly less information. While our approach makes use of model outputs (accuracy / perplexity) and requires calibration for the hyperparameter p, prior approaches like gradient variance and distance from initialization make use of information about the model weights, which ours does not. We also acknowledge that ATC was primarily designed for OOD settings, though we adapt their insights to the IID setting. We used these baselines because they are, to the best of our knowledge, the most well known/effective existing approaches for predicting generalization in prior works. We have added a note in our paper highlighting the differences in assumptions between our approach and these prior approaches.

We thank the reviewers for the response. To address these concerns, we present new experiments showing that the value of p generalizes to held out test sets and training runs, as well as a comparison to the first epoch training accuracy with p.

It is quite weird that p is chosen by fitting R^2 with a series of experiments and then using the fitting performance as an evaluation metric.

We present new experiments showing that p does not need to be fitted on the dataset on which we ultimately evaluate on. The value of p generalizes quite well to heldout test datasets, as well as heldout training runs.

First, we split the test set of the original dataset into halves, calibrating p on one half of the test set (calibration test set), and evaluating the prediction results on the other half (holdout test set) (https://imgur.com/a/8X9WNEm , Fig. 9 in our submission). We can see that the value of p generalizes extremely well, achieving a R^2 of 0.928 between pre-memorization train acc and test acc evaluated on the holdout test set. Thus, this experiment shows that the high predictive power of pre-memorization train accuracy does not stem from any sort of data leakage from calibrating p.

We also experimented with splitting our training runs into calibration training runs and heldout training runs. We calibrate p only on the calibration training runs, and evaluate on the heldout training runs (https://imgur.com/a/HBDtzx3 , Fig. 8 in our paper). We can see that using only 1-3 training runs for calibration, we are able to arrive at a value of p which achieves a R^2 of >0.86 on heldout training runs. These experiments show that the high predictive power of pre-memorization train accuracy does not stem from “overfitting” p to specific training runs during calibration.

p needs to be reselected under different data distributions. In practical applications, training data often requires multiple iterations. How should this metric guide experiments?

In Section 5.2 of our work, we discuss how pre-memorization train accuracy can be used to guide data curations. Our experiments show that our approach leads to a 1.5-2x improvement in sample efficiency compared to IID data curation, and outperforms other standard data curation approaches.

In these experiments, the data distribution changes as we scale up the amount of training data. However, we only calibrated the value of p once, using the smallest version of the dataset. If the task does not shift significantly, we find the value of p to generalize relatively robustly.

In the curve of fitting performance changes with p provided by the author, R^2 sometimes reaches quite low values (<0.4), making it difficult for me to judge when using average train accuracy only after the first epoch to predict the final test accuracy (R^2=0.5). Could the author also plot a similar curve of fitting performance changes with p only after the first epoch to illustrate the importance of using multiple epochs?

We experimented with using the metric $\text{Acc}(f_{\theta 1}(y|x_i), y_i) * 1[\text{Perp}(f_{\theta 1}(y|x_i), y_i) > p]$ to predict test accuracy, where $f_{\theta 1}$ refers to the model after 1 epoch of finetuning (https://imgur.com/a/Yfm0vUe). Sweeping across different values of p, we find that the R^2 values does not surpass 0.51. The most optimal values of p is when $1[\text{Perp}(f_{\theta 1}(y|x_i), y_i) > p] = 1$ for all examples. Thus, pre-memorization train accuracy, with R^2 >= 0.9, serves as a much better predictor of generalization.

For difficulty levels 4-5 in MATH: Given that elementary and middle school level math problems have almost been resolved, if the method cannot be proven to extend to more difficult problems, it will result in a loss of soundness.

The reason we do not perform our data curation experiments with levels 4-5 in MATH is because they are too difficult for the size of models that we use (8-9 billion parameters). This does not mean that our findings do not generalize to difficult problems in general. If we were able to train with bigger, more capable pretrained models, then they would be able to achieve higher performance on levels 4-5, which would provide better signal for our comparisons.

For academic papers in the field of LLMs, it is standard practice for researchers to experiment on smaller models with simpler tasks (where there is more signal), due to computational constraints. It is generally accepted that experimental results from 8-9 billion parameter models are likely transfer to bigger models on harder tasks.

We thank the reviewer for the nice feedback! We answer the question below.

It seems that for small models, they may never end up memorizing the reasoning trace, and the notion of a pre-memorization accuracy does not exist. Are there nevertheless proxy metrics available in those cases? At what scale does the generalization emerge?

We find pre-memorization train accuracy to be predictive of test accuracy, regardless of whether the model memorizes the training data. One example is models trained with a learning rate of 5e-7, which, as illustrated in Fig. 2, does not memorize any training data by the end of training (no yellow in last row). In Fig. 3, we can see that the pre-memorization train accuracy of models trained with 5e-7 learning rate are still highly predictive of test accuracy. For situations where the model does not memorize the training data, the pre-memorization training accuracy is generally equivalent to the training accuracy.

It would be interesting to consider a more diverse set of active learning metrics. For example, can influence functions be used to map examples to train data (as proposed in [1])?

We believe that influence function based approaches for data curation would likely be effective, but we did not implement this approach due to its computation and algorithmic complexity, such as the need to compute the inverse-Hessian-vector product for each query. In contrast, our approach only makes use of the model’s outputs, which allows us to treat the model as a “black box”. We believe simplicity is one of the advantages of our approach, which enables practitioners to quickly implement in their own training workflows.