Source-linked AI summary
Functional Subnetwork Localization for Efficient Continual Learning in Large Language Models
Shuoran Jiang, Wentao Zhu, Yang Xiang, Dongfeng Li, Yunan Zhang, Mengchen Zhao, Yuhui Guo, Kai Chen, Youcheng Pan, Xiangping Wu, Qingcai Chen
TL;DR
Continual learning in LLMs must incorporate new information while retaining prior knowledge, and the paper investigates whether pretrained models contain cognition-associated subnetworks that support localized adaptation. It proposes CiMEO to route instructions and restrict updates to those regions. Across five benchmarks, CiMEO achieves the highest average performance among compared methods while balancing knowledge acquisition and retention.
Problem
Continual learning in LLMs must incorporate new knowledge while preserving previously acquired knowledge, but updates can overwrite earlier representations.
Source: §I · p. 1Method
CiMEO uses gradient attribution to identify low-overlap cognition-associated subnetworks, then routes each instruction to a relevant subnetwork and restricts parameter updates and optimizer states to selected regions.
Source: §I · p. 2Results
Across five benchmarks, CiMEO achieves the highest average performance among compared methods while balancing knowledge acquisition and retention.
Source: §V · p. 9Takeaways & Limitations
The findings highlight functional subnetwork localization as a promising approach to efficient and stable continual learning without extensive modifications to pretrained models.
Source: §V · p. 9
Abstract
from arXiv · showhide
Biological intelligence continually adapts to new information while retaining prior skills through localized plasticity. In contrast, continual learning in large language models (LLMs) remains vulnerable to catastrophic forgetting. A promising direction is to identify functionally specialized subnetworks within LLMs that enable targeted and resource-efficient adaptation. Using gradient-attribution profiling, we localize cognition-associated functional subnetworks exhibiting extremely low mutual overlap. Building on this finding, we propose Function-isolated Memory-Efficient Optimization (CiMEO), a continual learning framework that routes each training example to its corresponding functional subnetwork, restricting parameter updates and optimizer states to selected regions. Across 1.4 million sequential instructions, CiMEO achieves the highest average performance among the compared methods on five generative benchmarks spanning programming, examination-style tasks, and mathematical reasoning, while mitigating catastrophic forgetting. Subspace-overlap and principal-angle analyses further indicate that localized updates help preserve pre-trained representations, highlighting functional subnetwork localization as a practical approach to efficient and stable continual learning in LLMs.
I. INTRODUCTION
Continual learning in LLMs must add knowledge while preserving prior capabilities, but global updates can overwrite earlier representations and dense optimizer states make repeated adaptation costly. The paper identifies low-overlap cognition-associated subnetworks and proposes CiMEO to route updates and optimizer states to relevant regions.
I. INTRODUCTION
- Continual learning for LLMs must incorporate new knowledge while preserving previously acquired knowledge, yet global updates can overwrite earlier representations.The paper frames this as a stability–plasticity tension: models must acquire new behaviors while retaining old ones.Source: §I · p. 1
- Existing stability methods can add memory and computation costs: regularization stores importance estimates, rehearsal retains examples, and distillation requires a previous model and extra forward passes.Dense Adam-like optimizer moments further increase state memory during repeated posttraining.Source: §I · p. 1
- Biological learning is described as selectively recruiting task-relevant neuronal ensembles while other circuits remain comparatively stable, motivating localized rather than global adaptation.The figure contrasts this localized plasticity with dense global weight updates across sequential LLM training batches.Source: §I · pp. 1–2, Figure 1 · p. 2
- Gradient-attribution profiling identifies subnetworks across six Bloom-inspired cognitive categories, with average mutual overlap below 0.002%.The authors use this low overlap to motivate restricting updates to cognition-relevant regions.Source: §I · pp. 2–3
- CiMEO routes each incoming instruction to its corresponding functional subnetwork, updating selected parameters and maintaining optimizer states only for selected indices.The framework combines cognition-guided routing, functional subnetwork localization, and memory-efficient optimization without relying on replay buffers.Source: §I · p. 2
II. RELATED WORKS
Prior continual-learning work addresses interference through regularization, rehearsal, distillation, subspaces, parameter isolation, and functional localization. CiMEO differs by routing examples according to cognitive categories and selecting attribution-ranked coordinates, while its observed separation does not prove strict functional independence.
A. Continual Learning and Interference Mitigation
- Regularization, rehearsal, and distillation address forgetting through parameter constraints, replayed examples, or preserved predictions, but introduce additional storage or computation costs.Subspace methods instead limit interaction between task updates; CiMEO restricts updates to cognition-attributed weight coordinates.Source: A. Continual Learning and Interference… · p. 3
B. Parameter Isolation, Sparse Adaptation, and Routing
- CiMEO groups instructions by six Bloom-inspired cognitive categories, selects attribution-ranked pretrained coordinates, and uses a frozen cognition classifier to route updates.This differs from task- or dataset-based allocation and is paired with sparse optimizer-state storage.Source: B. Parameter Isolation, Sparse… · p. 3
C. Functional Localization and Gradient-Based Attribution
- Gradient-based identification and localized continual updates are not novel in isolation; CiMEO combines cognitive-category organization with instance-level routing and state-serialized AdamW.The authors characterize the observed separation as empirical evidence, not proof of strict functional independence.Source: C. Functional Localization and… · p. 3
D. Memory-Efficient Optimization for LLM Adaptation
CiMEO combines sparse optimizer-state storage with attribution-based selection of original weight coordinates and cognition-based routing. Its method begins with Bloom-labeled instructions, then uses those labels to identify subnetworks and direct adaptation.
D. Memory-Efficient Optimization for LLM Adaptation
- CiMEO selects cognition-relevant original weight coordinates for each incoming instruction and stores optimizer moments only for those coordinates.The design targets cross-category interference and persistent optimizer-state memory jointly; it does not guarantee lower activation memory or dense backward-pass computation.Source: D. Memory-Efficient Optimization for… · p. 4
III. METHOD
- The method collects a cognitive corpus, applies state-serialized AdamW, and defines continual learning based on CiMEO.The framework diagram also depicts attribution-based subnetwork identification and cognition-classifier routing.Source: Figure 2 · p. 4, §III · p. 5
A. Cognitive corpus collection
- Each plain-text instruction is assigned one of six Bloom taxonomy labels: memory, understanding, application, analysis, evaluation, or creation.The frozen instruction-tuned LLM receives each sample as a multiple-choice prompt to select the correct option.Source: A. Cognitive corpus collection · p. 5
- Beam search generates multiple candidate sequences, and the option with the highest selection frequency across runs is chosen as the final answer.The procedure retains the top-k sequences by normalized score at each decoding step.Source: A. Cognitive corpus collection · p. 5
B. Cognition-isolated subnetwork identification
The method estimates parameter importance for each Bloom-inspired cognition category from gradients, then selects and stores category-specific subnetworks for later sparse updates.
B. Cognition-isolated subnetwork identification
- Gradient attribution estimates each parameter’s importance to a cognition category by aggregating gradient magnitudes over labeled examples.The method uses the magnitude of each parameter’s gradient under cognition-labeled supervision.Source: B. Cognition-isolated subnetwork… · p. 5
- For each category, the procedure samples labeled instances and applies a multiple-choice cognition-classification prompt with a negative log-likelihood supervision loss.Source: B. Cognition-isolated subnetwork… · p. 5
- Within each weight tensor, it ranks elements by attribution and selects the top p% as that category’s subnetwork.The selection threshold is the (1-p)-percentile of attribution scores.Source: B. Cognition-isolated subnetwork… · p. 5
- The selected parameter positions are stored in a hash table to support sparse updates by cognition.The procedure yields six cognition-specific subnetworks for routing and memory-efficient optimization.Source: B. Cognition-isolated subnetwork… · p. 5
C. State-serialized AdamW
State-serialized AdamW represents sparse subnetwork parameters and optimizer states through a cognition-indexed hash table.
C. State-serialized AdamW
- The method serializes parameters, gradients, and AdamW moments through hash table H(c) for memory-efficient updates on cognition-isolated subnetworks.Source: C. State-serialized AdamW · p. 5
- H(c) maps the sparse parameters and their positions for cognition c, while the optimizer’s moving averages are defined over the corresponding gradients.The notation identifies one moving average on the gradients and another on their element-wise squares.Source: C. State-serialized AdamW · p. 5
A. Endogenous cognition recognizability in LLMs
The study profiles cognition-specific subnetworks in a pretrained LLM and examines how their overlap varies across categories and layers.
A. Endogenous cognition recognizability in LLMs
- The identification procedure uses Bloom-labeled instructions from AM-DeepSeek-R1-Distilled-1.4M to find cognition-relevant subnetworks.The overall procedure uses cognition-annotated instructions to identify isolated subnetworks in the model.Source: A. Endogenous cognition… · p. 5
- Across cognition pairs, average subnetwork overlaps are below 0.002%, suggesting an endogenous partition of cognition in pretrained LLMs.Figure 3a reports overlap rates between sparse subnetworks across cognition pairs.Source: A. Endogenous cognition… · p. 5, Figure 3 · p. 6
- Most layers show similarly low overlap, but the second and third layers have higher overlap across cognition pairs.The layer-wise pair curves show pronounced peaks at the second and third layers.Source: A. Endogenous cognition… · p. 6
- The authors relate the early-layer overlap peaks to weaker cognitive discrimination in general representations, with deeper representations becoming more task- and cognition-specific.Source: A. Endogenous cognition… · p. 6
B. Continual learning performance
CiMEO routes examples to cognition-specific subnetworks and restricts updates using a state-serialized optimizer. Across five benchmarks, it has the highest average performance among the compared methods, while the reported analyses indicate balanced gains and support for gradient-based selection.
B. Continual learning performance
- CiMEO routes each sample to a cognition-specific subnetwork and uses state-serialized AdamW for sparse updates.A frozen cognition classifier performs the routing; optimizer states are maintained for selected parameters.Source: B. Continual learning performance · p. 6
- The evaluation compares methods under a shared setup and uses five benchmarks spanning code generation, exams, broad subject understanding, and mathematics.The shared setup uses batch size 64 and learning rate 1e-5; PEFT methods use low-rank dimension 16.Source: B. Continual learning performance · p. 6
- 38.92% average performance is the highest among compared methods across five benchmarks, although CiMEO has no single-benchmark top score.On HumanEval-CN, Olora scores 29.27% and CiMEO scores 25.00%; CiMEO scores 6.98% on GaokaoBench 2010 and improves by 28.85% over baseline on CMMLU.Source: B. Continual learning performance · pp. 6–7, Table I · p. 7
- On four benchmarks, maximum-gradient attribution substantially outperforms the minimum-attribution and random-sampling variants.The reported comparison covers HumanEvalCN, GaokaoBench 2010, CMMLU, and C-Eval.Source: B. Continual learning performance · p. 7
- The authors interpret the cross-task gains as indicating more uniform performance and potential mitigation of degradation from task interference during sequential learning.They also suggest the gains may support knowledge accumulation and transfer across domains.Source: B. Continual learning performance · p. 7
C. Stability-plasticity trade-off performance
CiMEO balances resistance to forgetting with incremental learning across benchmarks, while its minimum-attribution and random variants often trade low forgetting for weaker learning gains.
C. Stability-plasticity trade-off performance
- Forgetting and incremental learning rates are the two metrics used to evaluate anti-forgetting and incremental learning.The predictions of the fine-tuned models are compared with those of the base model.Source: C. Stability-plasticity trade-off… · p. 7
- Across HumanEval-CN, CMMLU, and GSM8K, CiMEO has consistently low forgetting and ranks among the top methods for incremental learning, with a more favorable stability–plasticity trade-off.Its average forgetting rate is close to the best-performing methods.Source: C. Stability-plasticity trade-off… · p. 7, Figure 4 · p. 8
- CiMEO w/ min and CiMEO w/ random achieve very low forgetting on some benchmarks but generally have low incremental learning rates across tasks.The authors suggest these selection strategies do not effectively capture cognition-relevant parameters for continual adaptation.Source: C. Stability-plasticity trade-off… · pp. 7–8
D. Weight stability analysis
The weight-space analyses indicate that CiMEO largely preserves dominant pretrained subspaces and induces only minor rotations during continual learning.
D. Weight stability analysis
- The analysis uses singular value decomposition to examine weight-space transformations, comparing dominant subspaces of the base model and fine-tuned checkpoints.Subspace overlap is defined on a 0–1 scale, with 1 indicating identical subspaces and 0 indicating orthogonal subspaces.Source: D. Weight stability analysis · p. 8
- CiMEO achieves high subspace overlap comparable to the most stable baselines, suggesting it largely preserves dominant weight subspaces during continual learning.Higher overlap indicates stronger subspace stability and smaller changes in learned representations; the authors note its association with resistance to forgetting.Source: D. Weight stability analysis · p. 8, Figure 5 · p. 9
- CiMEO exhibits relatively small mean and maximum principal angles, indicating only minor subspace rotations during continual learning.The authors interpret this as building primarily on the original feature structure rather than substantially reorienting it.Source: D. Weight stability analysis · p. 9, Figure 5 · p. 9
V. CONCLUSION
The paper localizes cognition-associated parameter subsets and uses them in CiMEO, which combines routed localized updates with state-serialized optimization. Across five benchmarks, CiMEO achieves the highest average performance among compared methods while balancing knowledge acquisition and retention.
V. CONCLUSION
- Gradient-based attribution revealed distinct parameter subsets associated with six cognitive categories derived from Bloom’s taxonomy.CiMEO integrates cognition-guided routing, attribution-based parameter selection, and state-serialized AdamW optimization for localized adaptation with reduced optimizer-state requirements.Source: §V · p. 9
- Across five benchmarks, CiMEO achieves the highest average performance among compared methods while balancing knowledge acquisition and retention.Ablation studies validated the importance of subnetwork selection, and weight-space analyses indicated substantial preservation of pretrained representations.Source: §V · p. 9
- The findings support functional subnetwork localization as a promising approach to efficient and stable continual learning in LLMs.The authors suggest cognition-guided updates can mitigate interference without extensive modifications to pretrained models.Source: §V · p. 9