Math Neurosurgery: Isolating Language Models' Math Reasoning Abilities Using Only Forward Passes
Bryan R. Christ, Zack Gottesman, Jonathan Kropko, Thomas Hartvigsen
2024-10-23

Summary
This paper introduces Math Neurosurgery (MathNeuro), a method that helps identify and isolate the math reasoning abilities of large language models (LLMs) without affecting their general language skills.
What's the problem?
Understanding how LLMs perform math reasoning is important, but researchers have not fully explored how these models encode math skills. Current methods do not effectively isolate math-specific abilities, making it difficult to improve their performance in math without changing their overall language capabilities.
What's the solution?
The authors developed MathNeuro, which uses a technique called forward passes to determine which parts of the model are important for math reasoning. By removing parameters that are not needed for math tasks while keeping those necessary for general language tasks, they can enhance the model's math performance by 4-17% on specific math tests, like GSM8K. This method is efficient and can work with just one example to identify important parameters.
Why it matters?
This research is significant because it provides a way to improve LLMs' math reasoning abilities without compromising their language skills. By isolating and enhancing these specific parameters, future AI models can become better at solving math problems, which is crucial for applications in education and other fields where accurate reasoning is necessary.
Abstract
Math reasoning is a highly active area of Large Language Model (LLM) research because it is a hallmark of artificial intelligence. However, few works have explored how math reasoning is encoded within LLM parameters and if it is a skill that can be isolated within a model. Doing so could allow targeted intervention to improve math performance without altering non-math behavior and foster understanding of how models encode math reasoning. We introduce Math Neurosurgery (MathNeuro), a method for isolating math-specific parameters in LLMs using only forward passes. MathNeuro builds on existing work by using weights and activations to calculate parameter importance, but isolates math-specific parameters by removing those important for general language tasks. Pruning parameters MathNeuro identifies deletes a LLM's math reasoning ability without destroying its general language ability. Scaling these parameters by a small constant improves a pretrained or instruction-tuned LLM's performance by 4-17% on GSM8K while leaving non-math behavior unaltered. MathNeuro is also data efficient: most of its effectiveness holds when identifying math-specific parameters using a single sample. MathNeuro highlights the potential for future work to intervene on math-specific parameters.