When Llama Learns Math …

Continual learning in Large Language Models.


“So it is truly amazing that by pure thought, without doing any experiments whatsoever, one can conclude that the universe must have had a beginning.” Michio Kaku


In previous articles, we explored the evolution of continual learning as a field and considered some open questions for the current era. In this article, we’ll analyse the impact of continual learning in small-scale Llama/Gemma models through a simple experiment, create a hypothesis to explain our findings and discuss the takeaways for fine-tuning Large Language Models (LLMs) in practice.

Setup

First things first – let’s define the goal of the experiments and the setup being used. We start with a pre-trained checkpoint for an LLM and fine-tune on a dataset of target ability – in this article, we’ll consider the target ability as grade-school level math. Our goal is to understand how well the model can learn this new target ability in a continual manner – i.e. without catastrophic forgetting (CF) (i.e. loss) of pre-trained abilities. Given the wide and abstract nature of capabilities gained from pre-training, defining and tracking the performance of “pre-trained ability” as a whole is of course up for debate. For now, I’ve chosen grade-school level science, entry-level Python programming and general knowledge from among commonly used tasks for reporting pre-trained model performance.

LLM checkpoints: (1) Llama-3.2-1B (2) Gemma-3-1B

Datasets: (1) MetaMathQA for target ability fine-tuning (2) GSM8K for evaluating performance in math (3) ARC-E for evaluating performance in science (4) MBPP for evaluating performance in Python programming (5) MMLU for evaluating performance in general knowledge

Training: Standard cross-entropy fine-tuning with a learning rate of 1e-5 and an effective batch size of 128 (=4×32 gradient accumulation steps) for 2 epochs.

Code and Evals: Fine-tuning and evaluation code is taken from here.

Results

The table below depicts the percentage accuracy on the test split of each of the chosen tasks, of the pre-trained and fine-tuned model checkpoints of both LLMs.

LLM
Checkpoint
GSM8K
(Target task)
MBPP
MMLU
ARC-E
Llama-3.2-1BPre-trained6.0726.6038.5866.29
Llama-3.2-1BFine-tuned42.0018.8035.4565.28
Gemma-3-1BPre-trained2.509.0025.4671.13
Gemma-3-1BFine-tuned31.086.0025.3569.49
Percentage accuracy of pre-trained and fine-tuned LLM checkpoints on various tasks.

Now, for some interpretation:

(1) The pre-trained models are quite bad at math, as measured on the GSM8K task, but fine-tuning on MetaMathQA improves their performance on GSM8K quite a bit. This is as expected.

(2) The improvement in mathematical abilities is not free and comes at the cost of some forgetting of Python abilities, as measured on the MBPP task. Again, this is expected.

(3) On MMLU and ARC-E, there is much lesser forgetting (<3% and <2%, for Llama and Gemma, respectively) after the fine-tuning than on MBPP (8% and 3%, for Llama and Gemma, respectively). This is surprising, since intuitively, we don’t expect mathematical abilities to be aligned to either general knowledge or scientific reasoning, any more than to Python programming.

Why does learning math lead to forgetting of Python abilities but not as much of science and general knowledge?

Code versus English: One possible answer to this question is that the MBPP task involves understanding a coding question framed in English (such as “write a function for…”) and generating a Python code snippet in the response. Fine-tuning on MetaMathQA, however, involves understanding a mathematical problem description in English and generating complete English sentences when solving the problem. This might partially degrade the ability to generate Python. Since the MMLU and ARC-E tasks also require the generation of complete English sentences, the model is able to better retain its performance on these tasks.

Testing the hypothesis

If it is indeed a matter of code vs English, how can we avoid this kind of forgetting when fine-tuning?

(1) If we knew that Python programming was an ability we want the fine-tuned model to be good at, and we had access to the Python pre-training dataset, the classic replay approach of mixing in some Python examples during the fine-tuning stage should help.

(2) What if we didn’t have access to the Python pre-training data? Alternatives are to create a proxy task – such as by using the pre-trained/different LLM for synthetic data generation – or, to re-frame the mathematical reasoning task to also generate supporting Python implementation of the solution.

(3) What if we didn’t know which pre-trained abilities to test for when fine-tuning on our downstream task? In that case, testing on a wide array of tasks would provide an idea of the type of tasks that seem to be most affected by our downstream task. Subsequently, approach 1 (replay) or approach 2 (synthetic replay) from above can be explored.

Comparing replay and synthetic replay strategies would be an interesting study of its own – crucially, it can tell us whether it really is important to have access to the pre-training dataset when fine-tuning models. Let me know if you try/have tried it and have any insights 🙂

On the other hand, what if our hypothesis is wrong, i.e. the forgetting of Python abilities is not due to an underlying code vs English phenomenon, but rather just an artefact of the fine-tuning setup? Perhaps training for another epoch might cause just as much of a drop in performance in the other pre-trained abilities – how do we tackle that? While more sophisticated approaches than replay do exist for mitigating forgetting (e.g. gradient alignment), it is important to note that it may well be that our target task cannot be learnt without sacrificing some pre-trained abilities. In that case, we need to take an informed decision on whether the amount of forgetting observed is acceptable for our use case.

Takeaways

(1) When fine-tuning pre-trained LLMs, it’s important consider what kinds of pre-trained abilities we still want to retain and focus on evaluating performance on these/similar tasks before and after fine-tuning.

(2) Lack of access to pre-training data can necessitate innovative approaches to handle forgetting during fine-tuning.

(3) The underlying cause of forgetting is not guaranteed to be intuitive, just as forgetting itself is not guaranteed to be concentrated on one/few abilities.

One caveat of our experiment is that the results may be specific to these particular models, model size or fine-tuning setup. E.g. Mixture-Of-Experts models can be less prone to forgetting since the gating mechanism encourages specialisation and separation of parameters (though gating also makes these networks harder to train in general). Nevertheless, the general approach when fine-tuning for a task should be to track performance on a wide array of tasks and LLMs and then take action to handle any forgetting.