What Do Learning Dynamics Reveal About Generalization in LLM Reasoning? (2411.07681v2)

Abstract: Despite the remarkable capabilities of modern LLMs, the mechanisms behind their problem-solving abilities remain elusive. In this work, we aim 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 generalization behavior 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 reliably predict test accuracy, achieving $R^2$ of around or exceeding 0.9 across various models (Llama3 8, 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 behavior to its generalization, pre-memorization train accuracy can guide targeted improvements to training strategies. We focus on data curation as an example, and 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 outperforms other standard data curation techniques.

• The paper introduces pre-memorization train accuracy, which strongly correlates (R² > 0.9) with test accuracy across various model and dataset configurations.
• It demonstrates that learning rates significantly affect the transition from generating diverse correct outputs to rote memorization, impacting overall generalization.
• The study highlights that emphasizing training examples with low pre-memorization accuracy can enhance data efficiency by 1.5-2x on reasoning tasks.

Insights into Learning Dynamics and Generalization in LLM Reasoning

The paper, "What Do Learning Dynamics Reveal About Generalization in LLM Reasoning?" by Kang et al., investigates the learning mechanics of LLMs specifically in the domain of reasoning, shedding light on the nuances of generalization beyond mere memorization. The authors delve into the training dynamics of LLMs, offering a novel perspective on how pre-memorization accuracy, a key metric introduced in this paper, predicts generalization to unseen data.

Key Contributions and Findings

The paper's primary contribution is the concept of pre-memorization train accuracy. This metric measures a model's accuracy on training examples before it memorizes the target reasoning paths, aligning memorization with low perplexity in reproducing training solution traces. The high correlation between pre-memorization accuracy and test accuracy, with coefficients of determination exceeding 0.9 across various model-dataset configurations, underscores its predictive power. This finding effectively challenges prior assumptions that memorization alone undermines generalization capabilities.

The authors explore how learning rates influence the transition from generating diverse correct outputs to rote memorization. Their empirical analysis demonstrates that while model predictions were often near-perfect on the training set, the generalization abilities varied significantly, underscoring an intricate relationship between learning intricacies and downstream performance. The pre-memorization accuracy serves as a robust predictor across architecture sizes (Llama3 8B, Gemma2 9B), tasks (GSM8k, MATH), and different hyperparameter settings.

Interestingly, the research uncovers that examples exhibiting high pre-memorization train accuracy tend to yield more robust model behaviors against diverse input perturbations, suggesting an internalized, flexible reasoning mechanism rather than a brittle, memorized one.

Practical Implications and Data Curation

From a practical standpoint, the insights into pre-memorization train accuracy offer new avenues for data curation strategies. The paper's empirical evidence highlights that emphasizing training examples with low pre-memorization accuracy enhances data efficiency over traditional i.i.d. sampling methods, delivering significant gains in sample efficiency (1.5-2x improvements) on reasoning tasks. This strategy represents a departure from common curation methods, which often rely on heuristics or blanket difficulty estimates.

By identifying examples that do not readily align with memorized reasoning paths as candidates for focused training, data curation could be better tailored to enhance model robustness and generalization, offering exciting prospects for improving the training regimes for LLMs.

Theoretical and Future Directions

Theoretically, this work provides a framework to investigate how the dynamics of learning can be leveraged to improve generalization. The paper challenges the established narrative by demonstrating that memorization, when dissected with precision, reveals dimensions of learning efficacy linked to reasoning capabilities. Highlighting the connection between learning dynamics and generalization provides fertile ground for further research, potentially informing the optimization of training processes and model architecture designs.

Future directions could explore adaptive learning rates and curriculum learning strategies that refine the identification and prioritization of complex training examples. Additionally, expanding the scope to broader reasoning tasks and more diverse datasets might shed further light on the scalability and universality of pre-memorization train accuracy as a metric.

In conclusion, Kang et al.'s work advances the understanding of LLM learning dynamics, presenting a compelling case for nuanced data-driven approaches in the ongoing evolution of intelligent LLMs. As the AI community continues to unravel the black box of LLM training, these insights pave the way for more sophisticated, efficient, and generalizable learning systems.