MathNeuro

Paper

Codebase for Math Neurosurgery: Isolating Language Models' Math Reasoning Abilities Using Only Forward Passes

Overview

Math reasoning is an active area of Large Language Model (LLM) research because it is a hallmark of artificial intelligence and has implications in several domains, including math education. However, few works have explored how math reasoning is encoded within LLM parameters and if it is a skill that can be isolated within models. 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 computationally efficient method we use to isolate 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 filtering out those important for general language tasks. Through pruning parameters MathNeuro identifies, we delete a LLM's math reasoning ability without significantly impacting its general language ability. Scaling the identified parameters by a small constant improves a pretrained or instruction-tuned LLM's performance by 4-17% on GSM8K and 5-35% on MATH 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.

License and Intended Use

Our code is released under the GNU GPLv3 license. Our codebase also contains a copy of the Eleuther AI Evaluation Harness to run all experimental evaluations.

Getting Started

After installing Python>=3.10 and PyTorch (follow instructions here), to install the dependencies for this codebase, you can run:

pip install -U -r requirements.txt

Running MathNeuro Pruning Experiments

To run MathNeuro for identifying and pruning math-specific parameters, you will need to run the MathNeuro.py file. Here is an example of how to conduct the GSM8K pruning experiment for Llama 3.2 1B IT:

python MathNeuro.py --model meta-llama/Llama-3.2-1B-Instruct --save_path /results_path --train_dataset data/gsm8k.csv --eval_datasets race mmlu --calibration_datasets data/race.csv data/mmlu.csv --eval_dataset_subset 200 --calibration_dataset_names Race MMLU --train_lm_eval_task gsm8k_cot --pre_train_eval

Running MathNeuro Scaling Experiments

To run MathNeuro for identifying and scaling math-specific parameters, you will need to run the MathNeuro.py file specifying your desired scalar. Here is an example of how to conduct the GSM8K scaling experiment for Llama 3.2 1B IT using a scalar of 1.1:

python MathNeuro.py --model meta-llama/Llama-3.2-1B-Instruct --save_path /results_path --train_dataset data/gsm8k.csv --eval_datasets race mmlu --calibration_datasets data/race.csv data/mmlu.csv --eval_dataset_subset 200 --calibration_dataset_names Race MMLU --train_lm_eval_task gsm8k_cot --pre_train_eval --scalar 1.1

Running MathNeuro Ablations

To run the ablations for MathNeuro (Wanda and random identification), you will need to run the MathNeuro_Ablations.py file. Here is an example of how to conduct the GSM8K pruning experiment ablations for Llama 3.2 1B IT:

python MathNeuro_Ablations.py --model meta-llama/Llama-3.2-1B-Instruct --save_path /results_path --train_dataset data/gsm8k.csv --eval_datasets race mmlu --calibration_datasets data/race.csv data/mmlu.csv --eval_dataset_subset 200 --calibration_dataset_names Race MMLU --train_lm_eval_task gsm8k_cot --pre_train_eval

Running LAPE

To run the LAPE comparison method, you will need to run the LAPE.py file. Please note that if you want to run this experiment for a model other than one from the Phi 1.5, Gemma 2, or Llama 3 families, you will need to create a customized implementation of its forward loop. Here is an example of how to conduct the GSM8K pruning experiment for Llama 3.2 1B IT using LAPE:

python LAPE.py --model meta-llama/Llama-3.2-1B-Instruct --save_path /results_path --train_dataset data/gsm8k.csv --eval_datasets race mmlu --calibration_datasets data/race.csv data/mmlu.csv --eval_dataset_subset 200 --calibration_dataset_names Race MMLU --train_lm_eval_task gsm8k_cot --pre_train_eval

Citation

@misc{christ2024mathneurosurgeryisolatinglanguage,
      title={Math Neurosurgery: Isolating Language Models' Math Reasoning Abilities Using Only Forward Passes}, 
      author={Bryan R. Christ and Zack Gottesman and Jonathan Kropko and Thomas Hartvigsen},
      year={2024},
      eprint={2410.16930},
      archivePrefix={arXiv},
      primaryClass={cs.CL},
      url={https://arxiv.org/abs/2410.16930}, 
}