PaperlayerView exact revision on arXiv

Source-linked AI summary

CPR for LLMs: Critical-Point Routing against Catastrophic Forgetting in Domain Adaptation

Kwangmin Ki, Yunhun Nam, Jongheon Jeong, Jaehyung Kim

arXiv:2608.30158v1cs.CLcs.AI

TL;DR

Domain-specific SFT can cause catastrophic forgetting, while existing loss-based mitigations remain constrained by a domain-generality trade-off. CPR decouples the capabilities by routing between a base model and an SFT expert at critical tokens, using a hierarchical router with stabilized, threshold-gated dispatch. Across four model-domain settings, it surpasses SFT in domain performance, nearly recovers general capabilities, and invokes the expert on only about 30% of tokens.

  • Problem

    Domain-specific SFT often degrades general capabilities, and existing approaches remain bound to a domain-generality trade-off.

  • Method

    CPR routes between the original base model and an SFT expert at tokens where the base fails but the expert succeeds, using hierarchical routing, momentum smoothing, and threshold-gated dispatch.

  • Results

    Across four model-domain settings, CPR consistently surpasses SFT in domain performance while fully recovering its general-capability drop and reducing expert invocation to about 30% of tokens.

  • Takeaways & Limitations

    CPR provides a model-level alternative to single-weight adaptation by selectively combining base and domain expertise rather than compressing both capabilities into one model.

  • Takeaways & Limitations

    CPR requires simultaneously holding the base and expert models, increasing GPU memory overhead and leaving latency above single-model inference.

Abstract

from arXiv · show

Supervised fine-tuning (SFT) is the de facto standard for adapting large language models (LLMs) to target domains, but it often degrades the model's general capabilities, a phenomenon known as catastrophic forgetting. Existing approaches typically modify the SFT loss to mitigate forgetting, but they inevitably operate along a domain-generality trade-off. In this work, we step outside this trade-off by decoupling the two capabilities at the model level: we keep the original base model for general capability, and selectively invoke the SFT expert only when domain-specific knowledge is required. Specifically, we propose CPR (Critical-Point Routing), a token-level routing framework between a base model and its expert derivative, based on critical tokens where the base model fails but the expert succeeds. We train a lightweight hierarchical router that estimates the expert-call probability per token, and pair it with a tailored inference procedure that combines momentum smoothing and threshold gating. Across diverse model-domain configurations, CPR achieves state-of-the-art across all settings, surpassing SFT expert by 1.4-5.5% in domain performance while recovering its general-capability drop from 3.4-14.5% to at most 0.5%, with minimal overhead from invoking the expert on only one-third of tokens.

1 Introduction

Domain-specific SFT can improve target-domain performance while degrading general capabilities, leaving existing methods bound to a trade-off. CPR decouples these capabilities by routing an SFT expert only where needed, with smoothing and thresholded dispatch for efficient inference.

  • Motivation: Domain-specific SFT often degrades general capabilities even as it improves performance on target domains, creating a central adaptation challenge.The affected foundations include language understanding, instruction following, and multi-step reasoning.
  • Motivation: Existing approaches modify the SFT loss but compress domain and general capabilities into one weight set, preserving a domain-general trade-off.Figure 1 characterizes this single-model paradigm as unable to fully eliminate the trade-off.
  • CPR: CPR decouples capabilities at the model level by routing an SFT-tuned expert only on critical tokens where the base fails but the expert succeeds.The remaining tokens are handled by the original base model.
  • CPR: A lightweight hierarchical router uses query-level and token-level signals, while momentum smoothing and threshold-gated three-way dispatch stabilize inference.Dispatch can select the base, blend distributions, or select the expert.
  • Results: Across four model-domain settings, CPR attained the highest overall average, surpassed vanilla SFT in domain performance, and nearly recovered SFT’s general-capability loss.The evaluation used two LLMs, two domains, and six benchmarks spanning domain and general capabilities.
  • Results: CPR invoked the expert on only ~30% of tokens while matching or exceeding performance, reducing latency relative to always-on collaborative decoding.Selective invocation avoids the per-step dual-model cost.

2 Related Work

Related work addresses catastrophic forgetting mainly through loss modification or inter-model routing. CPR builds on fine-grained routing while applying token-level adaptive dispatch specifically to SFT-induced forgetting.

  • Catastrophic Forgetting in Domain SFT: Research on catastrophic forgetting reports that domain-specific SFT degrades LLM general capabilities across model sizes and continual instruction-tuning settings.Prior work attributes this partly to fine-tuning skewing implicit task inference toward the SFT distribution.
  • Catastrophic Forgetting in Domain SFT: Prior mitigation methods modify next-token training loss through probability rescaling, entropy-based gating, or representation regularization.These approaches target destructive gradients or preserve representations during adaptation.
  • Inter-model Routing for LLM Decoding: Inter-model routing commonly dispatches an entire query between models according to a cost-quality trade-off.RouteLLM and FrugalGPT exemplify query-level routing and cascading.
  • Inter-model Routing for LLM Decoding: Other routing methods combine model outputs at finer granularity, including latent per-token deferral, patch-level checkpoint switching, and token-level expert selection.These methods motivate finer-grained capability fusion but target heterogeneous strengths rather than SFT forgetting.

3 Method

CPR routes generation token by token between a frozen base model and domain expert, using critical-point supervision and a hierarchical router. Momentum smoothing and threshold-gated dispatch stabilize decisions while selectively invoking the expert.

  • 3.1 Problem Setup: CPR combines the base model and domain expert through token-level routing, assigning each token to the model best suited to generate it.The expert is selectively invoked rather than used for every token.
  • 3.2 Training: The hierarchical router combines a question-level macro encoder with a per-token micro router to estimate the probability of invoking the expert.The macro encoder mean-pools base hidden states over the question span, while the micro router uses the current base hidden state and macro summary.
  • 3.2 Training: Critical points are tokens where the base fails but the expert succeeds, and teacher-forcing comparisons automatically create labels for router training.Tokens correctly predicted by the base receive z_t = 0 regardless of the expert’s prediction.
  • 3.2 Training: The base and expert are frozen while only router parameters are trained with re-weighted binary cross-entropy to address sparse critical-point labels.The positive weight uses the corpus-level ratio N_0/N_1 in square-root form.
  • 3.3 Inference: Momentum smoothing replaces raw expert-call probabilities with an exponential moving average, reducing short-range dispatch noise.The default momentum factor is α = 0.5.
  • 3.3 Inference: Two thresholds divide the smoothed probability into Hard Base, Soft Blend, and Hard Expert regimes, while skipped expert steps are handled through batched KV-cache catch-up.The default thresholds are τ_low = 0.35 and τ_high = 0.65; the base still runs at every step.

4 Experiments

CPR is evaluated across four model-domain settings and six benchmarks, consistently improving domain performance while preserving general capability. Analyses show that its hierarchical routing, selective expert invocation, and inference design jointly support robustness and efficiency.

  • Setups: CPR evaluates two open-source LLMs adapted to math and medical domains across six benchmarks covering domain and general capabilities.The four model-domain combinations use Gemma3-4B and Llama3.1-8B, with domain and general capability averages defined over separate benchmark groups.
  • Main Results: CPR attains the highest overall average across all model-domain combinations, improving over the base model by 4.2% to 10.8%.It nearly recovers the SFT expert’s general-capability loss while surpassing the expert in domain performance.
  • Main Results: 14.5% general-capability loss for the Llama3.1-8B math SFT expert is reduced to 0.5% with CPR, while domain improvement rises from 16.6% to 22.1%.The latter corresponds to an additional 5.5% gain over the expert.
  • Main Results: 8.7% to 10.8% overall improvement over the base on math exceeds the 1.6% to 2.9% achieved by query- and patch-level routing baselines.The comparison uses the same base model and SFT expert for all routing methods.
  • Analyses: The hierarchical router, momentum factor, and three-way thresholds play complementary roles, with the default configuration generalizing without per-setting tuning.Removing the macro or micro router harms different capability axes, while hard switching reduces math accuracy and fully soft blending increases latency.
  • Analyses: With only 1K domain examples, CPR improves GSM8K accuracy by 10.62% while maintaining MMLU above the base level.Across 1K, 2K, 4K, and 8K training subsets, CPR achieves the highest overall average.
  • Analyses: Only about one-third of tokens invoke the expert, with higher invocation on domain than general benchmarks and routing concentrated on tokens needed for domain reasoning.CPR consequently defaults to the base on most general-query tokens, while its latency remains below ensemble decoding: 1.40× versus 1.89× for Gemma3-4B and 1.44× versus 2.25× for Llama3.1-8B.
  • Analyses: CPR generalizes to external experts when the expert outperforms the base on a meaningful subset of tokens, recovering general-capability drops of 18.2% and 23.3% to +0.7% and -0.6%.Across the two external-expert settings, CPR retains most domain gains and produces overall improvements of +17.7% and +1.8% over the respective bases.

5 Conclusion

CPR addresses catastrophic forgetting by decoupling domain and general capabilities through token-level routing between a base model and an SFT expert. It selectively invokes the expert, recovers general capability, reduces latency, and generalizes to external experts.

  • CPR decouples domain and general capabilities by routing tokens between the original base model and an SFT expert.The router invokes the expert where the base fails, while the base retains general capability elsewhere.
  • CPR consistently surpasses SFT in domain performance while fully recovering its general-capability drop across four model-domain settings.
  • CPR substantially reduces latency by selectively invoking the expert instead of using always-on collaborative decoding.
  • Token-level analyses show that expert calls concentrate on domain-relevant tokens, providing a mechanistic view of CPR’s capability preservation.
  • CPR generalizes to external experts without additional SFT.

Limitations

CPR’s model-level decoupling improves capability preservation but introduces memory, latency, and expert-availability constraints. Its effectiveness may decline when the expert-base performance gap is narrow or training data is scarce.

  • CPR requires simultaneously holding the base model and domain expert, increasing GPU memory overhead relative to single-model approaches.
  • Although the expert is invoked on only ~30% of tokens, the base must always run for router hidden states, leaving latency above single-model inference.
  • Skipped-token transitions incur key-value cache catch-up costs, although CPR remains faster than ensemble and contrastive decoding.
  • CPR assumes an expert outperforming the base on a meaningful token subset, so effectiveness may decline with a narrow gap or scarce domain data.

Broader Impact and Ethical Implications

CPR may support domain-specialized deployment without sacrificing general capabilities, but deployment requires careful expert vetting. The medical experiments are benchmark evaluations rather than clinical validation.

  • Preserving general reasoning is particularly valuable in high-stakes domains where brittleness can lead to harmful outputs.
  • CPR inherits biases and failure modes shared by the base and expert, which routing does not correct.
  • Because CPR can use arbitrary third-party experts post-hoc, expert models require careful vetting before deployment.
  • The medical experiments are intended only for research evaluation on public benchmarks and do not constitute clinical validation.
  • The appendix provides additional details on datasets, baselines, implementation, and prompt templates.

A.1 Datasets

The experiments construct math and medical experts using vanilla SFT, evaluate domain and general benchmarks under specified prompting protocols, and compare single-model and multi-model baselines.

  • Datasets: Domain experts are constructed with vanilla SFT on math and medical data.
  • Datasets: The math training set uses GSM8K’s approximately 8K grade-school problems with step-by-step rationales formatted as 0-shot CoT.
  • Datasets: The medical training set samples 8K PubMedQA examples, balanced between yes and no answers and formatted as 2-shot CoT.
  • Datasets: Evaluation covers math, medical, and general benchmarks, including GSM8K, ASDiv, SVAMP, PubMedQA, MedQA, CareQA, MMLU, CommonsenseQA, and ARC-C.
  • Evaluation Protocol: Math benchmarks use 0-shot CoT, while medical and general benchmarks use fixed 2-shot CoT demonstrations for controlled comparisons.
  • Baselines: Baselines include query-level RouteLLM, patch-level Switch Generation, fixed Ensemble and Contrastive Decoding, and loss-level DFT, EAFT, and LfU.
  • Baselines: Ensemble and Contrastive Decoding invoke both models at every decoding step, unlike CPR’s selective expert invocation.

A.3 Implementation Details

The appendix specifies reproducible training, inference, prompting, and baseline-comparison settings, including controlled router evaluations and the choice of official routing checkpoints.

  • Compute: All experiments use NVIDIA RTX A6000 GPUs, with each model fitting on one GPU and parallel execution across GPUs.
  • Router Training: CPR router training freezes the base and expert models, trains only router parameters, and uses identical two-layer MLP dimensions across domains.The macro encoder and micro router each use hidden dimension 256.
  • Critical-point Labeling: Critical-point supervision compares greedy base and expert predictions against ground-truth tokens using only the vanilla-SFT training data.Labels are assigned at each generation step without additional supervision data.
  • Inference Settings: All methods use greedy decoding with temperature 0 and sufficiently large generation limits, while CPR uses α = 0.5, τlow = 0.35, and τhigh = 0.65 by default.
  • Routing Baselines: Re-trained RouteLLM and Switch Generation routers underperform their official checkpoints, so the main experiments use the officially released versions.The authors attribute this likely to insufficient scale of the approximately 8K SFT dataset for learning effective routers.

B.2 Additional Efficiency Results

Expert invocation remains around one-third of tokens across benchmarks, with more calls in-domain than out-of-domain, indicating concentrated routing where domain knowledge is needed.

  • Around one-third of tokens invoke the expert across benchmarks, while in-domain queries consistently receive more expert calls than out-of-domain queries.The reported gap indicates that routing concentrates expert computation on tokens requiring domain knowledge.

B.3 Additional Ablation Results

Additional experiments support CPR’s hierarchical routing design, preferred inference hyperparameters, cross-domain generalization, and lightweight deployment variant.

  • Router Architecture: The hierarchical router outperforms either the macro encoder or micro router alone across all reported benchmarks.The design combines query-level and token-level routing components.
  • Momentum Decay: α = 0.5 outperforms α = 0.0 and α = 0.9, while the extremes respectively ignore history or delay switching.The intermediate momentum setting achieves the strongest overall accuracy.
  • Dispatch Thresholds: The default thresholds (τlow, τhigh) = (0.35, 0.65) outperform hard switching and fully soft blending in the 3-way dispatch.Hard switching removes soft blending, whereas fully soft blending always invokes the expert and increases latency.
  • Finance Generalization: CPR generalizes to finance, retaining most FinQA gains while substantially recovering MMLU capability and achieving the highest overall average.The finance evaluation uses FinQA in-domain and MMLU out-of-domain.
  • Open-ended Instruction Following: CPR limits the MMLU drop to 0.66 points in open-ended instruction following, compared with a 6.47-point drop for the SFT expert.The SFT expert has the higher AlpacaEval win rate, while CPR provides a more favorable adaptation-preservation balance.
  • Lightweight Variant: The LoRA expert variant reduces GPU memory from 16.60 GB to 8.83 GB while achieving an overall average of 56.87 versus 58.74 for the full expert.Latency increases modestly from 1.40× to 1.49×.

C Additional Qualitative Results

Qualitative results visualize CPR’s token-level dispatch and support the claim that single-model baselines remain on a domain-generality trade-off.

  • The qualitative appendix supplements quantitative evaluation with routing visualizations and a trade-off plot across four model-domain settings.
  • Token-level Routing: In-domain examples concentrate expert calls on mathematical numerals and operators or biomedical terminology, whereas out-of-domain examples invoke the expert more sparsely.Expert tokens are red, base tokens green, and intermediate shades indicate soft blending.
  • Token-level Routing: The routing patterns provide a mechanistic view of CPR’s invocation gap by directing expert computation to tokens needing domain knowledge.
  • Domain-General Trade-off: Single-model baselines show negative domain-versus-general accuracy slopes in every setting, with no baseline escaping the trade-off.Each baseline occupies a different point along the same empirically observed trade-off.
Loading 2608.30158v1…