From Memorization to Reasoning in the Spectrum of Loss Curvature

Abstract: We characterize how memorization is represented in transformer models and show that it can be disentangled in the weights of both LMs and vision transformers (ViTs) using a decomposition based on the loss landscape curvature. This insight is based on prior theoretical and empirical work showing that the curvature for memorized training points is much sharper than non memorized, meaning ordering weight components from high to low curvature can reveal a distinction without explicit labels. This motivates a weight editing procedure that suppresses far more recitation of untargeted memorized data more effectively than a recent unlearning method (BalancedSubnet), while maintaining lower perplexity. Since the basis of curvature has a natural interpretation for shared structure in model weights, we analyze the editing procedure extensively on its effect on downstream tasks in LMs, and find that fact retrieval and arithmetic are specifically and consistently negatively affected, even though open book fact retrieval and general logical reasoning is conserved. We posit these tasks rely heavily on specialized directions in weight space rather than general purpose mechanisms, regardless of whether those individual datapoints are memorized. We support this by showing a correspondence between task data's activation strength with low curvature components that we edit out, and the drop in task performance after the edit. Our work enhances the understanding of memorization in neural networks with practical applications towards removing it, and provides evidence for idiosyncratic, narrowly-used structures involved in solving tasks like math and fact retrieval.

• The paper demonstrates that mapping loss landscape curvature reveals memorization is linked to sharp weight directions, while reasoning benefits from flatter regions.
• It introduces a weight-editing technique that selectively ablates memorization-related components, outperforming methods like BSN in preserving generalization.
• Experimental results show that tasks such as logical reasoning remain robust despite curvature perturbations, whereas arithmetic and fact retrieval are more sensitive.

From Memorization to Reasoning in the Spectrum of Loss Curvature

Introduction

The paper "From Memorization to Reasoning in the Spectrum of Loss Curvature" explores the interplay between memorization and reasoning capabilities in neural networks, specifically Transformer models such as LLMs (LMs) and Vision Transformers (ViTs). It investigates how the curvature of the loss landscape, characterized via the Kronecker-Factored Approximate Curvature (K-FAC), can be leveraged to disentangle these cognitive processes and contribute to model editing methodologies for suppressing memorization while preserving reasoning capabilities.

Disentangling Memorization via Loss Curvature

The authors present an approach where K-FAC is used to approximate the Hessian's eigenbasis to identify directions of high and low curvature in model weight space. This decomposition reveals that memorization is associated with sharper directions for specific instances, while generalization aligns with flatter directions across datasets. This insight is operationalized in a weight-editing technique that suppresses memorization by retaining high-curvature components over multiple data points.

Figure 1: Overview of the approach with activations and gradients leading to a K-FAC-based decomposition revealing task-specific interactions.

Practical Model Editing to Suppress Memorization

The study introduces a model weight-editing approach that selectively ablates components within the curvature spectrum to suppress memorization. By retaining only specific portions of the eigenbasis deemed essential for generalization, the authors demonstrate that model coherence is maintained while minimizing memory-based recitation. This technique is shown to outperform existing methods like Balanced Subnet (BSN), especially in terms of generalization to unseen data without requiring a supervised forget set.

Figure 2: Disentanglement in weight space between memorized and non-memorized data using K-FAC, showing distinct activation patterns.

Impact on Various Task Performances

The framework was applied to assess the sensitivity of different cognitive tasks to weight perturbations within this eigenbasis framework. The results illustrate a spectrum where arithmetic and closed-book fact retrieval are more brittle, suggesting reliance on low-curvature directions, while logical reasoning and open-book retrieval tasks remain robust.

Figure 3: Sensitivity analysis of task performances with respect to ablation of curvature-based weight directions.

Implications and Future Directions

The findings provide a deeper understanding of how neural networks encode memorization versus reasoning by linking these processes to weight space curvature. This understanding carries significant implications for model training and editing, offering pathways toward creating models that better generalize while being less prone to verbatim data recitation. Future work could extend this research to fine-tune models' reasoning capabilities further and explore the balance between memorization and reasoning, potentially leading to more efficient training paradigms.

Conclusion

By leveraging the spectrum of loss curvature, this study offers a robust framework for understanding and manipulating the memorization and reasoning dichotomy in neural models. This has practical applications in creating more resilient AI systems and opens up avenues for further exploration into improving AI's reasoning faculties through weight space analysis.

Figure 4: Eigenvector activation ratios indicate task interactions with curvature bands, reflecting task-specific processing differences.