What Do Learning Dynamics Reveal About Generalization in LLM Mathematical Reasoning?
Authors: Katie Kang
Amrith Setlur
Dibya Ghosh
Jacob Steinhardt
Claire Tomlin
Sergey Levine
Aviral Kumar
Venue: ICML 2025 poster (2025)
Keywords: reasoning
generalization
memorization
Primary area: deep_learning->large_language_models

Abstract: Modern large language models (LLMs) excel at fitting finetuning data, but often struggle on unseen examples. In order to teach models genuine reasoning abilities rather than superficial pattern matching, our work aims to better understand how the learning dynamics of LLM finetuning shapes downstream generalization. Our analysis focuses on reasoning tasks, whose problem structure allows us to distinguish between memorization (the exact replication of reasoning steps from the training data) and performance (the correctness of the final solution). We find that a model's performance on test prompts can be effectively characterized by a training metric we call pre-memorization train accuracy: the accuracy of model samples on training queries before they begin to copy the exact reasoning steps from the training set. On the dataset level, this metric is able to almost perfectly predict test accuracy, achieving $R^2$ of $\geq 0.9$ across various models (Llama3 8B, Gemma2 9B), datasets (GSM8k, MATH), and training configurations. On a per-example level, this metric is also indicative of whether individual model predictions are robust to perturbations in the training query. By connecting a model's learning dynamics to test performance, pre-memorization train accuracy can inform training decisions, such as the makeup of the training data. Our experiments on data curation show that prioritizing examples with low pre-memorization accuracy leads to 1.5-2x improvements in data efficiency compared to i.i.d. data scaling and other data scaling techniques.

What Do Learning Dynamics Reveal About Generalization in LLM Mathematical Reasoning?
Katie Kang'
Amrith Setlur 2 Dibya Ghosh
Jacob Steinhardt
Claire Tomlin
Sergey Levine! Aviral Kumar
Abstract
1. Introduction
Large language models (LLMs) have demonstrated remark- able problem-solving capabilities, yet the mechanisms by which they learn and generalize remain largely opaque. For instance, consider set of LLMs, each derived from the same pretrained model and finetuned on the same reasoning dataset but with varying learning rates (Fig. 1). While sev- eral of these models reach near-perfect accuracy on training data, their test performances raises the question: what factors in an LLM's finetuning training lead to differences in its generalization behavior? Understanding these factors could help us design better training methods that foster genuine reasoning abilities in models, rather than mere pattern matching.
Modern large language models (LLMs) excel at fitting finetuning data, but often struggle on un- seen examples. In order teach models genuine reasoning abilities rather than superficial pattern matching, our work aims to better understand how the learning dynamics of LLM finetuning shapes downstream generalization. Our analy- sis focuses on reasoning tasks, whose problem structure allows us to distinguish between memo- rization (the exact replication reasoning steps from the training data) and performance (the rectness of the final solution). We find that a model's performance on test prompts can be ef- fectively characterized by a training metric we call pre-memorizatiin train accuracy: the accu- racy of model samples on training queries before they begin to copy the exact reasoning steps from the training set. On the dataset level, this met- ric is able almost perfectly predict test accu- racy, achieving R² of 2 0.9 across various mod- els (Llama3 8B, Gemma2 9B), datasets (GSM8k, MATH), and training configurations. On a per- example level, this metric is also indicative of whether individual model predictions are robust perturbations in the training query. By connecting model's learning dynamies to test performance, pre-memorization train accuracy can inform train- ing decisions, such as the makeup of the training data. Our experiments on data curation show that prioritizing examples with low pre- mremmmoorizatinn accuracy leads to 1.5-2x improvements data ef- ficiency compared i.i.d. data scaling and other data scaling techniques.
vastly different. This
We foeus on mathematical problem-solving tasks, whose structure is particularly amenable for investigating this ques- tion. In reasoning tasks, models are trained to generate both final answer and intermediate reasoning steps. Although each problem has a single correct answer, the reasoning steps in the target solution trace represent just one of many valid ways to solve a problem. Therefore, a model that has memorized the training data is likely to replicate exact reasoning steps from the training data, while model with general problem-solving skills may produce the correct final answer but follow different reasoning path. By analyz- ing model responses on training queries, focusing on both the accuracy of the final answer and the similarity of the response to the target solution trace, we can gain insights into the generalizability of the model's learned solution.
Our findings reveal that, while LLMs often fully memo- rize the finetuning dataset by the end of training, model predictions for training queries prior to memorization are strongly indicative of final test performance. For certain examples, models first learn to generate diverse solution traces (distinct from the target solution trace) that lead to the correct final answer, before later memorizing the target solution trace. For other training examples, models only produce incorrect responses before memorizing the target trace. To capture this distinction, we introduce the concept of pre-memorization train accuracy: the highest accuracy a model achieves on a training example through the course
'Equal contribution 'UC Berkeley ²CMU. Correspondence to: Katie Kang <katickang @ cxes. becs.bbrrkley..du...
Proceedings of the ndnd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Generalization in LLM Mathematical Reasoning
0.6
0.6
0.5-
0.5
Test Accuracy
Tost Tost Accuracy
0.4-
0.4-
2c
0.3-
0.3-
2e-5
0.2
0.2
5
0.1
0.1
0.0
0.0
0.0
0.2
0.4
0.6
0.3
1.0
0.0
0.1
0.2
0.3
0.1
0.5
0.6
Train Accuracy
Pro-Memorization Train Accuracy
Figure 1. Relationship between train accuracy (left), pre- ne-memorization train accuracy (right), and test accuracy for models finctuned on GSM8k using Llama3 8B. Each line represents a training run, and each point represents an intermediate checkpoint. Pre-memorization train accuracy strongly correlates with test accuracy, while train accuracy does not.
of training before exactly memorizing the target solution trace. We find that a model's average pre-memorization train accuracy is highly predictive of the model's 's test accu- racy, as illustrated in Fig. 1. Our experiments show that this phenomenon holds across different models (e.g.. Llama3 8B (Dubey et al., 2024), Gemma2 9B (Team et al., 2024)), tasks (e.g., GSM8k (Cobbe al., 2021), MATH (Hendrycks et al., 2021)), dataset sizes, and hyperparameter settings, with coefficients of determination around or exceeding 0.9.
2. Related Works
A number of works have studied the phenomenon of mem- orization during training, but consider different definitions of memorization. One definition quantifies memorization with the "leave-one-out" gap. i.e., how much a model's prediction for an example changes if we were to remove it from the training data (Feldman & Zhang, 2020; Arpit et al., 2017: Zhang et al., 2017). Using this definition, some works argue that more memorization during training leads to worse generalization (Bousquet & Elisseeff, 2000), while others contend that memorization is actually necessary for generalization in long-tail distributions (Feldman, 2020). These works generally produce worst-case bounds on gen- eralization error within some class of training distributions. In contrast, our work presents a direct, empirical connec- tion between a model's train behavior and its test accuracy without relying on the computationally expensive "leave- one-out" metric. In the context of language models, oth- ers have defined memorized examples those where the model's output closely matches examples in the training data (Carlini et al., 2021; Tirumala et al., 2022: Inan et al., 2021; Hans et al., 2024), which is similar to our definition of memorization. However, these works mainly focus on privacy and copyright concerns, rather than connections to generalization.
We further find that the pre-memorization train accuracy can provide insights into the robustness of model predic- tions at a per-example level. For train examples with low pre-memorization accuracies, adding small perturbations to the training prompt causes the accuracy of model predic- tions significantly degrade. In contrast, for train examples with high ne-memoorzation accuracies, models gener- ally able to maintain high performance under perturbations. Thus, by measuring pre-memorization accuracy, we can identify specific training examples for which a model's pre- dictions are not robust, which can inform targeted improve- ments to the training strategy. As an example, we leverage findings to guide data curation. Our experiments show that training on data distributions that prioritize examples with low pre- memorization accuracy leads to 1.5-2x im- provement in sample efficiency over i.i.d sampling, and outperforms other standard curation techniques.
Beyond memorization, a number of prior works have stud- ied how other aspects of the learning process relate to generalization. Some works focus on metrics related to model complexity, such as VC dimension or parameter norms (Neyshabur et al., 2015; Bartlett al., 2019), while other works focus on empirically motivated measures, such as gradient noise (Jiang al., 2019) or distance of trained weights from initialization (Nagarajan & Kolter, 2019). Jiang et al. (2019) conducted comprehensive compari- son of these measures and found that none were consistently predictive of generalization, though their work primarily fo- cused on image classification. Other approaches have used unlabeled, held-out data to predict generalization, leverag-
The main contributions of this work are as follows: (1) we in- troduce the concept of pre-memorization train accuracy, and show that it is highly predictive of test accuracy for LLM rea- soning problems, (2) we show that pre-memorization train accuracy can also predict the robustness of individual model predictions for train examples, and (3) we leverage our ob- servations to improve the sample efficiency of data curation. By offering a deeper understanding of how a model's learn- ing dynamics shape its generalization, we hope our work can bring about more targeted and principled interventions for improving a model's 's reasoning capabilities.
2

---

Generalization in LLM Mathematical Reasoning
ing metrics such as the entropy of model predictions or the disagreement between different training runs (Garg al., 2022; Platanios et al., 2016; Jiang et al., 2021). Our find- ings show that pre-memorization accuracy can be a much stronger predictor of generalization in L.LM reasoning tasks.
train query may contain reasoning steps that differ from the target solution trace, while still arriving at the correct final answer. To quantify this difference, we will measure the distance between a model's prediction fo(y/az) and the target reasoning trace y/s with perplexity, defined as (erp(fo(y)a;) vi) = exp( -1 log(fo((i(zz))) where n; is the number of tokens in y/-
Finally, our work seeks to improve data curation, which has also been studied in number of prior works. Specific to LLM finetuning, prior data curation approaches largely fall into three categories: optimization- based, model-based and methods
m;
4. Connecting Learning Dynamics to Generalization
heuristic-based approaches. Optimization-basee frame data selection as an optimization problem, where the objective is model performance, and the search space con- sists of the training data distribution (Engstrom et al., 2024; Grosse et al., 2023). Model- approaches leverage char- acteristics of the learning process (Mekala et al., 2024; Liu et al., 2024), such comparing the perplexity of exam- ples (Li et al., 2023). Lastly, sauristic-based methods rely on simpler criteria, such as difficulty scores generated by off- the-shelf LLMs such as GPT, to classify desirable training data (Chen et al., 2023; Lu et al., 2023; Zhao et al., 2023). Our data curation approach aligns most closely with model- based strategies, as we use the model's pre-memorization accuracy, a characteristic of the learning process, inform the selection of training examples. Our experiments show that pre- meemiieorizttionn train accuracy can serve as an ef- fective metric for data curation, and outperforms previous approaches.
In this section, we will investigate the relationship between a model's learning dynamics during finetuning and its ability to generalize. Our findings show that, while models tend to memorize most of the training data after some number of epochs, their generated samples display varying levels of accuracy before memorization occurs. We find that this accuracy before memorization has a strong connection to the model's 's downstream generalization behavior.
4.1. Characterizing the Learning Dynamics of LLM Reasoning Finetuning
We begin by more precisely characterizing an LLM's learn- ing process when finetuning on reasoning tasks. We focus two key aspects of the model's behavior when presented with train queries: (1) whether the model's samples arrive at the correct final answer, and (2) the distance between the model's prediction and the target solution trace, mea- sured by perplexity. These two metrics, visualized in Fig. 2, offer different perspectives on the model's behavior, be- cause while there only one correct final answer for each query, there may exist many different valid reasoning traces. Tracking both metrics through the course of training allows us to measure how effectively the model is able to solve training queries, and the extent to which this is accomplish by replicating the target solution trace.
3. Preliminaries
We focus on training LLMs to perform reasoning tasks via finetuning, We are provided with training dataset Duain = (x), where queries H, are drawn from P(z) and solution traces y; are drawn from P(y) x)) We assume the test dataset, Dust. is generated from the same distribution as the training data. The model finetuned by minimizing next-token prediction loss. We denote the finetuned model as ()), and model predictions as ÿ fo(y)zz)
In Fig. 3, we visualize the learning progression, as charac- terized by the two metric described above, for three models finetuned on GSMSK. Each model is trained for six epochs, with a distinct peak learning rate that decays to zero by the end of training. As expected, training accuracy improves over time as the model minimizes the loss (color gradient from dark to light), and the distance between predictions and target solution traces decrease (from pink to yellow). For some learning rate settings, models approach pear-perfect accuracy by the end of training. and their predictions closely match the target reasoning traces (mostly yellow in bottom row). However, during early stages of training, we observe significant differences in model behavior. For some train queries, models initially produce incorrect samples (black), and then directly transition to replicating the target trace (yellow). For other examples, models first learn to generate correct answers with solution traces that differ from the tar-
In reasoning tasks, solution traces y consist of both inter- mediate reasoning steps and a final answer, denoted as Ans(y). Our goal in a reasoning task is for the model to generate solution traces with the correct final answer when faced with previously unseen queries. We measure the accuracy of model samples for a given query x; using (e((ff(((z))) v)) = fo( [1 (A.)(j.)
Ans(y.)) In our experiments, we approximate this accuracy by sam- pling from the model with temperature of 0.8 and averag- ing the correctness attained by the samples.
While different solution traces drawn from P(ylæ) should all have the same final answer, the target solution trace y of an example represents only one of many valid so- lution traces for solving z;. Thus, model samples for
3

---

Generalization in LLM Mathematical Reasoning
Teaming Progressionn
Learting Progression 2
Train Example
Momortzatoon Tareshold
Prompt: Natalia sold clips to 48 of her friends in April, and then she sold half as many clips in May. How many clips did Natalia sell altogether in April and May?
Train Acauresy
Pre Memorization Accuracy
Solution Trace: Natalia 48/2 = <<48/2-24>>22 clips May. Natalia sold
48+24 <<48+24-72b>72 clips altogether April and May. #860
0.75
Model Samples
Accuracy
A; In May, Natalia sold half as many clips as in April, so she sold 48 / =
0.5
<<48/22 2222222224 clips. She sold allogether 48 24 <<dB-24=72>372 clips in April and May. #### 72
B: In April, Natalia sold = <<482=-9>>>% clips. In May, Natalia sold /
0,25
<<96/ /2 22-8>>488 clips. Natalia sold 96 <<96145 144>>144 clips
altogether April and May. #6ие 144
C: Natalia sold 48/ 2 <<48/228/ 2<18/2=24>>21 clips in May. Natalia sold 48+24
<<48+24-72-0 22=722>72 clips altogether in April and May. #### 72
figh
Medium
Low
Perplexity
Figure 2. Visualizations of different learning progressions, measured by the accuracy of model samples (light vs. dark) and the perplexity of target solution traces under model predictions (pink vs. yellow). Right side presents examples model samples with (A) high aciuracy+high perplexity, (B) low accuracy+high perplexity, and (C) high accuracy+low perplexity. Black text represents exact match with the target solution trace, while grey text represents parts that do not match.
LR 2e-4
LR 2e-5
LR 5e-7
2
Booch
3
4-
5-
6-
Examples
Examples
Examples
Figure 3. Predictions of 3 different models through the course of training. X-axis represents individual training examples. Y-axis represents training epoch. Color represents model predictions for each example in terms accuracy and perplexity (legend in Fig. 2).
trace (pink), before later transitioning to fully replicating the target trace (yellow).
Perp (folla))) vi) < where is threshold (fixed across examples). We further define modified measure of accu- racy, whose value is masked to zero if the model's prediction for that example is considered memorized, as follows:
In this work, we will refer to model predictions with low dis- tance to target solution traces as memorization. We can see that when finetuned with different learning rates, different models exhibit different capacities for generating accurate samples before memorizing target solution traces (amount of pink). Models with low accuracy before memorization may be largely learning verbatim mappings from training queries to target traces, which would not generalize to new queries. contrast, models with high accuracy before mem- orization demonstrate an ability to arrive at correct answers through varied reasoning paths, suggesting that they have developed more generalizable problem-solving capabilities.
(askedAcc(fe(ylve.)) 94. =Acc((fe(y)z.)) vi) ((Perp(fo w) pl.
Now let fe.,, denote the model epoch of training. Us- ing our definition of masked accuracy, we define the pre- memorization accuracy as follows:
PreMemAce (f)
=min
max MaskedAcc(So.. (v) (1)24)) Vi. 15m Sm
(cl vi)
To better quantify this phenomenon, we introduce a met- ric called rreememorization accuracy. We consider train example (ze, yi) E Diran to be memorized by fo(y)æ) )) if
This quantity can be roughly interpreted as the best accu- racy that the model achieves for a training prompt thus far

---

Generalization in LLM Mathematical Reasoning
GSM8k / Llama3 8B
MATH Llama3 8B
0.16
0.5
0.12
LR
Tesl. Accuracy
0.4
2e-4
5e-4
0.3
0.08
2e-5
5e-5
0.2
5e-6
0.04
5e-7
0.0
0.00
Fpochs
0.1
0.4
0.5
0.00
0.04
0.08
0.12
0.16
GSM8k Gemma2 9B
MATH Gemma2 9B
0.7
0.35
Dalasel
Size
Test Accuracy
0.6
0.25
Full
Half
0.5
Quarter
0.15
0.4
0.3
0.05
0.3
0.4
0.5
0.6
0.7
0.05
0.15
0.25
0.35
Pre-Memorization Train Acouracy
Pre Memorization Train Accuracy
Figure Evaluating the relationship between pre- pre-memorization train accuracy and test accuracy. Each line corresponds a training run, and each marker corresponds to specific checkpoint. Pre-memorization train accuracy strongly predict test accuracy across tasks, models, and training settings.
in training before it memorizes the target trace. Unlike standard accuracy, which evaluates performance at specific checkpoints, pre-mmmorization accuracy evaluates the en- tire training process up epoch m. There is an additional minimum taken with the accuracy of model predictions at epoch m, which compensates for examples whose accura- cies decrease through training (though this is uncommon).
sweeping across range of values (see Appendix A).
In Fig. 4, we plot the pre-memorization training accuracy and test accuracy across different training runs. We used Llama3 8B and Gemma2 9B as base models and GSM8K and MATH the reasoning tasks. To evaluate different generalization behaviors, we finetuned the models by adjust- ing the peak learning rate (ranging from 5e-7 to 5e-4), the number of training epochs (1, 3, 6), and the dataset size (full, half, or quarter of the original dataset). We use the same value p within each plot. full list of the training runs in our experiments and other details can be found Ap- pendix B. We observe a strong linear relationship between pre-memorization training accuracy and test accuracy, with the results closely following the x line across different models, tasks, and hyperparameter settings. More quantita- tively, the coefficients determination associated with each plot are 0.94 (GSM8k Llama), 0.95 (MATH Llama), 0.97 (GSM8K Gemma), and 0.88 (MATH Gemma). Our results show that pre-memorization training accuracy is reliable predictor of test accuracy.
4.2. Pre-Memorization Train Accuracy Strongly Predicts Test Accuracy
We next use aremmioooizttion accuracy to analyze the con- nection between learning dynamics and downstream general- ization. We find that a model's average pre-memorization train accuracy is highly predictive of its test accuracy across variety of training runs and checkpoints. More concretely, find that there exists value of p for which a model's average pre-memorization train accuracy, Duin [PreMemAce(fo, (y]a;;) P)), closely approxi- mates the model's 's test accuracy, ( The value of the memorization threshold p is fixed across examples and training parameters, but may need to re- calibrated for different tasks or models. We calibrate p by
(y]=;),, ys]]-
As discussed in Section 2, various metrics have been pro- posed previous studies to predict the generalization gap,
5

---

Generalization in LLM Mathematical Reasoning
GSM8k / Llama3 8B
GSM8k Llama3 8B
GSM8k / Llama3 8B
GSM8k / Llama3 8B
Generalization Gap
0.4-
0.4
0.4
0.4-
/
0.2
0.2
0.2
0.2
0.0
0.0
0.0
0.0-
0.0
0.2
0.4
0e-10
4e-10 Gradient Variance
2e-10
100
10a
10
0.1
0.3
0.5
0.7
Train Acc- Pre-Mem Acc
Dist from Init
Train Acc. - ATC
Figure 5. Evaluating different generalization metrics vs. the ground truth generalization gap for models finetuned on GSM8k using Llama3 8B (legend in Fig. 4).
the difference between train and test accuracy. In Fig. 5, we compare several of these existing metries, including gradient variance (Jiang et al., 2019), distance between current model weights and initialization (Nagarajan & Kolter, 2019), and an estimate of test accuracy via Average Thresholded Con- fidence (ATC) (Garg et al., 2022) (details in Appendix C). The correlation coefficients associated with each metric (left to right) are 0.98, -0.72, 0.59, -0.04, which shows that the prior metrics do not correlate as strongly with test accuracy as our proposed metric.
preambles to the solution trace-phrases such as "First" or "We know that deviate from the target solution trace, which we visualize in Fig. 6. Because these generie phrases are plausible preambles to valid reasoning traces, we would expect a model which has learned robust solution to an example to still be able to arrive at the correct final answer. In contrast, if the model is unable to produce the correct final answer given these generic phrases, then the model likely to have learned only regurgitate the training response.
In Fig. 6, we show the prediction behavior of two models, both trained for six epochs with learning rate of 2e-5, on the GSM8K and MATH datasets. We can see that while model predictions are near-perfect for unaltered training prompts, their accuracy significantly degrades when pre- sented with perturbed prompts. Furthermore, we see that the accuracy of train examples with low pre-memorization train accuracy tends degrade much more than those with high pre-memorization train accuracy. These findings sug- gest that pre-memorization train accuracy can predict the robustness of model predictions for individual train exam- plcs. Note that while our perturbation analysis makes use of
5. Per-Example Analysis of Generalization
In this section, we go beyond aggregate test accuracy and show that tracking per-example pre-memorization accuracy offers a window into the model's behavior at the level of individual training examples. Specifically, we find that the pre-memorization train accuracy of given example is pre- dictive of the robustness of the model's prediction for that example. This example-level accuracy helps us identify subsets of the training data for which the model struggles to learn robust solutions and offers opportunities to improve training through targeted interventions. We explore how this insight can inform data curation strategies, showing that prioritizing examples with low pre-memorization train accuracy during data collection can lead to significant im- provements over i.i.d. data collection and other common data curation methods.
dask-dependent preambles, pre- memorization train accuracy does not require any domain knowledge. Therefore, pre-memorization train accuracy provides practical way to identify fragile examples where the model may have learned overly specific or non-robust patterns, which offers practical applications for improving model generalization.
5.1. Predicting Model Robustness with Pre-Memorizatiin Train Accuracy
5.2. Curating Data with Pre-Memorization Train Accuracy
We begin by examining the relationship between an indi- vidual example's 's pre-memorization train accuracy and the robustness of the model's predictions for that example. Our findings show that model predictions tend be less robust for train examples with low pre-memorization accuracy.
By offering insights into the robustness of individual model predictions, pre-memorization train accuracy can provide targeted guidance for improving a model's 's generalization. In this section, we explore data curation as a practical ap- plication of our findings. Prior work has suggested that focusing "harder" examples, where the model struggles to learn robust solutions, can lead to more sample-efficient improvements (Li et al., 2023; Chen et al., 2023). However, identifying useful metrics for determining example diffi-
To assess the robustness of model predictions, we analyze how the model responds to small perturbations in the input prompt. We present the model with both the original training queries, as well as training queries appended with short
6

---

Generalization in LLM Mathematical Reasoning
Robust Solution
Non-Robust Solution
## 105
Train Solution Trace
First, ...
First,
Train
Natalia ...
Train Prompt
Natalia
####72
Prompt
#####72
We know that ...
We know that
####633
Figure 6. Visualization of the robustness of model predictions to perturbations in the prompt, including the original training prompt (purple), original prompt "First" (pink), and original prompt + "We know that" (teal). A robust model prediction would arrive at the correct final answer even if the perturbations changes the reasoning steps. In contrast, a non-robust model prediction produces incorrect final answer when the prompt diverges from the training data.
GSM8k Train Prompts
GSM8k Train Prompts + "We know that"
GSM8k Train Prompts
*FFirst"
Sample Accuracy
1.0
1.0
1.0
0.8
0.8
0.8
0.6
0.6
0.6
0.4
0.4
0.4
0.2
0.2
0.2
0.0
0.0
0.0
0
0.2
0.4
0.6 0.8
0
0.2
0.4
0.6 0.8 1
0.2
0.4
0.6 0.8 1
MATH Train Prompts + "Firsst"
MATH Train Prompts
MATH Train Prompts
"We know that"
Sample Accuracy
1.0
1.0
1.0-
0.8-
0.8
0.8
0.6
0.6
0.6
0.4
0.4
0.4
0.2
0.2
0.2
0.0
0.0
0.0
0
0,2 0.4 0.6 0.8
0.2 0.4 0.6 0.8
0.2
0.4 0.6 0.8
Pre-Mem Train Acc
Pre-Mem Train Acc
Pre-Mem Train Acc
Figure 7. Accuracies of model samples (y-axis) when faced with the original prompt (left) and prompts with perturbations (middle, right). The x-axis represents bins of pre-memorization train accuracies associated with each prompt. Solid line denotes the average, and violins denote distributions within each bin. While the accuracy of model samples is almost perfect when faced with original prompts, significantly degrades when faced with prompts with perturbations. Furthermore, the degradation of accuracy is much more significant for train examples with low pre-memorization accuracy than those with high pre-memorization accuracy, showing that per-example pre-memorization train accuracy can provide insight into the robustness of a model's individual predictions.
culty remains an open challenge. We investigate the use of pre-memorizatio train accuracy as a metric for guiding data curation, and find that it outperforms i.i.d. sampling and other standard data curation approach in sample efficiency for reasoning tasks.
erties. In our experiments, we take Dunin to be the original dataset, and collect new examples by using GPT to rephrase examples in the original dataset, similar to the procedure in (Setlur al., 2024). By only collecting new examples that derive from the specified empirical distribution, ensure the new dataset approximates P" P'(xx). This setup can also be used when collecting new human- generated data, by providing the specified empirical distribution of examples as references for human labelers.
We will first more precisely define our data curation prob- lem. Given an existing set of N training examples with queries distributed as P(z), we aim to collect N' exam- ples, denoted as D'min to augment the dataset. The goal to specify new distribution P'(x) that maximizes the test performance of model trained on both the original and the newly collected examples. While defining the true distribution of queries can be challenging, we assume that by approximating it with an empirical distribution from the current dataset, we can collect new data with similar prop-
Our approach for data collection prioritizes examples with low pre-memorization accuracy. First, we calculate the pre-memorization accuracy for each example in the current dataset and then define P'(x)) as the distribution of exam- ples whose pre-memorization accuracy falls below a certain threshold t. We then collect new data according to this dis-
7

---

Generalization in LLM Mathematical Reasoning
Algorithm I Our Data Collection Process
1: Input: N' = N + +... N t
2: Output: Updated dataset Dirain
3: Initialize Diainn
0
4: for to do
5: Train model Duain Diain
6: Evaluate model on Dirain and compute pre-memorization accuracy for each example
7: Set (rr as the distribution of examples with pre-memorization accuracy below
8: Collect NJ new examples from P/ (x) and add them to D
9: end for
GSM8k I.lmm33 AB
MATH Levels 1-3 DeepSeekMath 7B
0.68
>xx
0.54
1.5x
0.66
Test, Accuracy
0.50
0.64
Ours
Hourlstic
0.60
0.58
0.42
0
40000
83000
120000
20000
40000
60000
Number of Train Examples
Number Train Examples
Figure 8. Comparison of different approaches for data curation. Each line represents a different data curation approach, and each point represents different training run. Our approach acheived the best sample efficiency compared to the other approaches.
tribution. If N' is large, we can split the data collection process into multiple iterations +...+ + N₁ = N'). In each iteration, we collect N new examples according to P. (x), retrain a model on the combined dataset, calculate the pre-memorization accuracy with the model, and update Pi+1 (z)) for the next round of data collection. This process summarized in Algorithm 1.
compared to i.i.d scaling in GSM8K, and 1.5x sample ef- ficiency on MATH levels 1-3. Furthermore, we find the gap in performances increases with dataset size, which sug- gests that better data curation metries may become more important as models become more capable. These results highlight the effectiveness of pre-memorization accuracy as criterion for targeted data collection, leading to enhanced generalization with fewer data points. We provide more details about our implementations in Appendix D.
We compare our strategy to i.i.d. sampling and two existing approaches commonly used in data curation. Both of these approaches propose a metric of example difficulty and pri- oritize difficult examples during data collection. The first Difficulty (IFD) (Li
6. Conclusion
metric, called EnstructionFFollovee al., 2023), computes the ratio between the perplexity of training labels given inputs and the perplexity of only labels using model finetuned for the task. The second metric uses heuristic notions of difficulty measured by external sources such as humans or more capable models (Chen et al., 2023; Lu et al., 2023; Zhao et al., 2023). For GSMSK, we use the number of lines in the target solution traces as a heuristic the difficulty levels
Our work studies the relationship between learning dynam- ies and generalization in L.MSs finetuned for reasoning tasks. We introduce the concept of pre-memorization train accuracy and show that it is a strong predictor of its test accuracy. We further show that model's per-example pre- memorization train accuracy can be an indicator of the ro- bustness model predictions for those examples. Finally, we leverage this insight for data curation, and show that priori- tizing examples with low pre-memorization train accuracy can be more effective than i.i.d. data scaling and other data curation techniques. We hope that by providing way at- tribute a model's generalization to specific aspects of the training process, our work can enable the design of more effective and principled training strategies.
for difficulty, while for MATH, provided in the dataset itself.
In Fig. 8, we evaluate the different data curation approaches for finetuning on GSM8k with Llama3 8B and MATH lev- els 1-3 with MeepSeekMath 7B (Shao et al., 2024). Our approach outperforms all three prior approaches, achieving 2x the sample efficiency for the same target accuracy
8

---

Generalization in LLM Mathematical Reasoning
Impact Statement
Garg, S., Balakrishnan, S., Lipton, Z. C., Neyshabur, B., and Sedghi, H. Leveraging unlabeled data pre- dict out-of-distribution performance. arXiv preprint arXiv:2201.02234 2022.
This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.
Grosse, R., Bae, J., Anil, C., Elhage, N., Tamkin, A., Tajdini, A., Steiner, B., Li, D., Durmus, E.,, Perez, E., et al. Study- ing large language model generalization with influence functions. arXiv preprint arXiv:2308.03220 2023.
References
Arpit, D., Jastrzebski, S., Ballas, N., Krueger, D., Bengio, E., Kanwal, M. S., Maharaj. T., Fischer, A., Courville, A., Bengio, Y., et al. A closer look at memorization in deep networks. In International conference on machine learning, pp. 233 242. PMLR, 2017.
Hans, A., Wen, Y., Jain, N., Kirchenbauer, J., Kazemi, H., Singhania, P., Singh, S., Somepalli, G., Geiping, J., Bhatele, A., et al. Be like goldfish, don't memorize! mit- igating memorization in generative Ilms. arXiv preprint arXiv:2406. 2024.
Bartlett, P. L., Harvey, N., Liaw, C., and Mehrabian, A. Nearly-tight vc-dimension and pseudodimension bounds for piecewise linear neural networks. Journal of Machine Learning Research, 20(63): 2019.
Hendrycks, D., Burns, C.,, Kadavath, S., Arora, A., Basart, S., Tang, E., Song. D., and Steinhardt, J. Mcasuring math- ematical problem solving with the math dataset. arXiv 2021.
preprint arXiv:2103.0338
Bousquet, O. and Elisseeff, A. Algorithmic stability and generalization performance. Advances in neural informa- tion processing systems, 13, 2000.
Inan, H. A., Ramadan, O., Wutschitz, L., Jones, D., Rühle, V., Withers, J., and Sim, R. Training data leakage analysis in language models. arXiv preprint arXiv:2101.052 2021.
Carlini, N., Tramer, F., Wallace, E., Jagielski, M., Herbert- Voss, A., Lee, K., Roberts, A., Brown, T., Song, D., Erlingsson, U., et al. Extracting training data from large language models. In 30th USENIX Security Symposium (USENIX Security 21), pp. 2633-2650, 2021.
Jiang, Y., Neyshabur, B., Mobahi, H., Krishnan, D., and Bengio, S. Fantastic generalization measures and where find them. arXiv preprini
2019.
Jiang, Y., Nagarajan, V., Baek, C., and Kolter, J. Z. As- sessing generalization of sgd via disagreement. arXiv preprint arXiv:2106. 13799, 2021.
Chen, L., Li, S., Yan, J., Wang, H., Gunaratna, K. Yadav, V., Tang, Z., Srinivasan, V., Zhou, T., Huang, H., et al. Alpagasus: Training better alpaca with fewer data, arXiv preprint arXiv:2307.0m 20002300,087001, 2023.
Li, M., Zhang, Y., Li, Z., Chen, J., Chen, L., Cheng, N., Wang. J., Zhou, T., and Xiao, J. From quantity to quality: Boosting IIm performance with self-guided data selection for instruction tuning. arXiv preprint arXiv:2308. 12032, 2023.
Cobbe, K., Kosaraju, V. Bavarian, M., Chen, M., Jun, H., Kaiser, L.,, Plappert, M., Tworek, J., Hilton, J., Nakano, R., et al. Training verifiers to solve math word problems. arXiv preprint arXiv:2 10. 14168, 2021.
Liu, W., Zeng, W., He, K., Jiang, Y., and He, J. What makes good data for alignment? a comprehensive study of automatic data selection in instruction tuning, 2024. URL https: http//axxxvv...gaaaa// org/abs./......ooo.... /2t///3322222 15685.
Dubey, A., Jauhri, A., Pandey, A., Kadian, A., Al-Dahle, A., Letman, A., Mathur, A., Schelten, A., Yang, A., Fan, A., et al. The llama herd of models. arXiv preprint 2xXiv:2407.21783 2024.
Lu, K., Yuan, H., Yuan, 7., Lin, R., Lin, J., Tan, C., Zhou, C., and Zhou, lyzing supervised fine fine-tuning of large language models. In The Twelfth International Conference Learning Representations, 2023.
Engstrom, I.,, Feldmann, A., and Madry, A. Dsdm: Model- aware dataset selection with datamodels. arXiv preprint arXiv:2401. 12926, 2024.
instag: Instruction tagging for ana-
Feldman, V. Does learning require memorization? a short tale about a long tail. In Proceedings of the 52nd Annual ACM SIGACT Symposium on Theory of Computing, pp. 954 959, 2020.
Mekala, D., Nguyen, A., and Shang. J. Smaller language models are capable of selecting instruction-tunin train- ing data for larger language models. arXiv preprint arX.v:2402.1042 2024.
Feldman, V. and Zhang, C. What neural networks mem- orize and why: Discovering the long tail via influence estimation. Advances in Neural Information Processing Systems, 33:2881 13:2881-2291, 2020.
Nagarajan, V. and Kolter, J. Z. Generalization in deep networks: The role of distance from initialization. arXiv preprint arXiv: 10010016700007 2019.

---

Generalization in LLM Mathematical Reasoning
Neyshabur, B., Tomioka, R., and Srebro, N. Norm-based capacity control in neural networks. In Conference learning theory, pp. 1376-1401. PMLR, 2015.
Platanios, E. A., Dubey, A., and Mitchell, T. Estimating accuracy from unlabeled data: A bayesian approach. In International Conference Machine Learning, pp. 1416- 1425. PMLR, 2016.
Setlur, A., Garg, S., Geng, X., Garg, N., Smith, V., and Kumar, A. RI on incorrect synthetic data scales the effi- ciency of IIm math reasoning by eight-fold. arXiv preprint arXiv:2406. 14532, 2024.
Shao, Z., Wang, P., Zhu, Q., Xu, R., Song, J., Zhang, M., Li, Y., Wu, Y., and Guo, D. Deepseekmath: Pushing the limits of mathematical reasoning open language models. arXiv preprint arXhv:2402.03300 2024.
Team, G., Riviere, M., Pathak, S., Sessa, P. G., Hardin, C., Bhupatiraju, S., Hussenot, L., Mesnard, T., Shahri- ari, B., Ramé, A., al. Gemma Improving open language models at practical size. arXiv preprint arXiv:2408.001 18, 2024.
Tirumala, K., Markosyan, A., Zettlemoyer, L., and Agha- janyan, A. Memorization without overfitting: Analyzing the training dynamics of large language models. Ad- vances in Neural Information Processing Systems, 35: 38274-38290, 2022.
Zhang, C., Bengio, S., Hardt, M., Recht, B., and Vinyals, O. Understanding deep learning requires rethinking gen- eralization, 2017. URL httts:://aaiiiii /ttp///aavvv///a///v//eae//a.ava org/ /bs/
Zhao, Y., Yu, B., Hui, B., Yu, H., Huang, F., Li, Y., and Zhang, N. L. A preliminary study of the intrinsic relation- ship between complexity and alignment. arXiv preprint arXiv:2308.085. 05696, 2023.
10

---

Generalization in LLM Mathematical Reasoning
A. Selection of Memorization Threshold
find the threshold p by sweeping across 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. In Fig. we illustrate how the value of p influences the R² for predicting average accuracy. We can see that R² degrades smoothly with respect p, which makes is relatively easy to find good value of p by sweeping range of values.
This calibration process only requires small number of training runs (e.g. 1-3) to arrive at a robust value of which can generalize to new training runs on the same model and finetuning dataset, illustrated in Fig. 10. However, it is important the the training runs used for calibration exhibit some spread over test accuracies, and memorization during training.
Finally, we also show that the calibration process generalizes to new test examples. We divide the test set into two halves: calibration set, and a heldout test set. We calibrate p on the calibration test set, and evaluate the coefficient of determination of the heldout test set. In Fig. 11, on which it had not been calibrated, achieving high coefficient of determination.
that the value of p able to generalize robustly to new examples
GSM8k / Llama3 8B
MATH / Llama3 8B
0.9
0.8
0.8
0.6
0.7
&
0.4
0.6
0.5
0.2
0.4
-3.0
-2.8
-2.6
-2.4
-22
-2.0 -1.8
-2,2
-2.0
-1.8
-1.6
-1.4
-1.2
log(log(p))
log(log(p))
Figure 9. Relationship between the value of and the cocfficient of determination (R²) with respect to pre-memorization train accuracy and test accuracy. The R² is taken in aggregate of all the corresponding training runs in Fig. 4.
Calibration Training Runs
Heldout Training Runs
0.6
Se-5 opocas
0.6
2:20 opechad apects
0.5
0.5
Test Accuracy
0.4
0.4
oporas3)
Se7 opecha
0.3
0.3
sul
qurter
0.2
0.2
0.1
R² = 0.973
0.1
R² 0.868 log(log(p)) = - 224
log(log(p)) - 2.4
0.0
0.0
0.0
0.1
0.2
0.3
0.4
0.5
0.6
0.0
0.1
0.2
0.3
0.4
0.5
0.6
0.6
be-b epuchs
0.6
opocia
apecas
0.5
Ir 2e-1 epochs
0.5
Test Accuracy
365 opedua
0.4
0.4
36.7 opecha
openas
0.3
0.3
qurter
0.2
0.2
0.1
R² = 0.988
0.1
R²= 0.899
log(jog(p)) 2.2
log(log(p)) 2222
0.0
0.0
0.0
0.1
0.2
0.3
0.4
0.5
0.1
0.2
0.3
0.4
0.5
0.6
Pre-Memorization Train Accuracy
Pre-Memorization Train Accuracy
Figure 10. Calibrating subset of training runs, and evaluating Rt² heldout training runs using GSMSK and Llama3 8B. We see that calibrating on just 1-3 training runs was able to yield robust value of which leads to high It? on heldout training runs.

---

Generalization in LLM Mathematical Reasoning
Calibration Test Set
Heldout Test Set
0.6
0.6
0.5
0.5
Test Accuracy
0.4-
0,4
0.3
0.3
0.2
0.2
0.1
=
0.1
0.998
log(log(p)) 2.30
log(log(p)> -..0
0.0
0.0
0.0
0.1
0.2
0.3
0.5
0.0
0.1
0.2
0.4 0.5
0.6
Pro-Memorization Train Accuracy
Pre-Memorization Train Accuracy
Figure 11. Calibrating using a subset the test set (calibration test set), and ovaluating R² on heldout test set using GSM8k and
Llama3 8B. We can that calibrating just the calibration test set was able yield robust value of which leads high R on the heldout test set.
B. Section 4.2 Training Runs Details
In this section, will enumerate all training runs shown in Fig. 4 and their training details. For our half and quarter training runs, we fix the total number of training steps to be equivalent to training for epochs the full dataset.
B.1. GSM8k LLama3 8B
For all training runs with GSM8k and Llama3 8B, the AdamW optimizer, with linear decay learning rate scheduler with 20 warmup steps, a batch size of 128, and max gradient norm of 2.
Learning Rate
Epochs
Dataset Size
5e-5
6
full
2e-5
6
full
5e-7
6
full
2e-4
6
full
5e-5
3
full
2e-5
3
full
5e-7
3
full
2e-4
3
full
5e-5
1
full
5e-7
1
full
2e-4
1
full
2e-5
6
half
2e-5
12
quarter
B.2. MATH LLaaa3 8B
For all training runs with MATH and Llama3 8B, we use the AdamW optimizer, with a linear decay learning rate scheduler with 20 warmup steps, a batch size 24, and max gradient norm of 2.
12

---

Generalization in LLM Mathematical Reasoning
Learning Rate
Epochs
Dataset Size
5e-5
6
full
5e-7
6
full
2e-4
6
full
5e-5
3
full
5e-7
3
full
2c-4
3
full
5e-5
1
full
5e-7
full
2e-4
full
2c-5
6
half
2c-5
12
quarter
B.3. GSM8K Gemma2 9B
For all training with GSM8K and Gemma2 9B, the Adam optimizer, with cosine decay learning scheduler with (0.1 *total steps) warmup steps, batch size of 32, and gradient norm of 1.
Learning Rate
Epochs
Dataset Size
5e-4
6
full
5e-5
6
full
5e-6
full
5c-7
6
full
5e-4
3
full
5e-5
3
full
5e-6
3
full
5e-7
3
full
5e-4
full
5e-5
1
full
5e-6
full
5e-7
1
full
5e-5
6
half
5e-5
12
quarter
B.4. MATH Gemma2 9B
For all training runs with MATH and Gemma2 9B, the Adam optimizer, with cosine decay learning rate scheduler with (0.1 *tota steps) warmup steps, batch size of 32, and max gradient norm of 1.
Learning Rate
Epochs
Dataset Size
5e-4
6
full
5e-5
full
5e-6
6
full
5c-7
6
full
5e-4
3
full
5e-5
3
full
5e-6
3
full
5e-7
3
full
5e-4
full
5e-5
I
full
5e-6
I
full
5e-7
1
full
5e-5
6
half
5e-5
12
quarter
13

---

Generalization in LLM Mathematical Reasoning
C. Section 4.2 Prior Generalization Metrics
In this section we will more precisely describe each generalization metric.
C.1. Gradient Variance
We calculate the gradient of the model for 5 different minibatches, take the variance across the 5 samples for each element of each weight matrix, and take the average over each element of the model weights.
C.2. Distance from Initialization
We calculate the squared difference between each element of the model weights at initialization and after finetuning, and take the sun across all elements.
C.3. Average Thresholded Confidence (ATC)
ATC computes a threshold on score computed on model confidence such that the fraction of examples above the threshold matches the test accuracy. For the score, we use the likelihood of greedily sampled responses under the model. We calculate the the score over the training data using a model trained for 3 epochs using learning rate 2e-5, and calculate the threshold over the score using the test dataset. We then predict the test accuracies over different models our experiment by calculating the score associated with the training data using each model, and measuring the percentage of examples whose score surpass the threshold that previously calculated.
D. Section 5.2 Implementation Details
For our approach for data curation, we implemented the process described Algorithm 1, with iterations (n) and using threshold (t) 0.75 for both GSM8K and MATH.
For the IFD approach for data curation, calculated the IFD) using model that was train the test associated each dataset for 2 epochs. This is because, in order to calculated the IFD score, we need model which has been briefly trained for the task of interest, but which has not been exposed to the dataset for which we want to calculate the IFD score over. Note that this model is only used for calculating for the IFD score, and not used for evaluations our experiments, so there is no data leakage.
For both the IFD approach and the heuristic approach, we take P'(x)) be top 50 percentile examples for GSM8k, and top 75 percentile examples for MATH. We designed these percentiles to roughly match the percentile of examples that approach selects from.
For all training runs, we use the AdamW optimizer, with linear decay learning rate scheduler with 20 warmup steps, batch size of 128, a max gradient norm of 2, a learning rate of 2e-5, and epochs of training.
14


Reviews:
[{"review_id": "4K6cVrokRT", "reviewer": "Reviewer_tBbx", "summary": "This paper investigates how the learning dynamics of large language models (LLMs) during finetuning influence their generalization on reasoning tasks. The key contribution is the introduction of pre-memorization train accuracy, defined as the highest accuracy a model achieves on a training example before it begins to replicate the exact reasoning steps (memorization) from the training data. The authors demonstrate that this metric strongly correlates with test accuracy across models (Llama3 8B, Gemma2 9B), datasets (GSM8k, MATH), and training configurations (learning rates, epochs, dataset sizes), achieving coefficients of determination bigger than 0.9. Furthermore, this paper discusses the relationship between pre-memorization accuracy and the model prediction robustness, finding that at the per-example level, low pre-memorization accuracy indicates fragile model predictions. Perturbations to such examples (e.g., altering prompts) significantly degrade accuracy, whereas high pre-memorization examples remain robust. Additionally, this paper proposes a data curation algorithm based on the pre-memorization accuracy, which could yield 1.5-2 gains in data efficiency over i.i.d. sampling and outperforms heuristic-based and optimization-based data selection methods.\n## update after rebuttal\nI raised my score to 3 (boardline) since the authors' detailed response to me and other reviewers. However, some of my questions were not addressed and I still regard authors should take a deeper exploration about this very interesting phenomenon, standardize their method's detail and implement the pre-mem accuracy on additional datasets.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the feedback! In this rebuttal we provide a more detailed discussion of the implications of our experimental findings, and provide additional explanations for our experimental design. \n\n> The authors need ​a deeper analysis and discussion​ of pre-memorization accuracy. \n\nIn the following response, we will summarize the main motivation and findings of our paper, and provide new discussion regarding the *implications* of these findings and how they relate to the motivation. We acknowledge that our paper currently lacks discussion with regards to the implications of our work, and we will update it to include this in more detail. \n\n**Motivation**: Our work is motivated by the following phenomenon: models derived from the same pretrained model and finetuned on the same finetuning dataset can have vastly different test performance, even when their train accuracies are nearly perfect. Here, the optimization of the training loss is not the issue, which suggests that something about the models’ learning *dynamics* is causing the difference in models’ generalization. The goal of our work is to better understand the relationship between a model’s learning dynamics and generalization.\n\n**Findings**: The main finding of our work is that pre-memorization train accuracy (a metric of a model’s learning dynamics) is a good predictor of a model’s downstream generalization. We provide experiments showing this both in aggregate and on a per-example basis. \n\n**Implications**: Our findings suggest that models exhibit two different modes of optimizing the training objective: one which leads to downstream generalization and one which doesn’t (learning progression 1 vs 2 in Fig. 2). Train accuracy conflates the two models of learning. In contrast, pre-memorization train accuracy allows us to delineate between the two, and identify the examples that are learned in a “generalizable” manner and the examples that are not. This analysis sheds lights on the aspects of a model’s learning dynamics that lead to the differences in downstream generalization, and provides us with a tool that enables more nuanced interventions to the training recipe (e.g. via data curation) to improve generalization. \n\n> The authors categorize samples with high perplexity (ppl) and high accuracy as well-generalized. However, it is possible the model simply memorizes the final answer rather than learning reasoning steps. Is there any analysis or experiment addressing this concern?\n\nYes, in Sec 5.2, we present experiments where we perturb the prompt of the training examples, and measure how much the model’s predictions degrade from these perturbations. We find that model predictions for examples with high pre-memorization train accuracy tend to maintain high accuracy, while predictions for low pre-memorization train accuracy tend to degrade much more significantly. This experiment addresses the reviewer’s concern, as it shows that examples with low pre-memorization train accuracy actually do not generally maintain the ability to produce correct final answers for perturbed prompts. \n\n> Is pre-memorization accuracy effectively equivalent to accuracy under very small learning rates since it is measured early in training?\n\nNo, we do not believe that pre-mem accuracy will always be equivalent to train accuracy under very low learning rates, because prior works have found that smaller learning rates can lead to more memorization (see [1]). So in the general case, we don't think our experiments can be used to imply any relationship between learning rate and memorization.\n\n[1] Fitting Larger Datasets by Learning with Small Batch Sizes and Regularized Initialization\n\n> Figure 4 uses Llama3 8B and Gemma2 9B, but Figure 8 (right) substitutes Gemma2 9B with DeepSeekMath 7B. What motivated this inconsistency in model selection?\n\n\nDifferent authors contributed to different parts of the paper’s experiments. Due to a lack of communication at the time, we ended up choosing different models to use for our experiments. \n\n> In the data curation experiments, only levels 1–3 of the MATH dataset are used for training. Why was only a subset of MATH employed instead of the full dataset?\n\nMATH is hard for models at the scale we are looking at, so most examples end up being in the “hard” regime, and all data curation approaches end up selecting from the same data distribution. In order to highlight the difference between the different approaches, we focused on a subset of the dataset where a larger ratio of examples are in not the “hard” regime. \n\n> The process for choosing the perplexity threshold p lacks explanation and analysis...\nThe explanation of the data curation process lacks detail. \n\nWe provide a detailed explanation of the process we use to select p in Appendix A, and the process for data curation in Appendix D. Please let us know if there are any additional questions you have regarding the processes.", "claims_and_evidence": "See Other Strengths And Weaknesses and Questions", "methods_and_evaluation_criteria": "See Other Strengths And Weaknesses and Questions", "theoretical_claims": "There isn’t any theoretical analysis.", "experimental_designs_or_analyses": "See Other Strengths And Weaknesses and Questions", "relation_to_broader_scientific_literature": "This paper advances the understanding of LLM generalization in reasoning tasks by introducing ​pre-memorization train accuracy—a metric capturing model accuracy before memorizing training solution traces—and bridges gaps across multiple research areas.", "essential_references_not_discussed": "No", "other_strengths_and_weaknesses": "Strengths: \n\n1. The introduction of pre-memorization training accuracy is a fresh idea. It focuses on model performance before memorizing exact solutions, which hasn't been explored much before.\n\n2. The paper shows high $R^2$ values (≥0.9) across different models and datasets, which is impressive. \n\n3. The data curation method using pre-memorization accuracy improves sample efficiency, which is useful for real-world training scenarios.\n\nWeakness: \n\n1. The authors need ​a deeper analysis and discussion​ of pre-memorization accuracy. The robustness and data curation sections lack sufficient insight beyond stating the relationship between pre-memorization accuracy and test accuracy, especially since test samples could be interpreted as a form of perturbation on the training dataset.\n\n2. The process for choosing the perplexity threshold p lacks explanation and analysis. The reasons why p is strongly dependent on the task/model remain unclear.\n\n3. The experiments are restricted to mathematical datasets. Additional experiments with ​diverse datasets​ (e.g., logical reasoning or commonsense QA) are necessary to validate generalizability.\n\n4. The perturbations used (like adding \"First\") are minor and may not test true robustness against more complex or semantic changes.\n\n5. The explanation of the data curation process lacks detail. The rationale for hyperparameter choices (e.g., thresholds and iterations) is ​insufficiently explained or discussed.\n\n6. The authors categorize training samples into two types based on learning dynamics but fail to explain ​the underlying reasons​ why some samples exhibit one training pattern while others show another.", "other_comments_or_suggestions": "See Other Strengths And Weaknesses and Questions", "questions_for_authors": "1. The authors categorize samples with high perplexity (ppl) and high accuracy as well-generalized. However, it is possible the model simply memorizes the final answer rather than learning reasoning steps. Is there any analysis or experiment addressing this concern?\n\n2. The perplexity threshold p is selected by sweeping values, but the chosen p seems disproportionately large relative to the task/model. Is there an explanation for why such high thresholds are necessary?\n\n3. In Figure 1 (left), the curve for the learning rate 5e−7 closely aligns with y=x, suggesting a strong train-test accuracy correlation. Is pre-memorization accuracy effectively equivalent to accuracy under very small learning rates since it is measured early in training?\n\n4. Figure 4 uses Llama3 8B and Gemma2 9B, but Figure 8 (right) substitutes Gemma2 9B with DeepSeekMath 7B. What motivated this inconsistency in model selection?\n\n5. In the data curation experiments, only levels 1–3 of the MATH dataset are used for training. Why was only a subset of MATH employed instead of the full dataset?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Uyie3CxAiz", "reviewer": "Reviewer_ZcdX", "summary": "The paper studies how LLMs generalize in reasoning tasks and introduces \"pre-memorization train accuracy\" as a metric that predicts test accuracy. The key idea is that models first learn correct reasoning patterns before they start memorizing training examples, and this early accuracy is strongly correlated with final test performance. The authors show that this metric is predictive across different models and datasets. They also propose using this metric for data selection. The results suggest that tracking model performance before memorization kicks in can be useful for better training strategies.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the nice feedback! We answer the questions below:\n\n> How does pre-memorization accuracy behave in one-epoch fine-tuning? Many LLMs are fine-tuned in a single pass. Does the metric still hold if computed over early training steps?\n\nReasoning data tend to be rare, so it is common in prior work to do multiple passes on these datasets. For instance, [1] does some analysis on model performance and number of epochs of training, and found that pass@1 performance on the test set continues to increase for 50 epochs\n\nOur work also does analysis on both the one-epoch finetuning setting, as well as early-epochs in longer training runs (Fig. 4). We found that pre-memorization train accuracy is still highly predictive of test accuracy in these settings. \n\n[1] Training Verifiers to Solve Math Word Problems \n\n> Could this work for non-math reasoning tasks?\n\nWe believe our findings are likely to hold in problem settings where the target label includes both chain-of-thought reasoning (not uniquely correct) and a final answer (uniquely correct and unlikely to be guessed correctly)\n\n> Does this method work when models are weaker?\n\nWe have not tested. Weaker models (e.g. GPT 2) tend to perform very poorly on reasoning tasks, which is likely to make the analysis largely uninformative\n\n> Could pre-memorization accuracy be used to predict how well a model generalizes OOD?\n\nWe have not tested our work in the OOD setting, though we believe this would be an interesting direction for future study!", "claims_and_evidence": "The claims are mostly well supported by experiments. The paper shows strong correlations between pre-memorization accuracy and test accuracy, and the data curation experiments give good evidence that this metric can be useful for improving sample efficiency. One possible concern is that models in real-world fine-tuning setups often only train for one epoch, meaning they may not actually go through a clear pre-memorization and memorization phase. So it's not clear how well this metric would apply outside of multi-epoch settings.", "methods_and_evaluation_criteria": "The methods seem reasonable for studying generalization in reasoning tasks. The authors use standard reasoning datasets GSM8K and MATH which are a good fit since they allow distinguishing between correct reasoning and memorization. The main metric (pre-memorization train accuracy) is well defined. The evaluation is thorough, with different models, learning rates, and dataset sizes. That said, since most practical LLM fine-tuning is done in a one-pass setting, it might be more useful to measure pre-memorization accuracy over early batches instead of epochs.", "theoretical_claims": "There aren't any formal theoretical results in the paper.", "experimental_designs_or_analyses": "The experiments are generally sound, with strong empirical results. The authors systematically vary learning rates, dataset sizes, and models to validate their claims. However, the multi-epoch training setup is a bit unrealistic. Most LLM fine-tuning is done in a single pass, so it’s unclear if the metric holds in that scenario.", "relation_to_broader_scientific_literature": "The paper is closely related to research on memorization and generalization in deep learning. It connects well with past work on leave-one-out memorization metrics (Feldman & Zhang, 2020) and discussions on whether memorization helps or hurts generalization.", "essential_references_not_discussed": "Nothing that I am aware of.", "other_strengths_and_weaknesses": "Strengths:\n\n- Novel and well-motivated metric for understanding generalization.\n- Strong empirical results with high correlation to test accuracy.\n- Practical applications for data curation and improving sample efficiency.\n\nWeaknesses:\n\n- Limited model and dataset diversity (only two models, only math reasoning).\n- No theoretical justification for why pre-memorization accuracy is predictive.\n- The multi-epoch training assumption might limit applicability to real-world LLM fine-tuning.", "other_comments_or_suggestions": "- Would be interesting  to test if pre-memorization accuracy helps detect spurious correlations in training data.\n- The p-selection process still raised concerns about leakage. Maybe an alternative selection method would be more robust.\n- Typo: \"acheived\" -> \"achieved\"", "questions_for_authors": "1. How does pre-memorization accuracy behave in one-epoch fine-tuning? Many LLMs are fine-tuned in a single pass. Does the metric still hold if computed over early training steps?\n2. Could this work for non-math reasoning tasks?\n3. Does this method work when models are weaker?\n4. Could pre-memorization accuracy be used to predict how well a model generalizes OOD?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "1vLkGjTeOB", "reviewer": "Reviewer_833t", "summary": "This paper focuses on LLM reasoning tasks and proposes a metric called **per-memorization train accuracy**, which is the accuracy of model samples on training queries before they begin to copy the exact reasoning steps from the training set. The authors show that the proposed metric is predictive of the **test accuracy in dataset level**, as well as **model robustness to perturbation at a sample level**. They further use this metric to improve sample efficiency by prioritize training on samples with low pre-memorization accuracy.\n\n## update after rebuttal\nI thank the authors for the clarifications. I increased my score to 3 (provided that they make the title and scope of their claims clearer as math reasoning)", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the feedback! In the following rebuttal, we first provide a detailed description of how our metric is calculated and used in practice. Next, we address the concerns about the scope of our experiments. Finally, we provide clarifications for our experimental setup. We will update the final version of our submission to include these additional details. \n\n> Weakness: my concern is on the actionable point and implication of this paper. In Appendix A (Figure 9), the correlation between the proposed metric and test accuracy varies a lot with differnet threshold p, and the optmal p  is different for different datasets. For an unseen dataset, how would you predict test accuracy without calibrating the p w.r.t. true test accuracy? Can you clarify how the metric can be used in practice?\n\nWe would like to clarify that the main goal of our work is not to predict test accuracy using the model’s training behavior, since test accuracy can be measured relatively easily with a holdout dataset. Instead our goal is to 1) better understand aspects of a model’s learning dynamics which govern their downstream generalization, and 2) use our findings to improve model generalization. \n\nTowards goal 1, we analyze the model’s behavior through training, and find that models tend to have 2 modes of optimizing the learning objective: one which leads to downstream generalization and one which doesn’t (learning progression 1 vs 2 in Fig. 2). Train accuracy conflates the two models of learning. In contrast, pre-memorization train accuracy allows us to delineate between the two, and identify the examples that are learned in a “generalizable” manner and the examples that are not.\n\nTowards goal 2, we study how pre-memorization train accuracy can be used to improve data curation. While test accuracy measures a model’s generalization in aggregate, pre-memorization train accuracy can measure generalization for individual train examples. This allows us to determine which train examples are “hard” for the model, which can better inform the kinds of examples to collect that would most improve model performance. \n\nThe exact workflow for our data curation experiments is as follows. We assume access to a dataset split into a train/test set, using which we can calibrate p with a few (<3) training runs using test accuracy. We measure the pre-memorization train accuracy of the train examples. Using this information, we curate new data that is drawn from the same distribution as training examples with low pre-memorization train accuracy in the current iteration. \n\n> Evidence 1: while it is always annoying to be asked for results on more datasets, I think it is needed here. The paper is purely empirical and there are quite a few other datasets (besides the 2 math-only datasets) that authors could have evaluated to further strengthen their conclusion.\n\nThe experiments in our paper include more than 60 training runs across different models, tasks, dataset scale, and training settings. While we focused on two mathematical datasets, we believe the depth and comprehensiveness of our analysis is comparable if not exceeds the standards for empirical conference papers.\n\nWe unfortunately do not currently have the compute to perform experiments on additional datasets. However, we believe this should not be the reason for rejection, as many reasoning papers have been accepted at similar-tier conferences with comparable or even more limited dataset coverage. For instance:\n\n[1] RL on Incorrect Synthetic Data Scales the Efficiency of LLM Math Reasoning by Eight-Fold (Neurips) - Experiments: GSM8k, MATH\n\n[2] Scaling LLM Test-Time Compute Optimally Can Be More Effective Than Scaling Parameters for Reasoning (ICLR) - Experiments: MATH\n\n> Method 1: If I understand correctly, the pre-memorization accuracy for a single sample can only be either 0 or 1 (correct or wrong)? If so, why the accuracy in Figure 7 is continuous?\n\nWe take multiple samples of outputs for every input, and calculate the average accuracy\n\n> Experimental design 1: In terms of writing, I think the paper lacks necessary details here and there. For example,\n> The experimental setup of Section 4.1 is not mentioned, the models used in Figure 3 are not mentioned)\n\nThe models are trained using llama 2 8B on the full dataset with 6 epochs, and peak learning rates as described in the figure\n\n> Similar for Figure 6, the model used to train is not mentioned.\n\nWe used llama2 8b\n\n>In section 5.2 (L369 right), how and which GPT is used to rephrase examples in the original dataset? The author mentioned that they followed a similar procedure in Setlur et al. 2024, but it is important to briefly summarize (or give examples) how it is done.\n\nThe data used in our work was generated using the exact procedure as described in Appendix E of Setlur et al. 2024, which details the model, system prompt, and approximate cost. We are happy to add or clarify any more details needed.", "claims_and_evidence": "**Evidence 1**: while it is always annoying to be asked for results on more datasets, I think it is needed here. The paper is purely empirical and there are quite a few other datasets (besides the 2 math-only datasets) that authors could have evaluated to further strengthen their conclusion.", "methods_and_evaluation_criteria": "Most of the evaluation criteria makes sense to me, and I have a following question:\n\n**Method 1**: If I understand correctly, the pre-memorization accuracy for a single sample can only be either 0 or 1 (correct or wrong)? If so, why the accuracy in Figure 7 is continuous?", "theoretical_claims": "The paper is mostly empirical and there are no theoretical results. However, this is not considered as a weakness.", "experimental_designs_or_analyses": "**Experimental design 1**: In terms of writing, I think the paper lacks necessary details here and there. For example,\n\n\n- The experimental setup of Section 4.1 is not mentioned, the models used in Figure 3 are not mentioned)\n- Similar for Figure 6, the model used to train is not mentioned.\n- In section 5.2 (L369 right), how and which GPT is used to rephrase examples in the original dataset? The author mentioned that they followed a similar procedure in Setlur et al. 2024, but it is important to briefly summarize (or give examples) how it is done.\n\nOverall, adding these details will help the reader in developing a more concrete idea and therefore make the reading much smoother.", "relation_to_broader_scientific_literature": "- **Broader literature 1**: the idea of relying on the correlation of a prosed metric and the test accuracy to predict test accuracy broadly relates to works trying to predict out-of-distribution (OOD) performance in a similar manner. The reason is that there is need to estimate how the model at hand will perform on target domain when labeled OOD target data is hard to obtain.\n- **Broader literature 2**: the idea of using the proposed metric to *expand* the training dataset broadly relates to synthetic data and data selection. For the latter, many works try to reduce the training dataset by removing noisy or redundant ones and only keep the informative ones, which is a concept called Coreset in the relevant literature.", "essential_references_not_discussed": "The paper reminds me of Miller et al. 2021, where they show in-distribution accuracy correlates with OOD accuracy in many cases. Maybe worth discussing in related work.\n\nReference:  \nAccuracy on the Line: On the Strong Correlation\nBetween Out-of-Distribution and In-Distribution Generalization. Miller et al. 2021", "other_strengths_and_weaknesses": "**Strength**: the paper is generally well-organized and the idea is easy to understand. Understanding memorization and generalization is timely and important, and the proposed Pre-Memorization Train Accuracy shows storng correlation with the test accuracy on unseen samples.\n\n**Weakness**: my concern is on the actionable point and implication of this paper. \n- In Appendix A (Figure 9), the correlation between the proposed metric and test accuracy varies a lot with differnet threshold $p$, and the optmal $p$ is different for different datasets. For an unseen dataset, how would you predict test accuracy without calibrating the $p$ w.r.t. true test accuracy? Can you clarify how the metric can be used in practice?", "other_comments_or_suggestions": "Minor: L298 right side, it should be Fig. 7 instead of 6?", "questions_for_authors": "I would be happy to increase my score if the authors could address my questions in **Weakness**, **Evidence 1**, **Method 1**, **Experimental design 1**, especially the first two.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "kfg2Z0PiAE", "reviewer": "Reviewer_cGoj", "summary": "The paper introduces a novel metric called \"pre-memorization train accuracy\" that predicts how well large language models will generalize during reasoning tasks. The key insight is examining model performance on training examples before they are memorized verbatim. The authors show this metric strongly correlates with test accuracy (R² ≥ 0.9) across different models, datasets, and training configurations. They leverage this finding to improve data curation, achieving 1.5-2x better data efficiency compared to standard approaches.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the nice feedback! We answer the questions below. \n\n> How sensitive is the pre-memorization accuracy metric to the choice of memorization threshold? \n\nIn Fig, 9, we provide some some analysis on the sensitivity of the predictive power (R^2) of our metric to the choice of the memorization threshold. We found that R^2 degrades smoothly with respect to p, which makes it is relatively easy to find a good value of p by sweeping a range of values.\n\n> Does the correlation with test accuracy hold for much larger models (>70B parameters)? \n\nUnfortunately, we do not have the compute to test this right now, though we believe this would be an interesting direction for future work\n\n> How does the metric behave during few-shot learning versus full fine-tuning?\n\nOur metric cannot be directly applied in the few-shot learning setting. However, we believe understanding the relationship between the few-shot data and downstream generalization would be a very interesting direction for future study.", "claims_and_evidence": "From my experience, it would be better to add more ablation studies testing different definitions of memorization beyond perplexity thresholds.", "methods_and_evaluation_criteria": "The methodology is sound, with careful experimental design. The authors evaluate their metric through multiple lenses: aggregate test performance prediction, per-example robustness analysis, and practical data curation improvements. The evaluation on both mathematical (GSM8k) and general reasoning tasks demonstrates broad applicability.", "theoretical_claims": "I have verified the proofs in Section 2 and Appendix E. The formalization of language-thought modeling gap and the derivations of Propositions 2.4 and 2.7 are mathematically rigorous. The assumptions and conditions are clearly stated.", "experimental_designs_or_analyses": "The experiments are thorough and well-designed. The authors control for multiple variables like learning rate, dataset size, and model architecture. From my experience, it would be better to add experiments testing the metric's sensitivity to different random seeds.", "relation_to_broader_scientific_literature": "The paper effectively positions itself within relevant literature on memorization, generalization, and LLM reasoning. It builds upon and advances prior work on measuring memorization (Feldman & Zhang 2020) and understanding Chain-of-Thought reasoning (Wei et al. 2022).", "essential_references_not_discussed": "In my view, it would be better to add recent works on emergent abilities in LLMs and how they relate to memorization versus generalization. The discussion would benefit from incorporating recent papers on phase transitions in LLM learning.", "other_strengths_and_weaknesses": "A key strength is the practical utility of the findings - the metric provides actionable insights for improving training efficiency. One weakness is the reliance on perplexity thresholds for defining memorization, which may not capture all forms of memorization.", "other_comments_or_suggestions": "NA", "questions_for_authors": "1) How sensitive is the pre-memorization accuracy metric to the choice of memorization threshold? 2) Does the correlation with test accuracy hold for much larger models (>70B parameters)? 3) How does the metric behave during few-shot learning versus full fine-tuning?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: The paper proposes a pre-memorization train accuracy metric that measures model performance on training examples before memorization occurs, which is shown to correlate strongly with test accuracy.
Several experiments across different models (Llama3 and Gemma2) and math reasoning datasets (GSM8k, MATH) backup the claims.

Overall, the reviewers appreciated the novel perspective on analyzing LLM learning dynamics, the comprehensive experimental setup, and the practical applicability of the data curation technique.

Some concerns were raised, for instance the limited dataset diversity, with a focus mainly on math reasoning tasks, raising questions about generalizability to other reasoning domains. That said, I think math reasoning datasets are very popular and I think this is a good choice. I would encourage the authors to consider another dataset to demonstrate further generalization. I would also consider adding more in-depth discussion on the selection and impact of the perplexity threshold (p) and the robustness of the metric under varying training conditions (e.g., one-epoch fine-tuning, few-shot learning).

I recommend acceptance but strongly recommend the authors make the changes required by the reviewers to broaden the impact of the paper.
Decision: Accept (poster)