Towards Understanding Multi-Task Learning (Generalization) of LLMs via Detecting and Exploring Task-Specific Neurons

TL;DR

This work examines multi-task generalization in large language models from a neuron-centric perspective, identifying task-specific neurons in FFN modules via gradient attribution and validating their role with deactivation and fine-tuning experiments. It reveals that the overlap and parameter similarity of these neurons across tasks are strongly linked to generalization, and that transfer tends to emerge in deeper layers when knowledge is shared. Building on these insights, the authors introduce neuron-level continual fine-tuning (NCFT) and a weighted variant (W-NCFT) to mitigate catastrophic forgetting without adding new parameters. Extensive experiments on Llama-2-7b and Bloom-7b1 demonstrate meaningful improvements in continual learning and offer interpretable evidence for how multi-task learning is organized inside LLMs. The study provides practical strategies for improving generalization across tasks while preserving previously learned knowledge, contributing to the interpretability and reliability of multi-task LLMs.

Abstract

While large language models (LLMs) have demonstrated superior multi-task capabilities, understanding the learning mechanisms behind this is still a challenging problem. In this paper, we attempt to understand such mechanisms from the perspective of neurons. Specifically, we detect task-sensitive neurons in LLMs via gradient attribution on task-specific data. Through extensive deactivation and fine-tuning experiments, we demonstrate that the detected neurons are highly correlated with the given task, which we term as task-specific neurons. With these identified task-specific neurons, we delve into two common problems in multi-task learning and continuous learning: Generalization and Catastrophic Forgetting. We find that the overlap of task-specific neurons is strongly associated with generalization and specialization across tasks. Interestingly, at certain layers of LLMs, there is a high similarity in the parameters of different task-specific neurons, and such similarity is highly correlated with the generalization performance. Inspired by these findings, we propose a neuron-level continuous fine-tuning method that only fine-tunes the current task-specific neurons during continuous learning, and extensive experiments demonstrate the effectiveness of the proposed method. Our study provides insights into the interpretability of LLMs in multi-task learning.

Figures (6)

• Figure 1: Illustration of our research methodology. The entire framework consists of three components: Identification (task-specific neurons), Understanding (multi-task learning mechanisms of LLMs from the neuron level) and Exploration (neuron-level continuous fine-tuning method). The first component provides tools for mechanism understanding which in turn provides insights for the third component Exploration.
• Figure 2: Results on classification and generation tasks after fine-tuning different proportions of task-specific neurons. The red line indicates the average of the results on the in-domain (ID) test set and out-of-domain (OOD) test set with the same type of training task. For example, in subfigure (a), the red line shows the average of the blue and orange lines while the average of the purple and green lines in subfigure (b). The correspondence of the other colored lines to the test set is shown in the caption of Table \ref{['tab_3']}.
• Figure 3: The similarity of the task-specific neuron parameters between the test task and training tasks in different layers.
• Figure 4: Forgetting rates for eight stages on the Large Number of Tasks benchmark.
• Figure 5: Performance of Llama-2-7b on all tasks after deactivation or fine-tuning a particular class task-specific neurons. The element in the $i$-th row and $j$-th column is the performance change for task $j$ due to deactivation or fine-tuning of the task $i$ specific neurons.