Source-linked AI summary
Learn Here, Move Less Elsewhere: Input-Conditioned Plasticity from Retained-Domain Activation Atlases
Jiangtao Lin, Bangyang Wei, Yihang Ding, Siyi Liu, Yuhan Dong
TL;DR
Task-specific updates can rewrite answers beyond their training task, complicating efforts to preserve existing behavior. ATLAS uses retained activation geometry to condition a shared task residual, and on Qwen3-8B it achieves lower mean retained-output KL than seven published baselines at shared coding requirements.
Problem
Task-specific training can change a language model’s answers on unsupervised domains, including making some original answers incorrect, complicating continuity across model versions.
Source: §1 · p. 1Method
ATLAS routes tokens through retained activation neighborhoods and uses their centers and directions to condition a shared low-rank task residual during training and inference.
Source: §1 · pp. 1–2Results
On Qwen3-8B, ATLAS achieves lower mean retained-output KL than all seven published baselines at shared coding requirements, consistently across three training seeds.
Source: §8 · p. 9Takeaways & Limitations
Across five backbones and two retained domains, experiments demonstrate coding gains with reduced output movement, alongside fewer rewritten mathematical answers and more stable commonsense choices.
Source: §8 · p. 9
Abstract
from arXiv · showhide
Task-specific fine-tuning can rewrite a language model's answers beyond the training task, complicating updates that must preserve existing behavior. We introduce ATLAS, which turns retained-domain representations into an input-dependent rule for task adaptation. An activation atlas supplies local reference centers and directional filters to a shared low-rank residual. Target supervision learns the residual, while retained geometry shapes its action throughout training and inference. On Qwen3-8B, ATLAS achieves lower mean retained-output Kullback-Leibler (KL) divergence than all seven published baselines at shared coding-performance requirements, with consistent advantages across multiple training seeds. Structural comparisons identify the contributions of retained reference states and directional conditioning, and answer-level analyses show fewer rewritten mathematical answers and more stable commonsense choices. Experiments spanning five backbones and two retained domains further demonstrate coding gains with reduced retained-output movement. With compact storage and modest decoding overhead, ATLAS provides a practical mechanism for acquiring specialized skills while maintaining continuity in existing responses.
1 INTRODUCTION
Task-specific updates can alter answers on unsupervised domains, so adaptation must acquire a skill while controlling changes to existing behavior. ATLAS uses retained activation geometry to condition a shared task residual, and the paper reports lower retained-output KL and fewer rewritten answers in its evaluations.
Motivation
- Task-specific training can change mathematics and commonsense answers, including making some original answers incorrect, so updates must control effects on existing behavior.Other answers may change without changing correctness, which matters when applications rely on continuity across model versions.Source: §1 · p. 1
Approach
- ATLAS routes each token to nearby retained activation neighborhoods and uses their geometry to condition one shared task residual.The atlas stores centers and leading local variation, while target supervision trains the shared low-rank branch.Source: §1 · p. 1
- Retained geometry shapes the update during both training and inference, while target supervision sets the learning destination.The shared factors are reused across charts rather than adding separate task experts.Source: §1 · p. 2
Main contributions
- On Qwen3-8B, ATLAS achieves lower mean retained-output KL than all seven published baselines at shared coding requirements, with consistent advantages across three training seeds.Cross-family results extend the learning–retention benefit to GLM.Source: §1 · p. 2
- Answer-level analyses connect distributional preservation with fewer rewritten mathematical answers and more stable commonsense choices.The analysis distinguishes harmful flips, incidental corrections, and changes between distinct incorrect answers.Source: §1 · p. 2
2 RELATED WORK
Prior approaches preserve behavior through learning signals, geometric constraints, or input-dependent computation. ATLAS differs by using the current token’s local retained neighborhood as an explicit reference for a shared residual.
Preservation through learning signals
- Learning-signal methods shape optimization through parameter penalties, prediction preservation, memory-based gradient constraints, or token-loss selection and reweighting.The paper contrasts these optimization signals with conditioning a shared residual on retained representations.Source: §2 · pp. 2–3
Geometry of preserved computation
- Geometric methods constrain gradients, activation directions, or low-rank subspaces, but do not select a local retained reference for each token.The cited approaches include OGD, GPM, O-LoRA, CorDA, LoRA-Null, OPLoRA, and CLoRA.Source: §2 · p. 3
Input-conditioned adaptation
- Input-conditioned methods make transformations or routing depend on the current state, while ATLAS uses local retained centers and directions to condition one shared residual.The discussion names TopLoRA, SLIM, and PASs-MoE as related approaches.Source: §2 · p. 3
Comparison with prior methods
- At shared coding requirements on Qwen3-8B, ATLAS achieves 10.16% to 17.60% lower mean retained-output KL than seven published baselines.The comparison spans subspace constraints, input conditioning, and token-level strategies.Source: §2 · p. 3
Behavioral continuity across updates
- Answer-level analysis distinguishes harmful flips, incidental corrections, and changes between distinct incorrect answers, finding fewer rewritten math answers and more stable commonsense choices after code adaptation.The paper situates this analysis alongside prediction-churn and negative-flip measures of behavioral continuity.Source: §2 · p. 3
3 RETAINED GEOMETRY FOR TASK ADAPTATION
ATLAS combines a fixed atlas of retained activation neighborhoods with a shared trainable low-rank residual. Routed centers and directional filters condition the residual’s input, output, and scale for each token.
Retained geometry for task adaptation
- ATLAS inserts a residual at a selected transformer layer, with target supervision training the residual and retained geometry supplying input-dependent conditioning.The starting model and retained distribution are denoted f0 and R, and the target distribution is T.Source: §3 · p. 3
Constructing the atlas
- The atlas groups retained-calibration activations into local neighborhoods whose charts store centers and orthonormal PCA bases; the backbone and atlas stay fixed during target training.Only the shared residual factors are updated.Source: §3 · p. 3
Routing a token to local geometry
- Routing selects nearby charts and mixes their centers and local filters, producing a token-dependent reference and attenuation operator.The filter attenuates directions represented by the local PCA bases while preserving their common orthogonal complement.Source: §3 · p. 4
Filtering the task residual
- The routed filter is applied to both the input and output of a rank-r branch, determining what the branch reads and the directions it emits.Trainable factors A and B form the shared branch, and the positive scale s controls how ρ(h) grows with the filtered-state norm.Source: §3 · p. 4, Figure 2 · p. 4
Implementation
- Sharing A and B across charts keeps the trainable parameter count at 2dr regardless of geometric resolution; the Qwen3-8B configuration uses 131,072 trainable parameters.The same token’s charts and routing weights are reused during generation, tying cost to selected local ranks.Source: §3 · p. 4
4 EVALUATING LEARNING AND RETAINED BEHAVIOR
The evaluation tests code learning alongside retained behavior on mathematics and commonsense tasks. It compares methods on learning trajectories and measures retained distributional movement at shared coding requirements.
Tasks and models
- MBPP provides training code problems and EvalPlus evaluates generated programs with MBPP+ tests; GSM8K and CommonsenseQA measure retained mathematics and commonsense behavior.The 224 evaluation problems are separate from 120 training and 43 validation problems.Source: §4 · p. 5
- Qwen3-8B is the main backbone for baseline and residual-structure comparisons, while the study also covers five backbones, two retained domains, and three training seeds.Evaluation uses complete five-epoch trajectories and single-layer adaptation.Source: §4 · p. 5
Comparisons
- The comparison includes seven published methods plus unconstrained residual, retained-example training, output distillation, and a one-chart construction.Methods share a single-layer adaptation budget, and validation loss selects learning rates.Source: §4 · p. 5
Measuring learning and retained movement
- MBPP+ pass@1 measures code learning, while retained-output KL measures next-token distribution changes on shared teacher-forced contexts.The evaluation fixes reference distributions and non-padding context positions within each comparison.Source: §4 · p. 5
- At each shared coding requirement, each method contributes its lowest observed KL among checkpoints meeting that requirement, and seed ratios are geometrically averaged over the shared score range.The seven-method table uses a range common to every method and seed.Source: §4 · p. 5
5 LEARNING WITH LESS RETAINED-OUTPUT MOVEMENT
ATLAS reduces retained-output KL relative to published adaptation baselines at shared coding requirements, with the advantage extending across seeds, model settings, and a second model family.
Lower retained KL than seven published baselines
- At 13 shared coding requirements, ATLAS has lower mean retained KL than all seven published baselines on Qwen3-8B, with reductions of 10.16%–17.60% across eight configurations and every seed.Table 1 covers GSM8K and three seeds; its caption specifies 148–160 solved MBPP+ problems.Source: §5 · pp. 5–6, Table 1 · p. 6
Learning trajectories show the preservation gain
- Figure 3a compares three-seed geometric-mean KL at shared coding requirements, while its absolute KL axis tracks mathematical-output movement as coding performance changes.The caption describes requirements of 151–161 solved MBPP+ problems and pairwise shared requirements for each seed.Source: §5 · p. 6, Figure 3 · p. 7
Learning and retention across scales and families
- Across six evaluated settings, all 18 three-seed summaries and all 54 seed-level comparisons with Plain, Replay, and Output-KL favor ATLAS, with an aggregate KL reduction of 62.03%.The settings span five backbones and two retained domains.Source: §5 · p. 6
- On GLM-4-9B, ATLAS checkpoints average 70.54% MBPP+ accuracy versus 36.61% for the starting model, while aggregate retained KL is 67.24% lower than Plain, Replay, and Output-KL.The retained-output comparison uses shared coding requirements.Source: §5 · p. 6
6 HOW RETAINED GEOMETRY SHAPES THE RESIDUAL
Structural comparisons find retained centers provide the largest preservation gain over an unconstrained residual, with distance scaling and local directions refining the result.
Centers establish a useful reference
- Centering reduces KL by 38.44% relative to Plain, and distance scaling adds a 2.39% reduction; retained centers provide the largest improvement in the comparison.Table 2 compares residual structures at common coding requirements, with three seeds and 151–160 solved problems.Source: §6 · p. 7, Table 2 · p. 7
Local directions refine the conditioning
- Full ATLAS reduces KL by 4.64% relative to centering and 2.31% relative to calibrated distance scaling, with both gains present in every seed.PCA-derived amplitude improves on distance scaling by 1.62%; the complete construction has the lowest aggregate KL.Source: §6 · pp. 7–8
Residual and weight-adapter comparison
- ATLAS achieves 12.98% lower mean KL than native Weight-LoRA at common coding requirements, extending the comparison to a standard low-rank parameterization.Native Weight-LoRA updates the attention output projection with the same rank and parameter count.Source: Table 2 · p. 7, §6 · p. 8
Direction matters at equal effective magnitude
- At equal effective residual magnitude, ATLAS has lower KL in 11 of 12 comparisons across four Qwen backbones, with an aggregate reduction of 3.89%.Fitted local bases also outperform random bases in all 11 comparisons.Source: §6 · p. 8
Geometric resolution
- Qwen generally favors local charts, whereas GLM favors the global construction.The passage reports complete per-model chart-comparison results in Appendices B and D.Source: §6 · p. 8
7 FROM OUTPUT MOVEMENT TO ANSWER REWRITING
Answer-level analyses find that reduced output movement corresponds to fewer rewritten mathematics answers and more stable commonsense choices, while distinguishing correctness changes from other answer changes.
Answer-churn measures
- Answer churn counts changed normalized labels as correct-to-wrong, wrong-to-correct, or distinct-incorrect transitions, complementing teacher-forced KL on shared reference prefixes.Generated answers provide an answer-level view because each token shapes context for subsequent predictions.Source: §7 · p. 8
Fewer rewritten mathematical answers
- On Qwen3.5-9B, rewritten mathematical answers fall from 22.43% to 10.18%, a 54.60% reduction across nine matched comparisons.Correctness-transition reductions persist under changes to batch composition and question order.Source: §7 · p. 8
Conservation acts in both directions
- On Qwen3.5, harmful flips fall by 59.86% and incidental corrections by 59.57%, yielding closely balanced reductions across both transition types.The analysis therefore covers previously correct and previously incorrect answers.Source: §7 · p. 8
Stability within the incorrect-answer set
- In eight target-matched Qwen3-8B comparisons, changes between distinct incorrect answers fall from 123 to 96, a 21.95% reduction, while correctness-transition totals remain close at 248 and 251.The comparison tracks answer identity beyond correctness labels.Source: §7 · pp. 8–9
More stable commonsense choices
- On CSQA, ATLAS reduces choice churn by 38.42% while maintaining nearly the same conditional-choice accuracy as matched methods.Conditional choice selects the highest-likelihood option.Source: §7 · p. 9, Figure 4 · p. 9
8 CONCLUSION
ATLAS uses retained activation geometry to condition a shared low-rank task residual, and the conclusion reports reduced retained-output movement alongside coding gains and answer continuity.
8 CONCLUSION
- ATLAS uses local centers and directional filters to condition a shared low-rank task residual, with the atlas active throughout training and inference and no separate trainable branch per chart.On Qwen3-8B, it achieves lower mean retained-output KL than all seven published baselines at shared coding requirements, consistently across three seeds.Source: §8 · p. 9
- Retained centers provide the largest structural improvement over an unconstrained residual, while distances and local directions further refine preservation.The conclusion also reports coding gains with reduced output movement across five backbones and two retained domains, plus fewer rewritten mathematical answers and more stable commonsense choices.Source: §8 · p. 9
A.1 TARGET AND RETAINED DATA
- MBPP training, validation, and test use separate source partitions; the sanitized training set has 120 problems and validation set 43, while target scores use 224 evaluation problems.The retained KL subsets belong to their corresponding behavior evaluation sets; atlas fitting and Replay share training examples.Source: §A.1 · p. 14, Table 3 · p. 14
- GSM8K generation uses all 1,319 test problems, while distributional measurements use 512 retained examples; CSQA evaluates on 1,221 validation questions with a 512-question KL subset.For CSQA conditional choice, the evaluation considers five option labels under the shared prompt.Source: §A.1 · p. 14
A.2 TRAINING AND SELECTION
The experiments use fixed backbones and a shared low-rank residual, comparing retained-output KL at matched coding requirements. Retained activations form a local atlas, while method-specific settings define how the atlas and residual are constructed and evaluated.
A.2 TRAINING AND SELECTION
- The five backbones use five-epoch trajectories with three seeds, fixed backbone weights, and a standard ATLAS residual of rank 16 and scale 64.Qwen studies use a learning rate of 2 × 10^-4; GLM selects 4 × 10^-4 using Plain.Source: §A.2 · pp. 14–15, Table 4 · p. 15
A.3 INTEGER TARGET-SCORE COMPARISON
- At each shared integer coding requirement, the comparison averages each method’s minimum observed retained KL within a seed, then summarizes ratios across three seeds.The Qwen3-8B seven-method comparison uses the common range k = 148–160; supplementary continuous averaging weights intervals by score width.Source: §A.3 · p. 15
A.4 ACTIVATION SAMPLING AND LOCAL PCA SUPPORT
- The standard atlas fits 32 local charts from retained activations, using cluster-centered rank-16 PCA and routing each token to its four nearest centers.The main atlas uses 4,096 retained-domain vectors; a separate equal-support construction selects 128 nearest vectors per chart for fitting-time benchmarking.Source: §A.4 · pp. 15–16
A.5 WEIGHT-ADAPTATION MODULES AND METHOD PARAMETERS
- Weight-adaptation baselines update Qwen3-8B’s attention output projection, while ATLAS and residual Plain are inserted at decoder-layer output.LoRA-Null and CorDA-KPM initialize from retained activations; STM masks response tokens once using the starting model, and TALR reweights token loss using response-token NLL.Source: §A.5 · pp. 16–17
B COMPLETE LEARNING RESULTS
Across the learning trajectories, ATLAS lowers retained KL at shared coding requirements, including comparisons with published methods and structural variants. The section also examines target learning, validation-selected outcomes, and the sensitivity of its retention advantage to design choices and score-count variation.
B COMPLETE LEARNING RESULTS
- ATLAS improves mean MBPP+ scores over frozen Qwen scores on all four Qwen backbones, while local charts generally outperform Global-PCA on Qwen and the reverse holds for GLM.Mean ATLAS gains range from 2.32 to 18.87 percentage points across the Qwen backbones; GLM reaches 157–159 solved problems at its best checkpoint across seeds.Source: B COMPLETE LEARNING RESULTS · p. 17, Table 8 · p. 18
- ATLAS has lower mean retained KL in all 54 seed-level comparisons and at 743 of 760 target levels; its all-seed geometric reduction is 62.03%.Across settings, reductions are 58.56% across four Qwen backbones on GSM8K, 68.96% on Qwen3-8B with CSQA, and 67.24% on GLM-4-9B with GSM8K.Source: B COMPLETE LEARNING RESULTS · p. 17, Table 6 · p. 17, Table 7 · p. 18
C.1 IMPLEMENTATION AND TUNING
- The published-method comparison uses shared MBPP+ training problems, three seeds, five epochs, and approximately matched trainable parameter counts across eight configurations.The comparison includes seven published methods, with OPLoRA evaluated at k = 16 and k = 128.Source: §C.1 · p. 22, Table 9 · p. 22
C.2 INDEPENDENT VALIDATION SELECTION
- ATLAS has lower retained KL in all 24 independent validation-selected comparisons, with geometric mean reductions ranging from 15.47% to 88.00% across configurations.In the LoRA-Null/OPLoRA subset, the nine comparisons yield a 67.64% geometric KL reduction and a mean target-score difference of −0.099 percentage points; eight target-matched pairs all favor ATLAS, with a 50.32% geometric reduction.Source: §C.2 · pp. 22–23, Table 11 · p. 23
D.1 RESIDUAL STRUCTURES AT SHARED INTEGER TARGET SCORES
- Centering reduces KL by 38.44% relative to Plain; distance scaling adds a 2.39% reduction and PCA-derived scaling adds 1.62%, with each step favorable in every seed.Across the tested structures, ATLAS has lower KL in every seed except input filtering, whose aggregate reduction is 0.52%.Source: §D.1 · pp. 23, 25, Table 12 · p. 24
D.2 CENTERING AND CALIBRATED DISTANCE SCALING
- The distance scale is calibrated on retained atlas activations to match the sample-average amplitude of PCA-derived scaling, while retaining distinct tokenwise responses.The calibrated scale is approximately 78.6227 and is shared across seeds.Source: §D.2 · p. 25
D.3 SENSITIVITY TO INTEGER TARGET REQUIREMENTS AND SCORE COUNTS
- Under a conservative simultaneous one-problem count shift, ATLAS retains reductions of 3.62% and 1.10% for the two comparisons, with all three seeds favorable in both.Sensitivity scenarios hold trained checkpoints and KL values fixed while varying target requirements or hypothetical solved-problem counts.Source: §D.3 · p. 25, Table 13 · p. 25
D.4 LOCAL OPERATORS
- At equal mean MBPP+ across seeds, adding input filtering gives 7.90% lower KL and 40.94% lower displacement; learned charts also lower KL in 11/11 Random-local and 9/12 Shuffled-patch comparisons.The corresponding geometric reductions for learned charts are 5.53% and 4.13%.Source: §D.4 · p. 26, §D.5 · p. 26
E ANSWER IDENTITY, CORRECTNESS, AND OUTPUT FORMAT
The appendix tracks behavioral changes at matched coding performance, distinguishing correctness transitions from answer-identity and format changes. Across mathematics and commonsense, ATLAS shows fewer rewrites, while additional analyses examine comparison and decoding conditions.
E.1 CHECKPOINT PAIRING FOR THE BEHAVIORAL ANALYSES
- The behavioral analyses pair checkpoints by coding score and validation-NLL proximity; churn is (F + G + W)/n, with reductions equally weighted across matched pairs.The appendix identifies 38 checkpoint pairs and reports three training seeds as the training replicates.Source: §E.1 · pp. 26, 28, §E.2 · p. 28
E.3 QWEN3.5-9B
- F fell 59.86%, G fell 59.57%, and overall answer churn fell 54.60% across matched Qwen3.5 comparisons.The frozen model was correct on 896 questions and wrong on 423.Source: Table 18 · p. 29, §E.3 · p. 29
E.4 GLM-4-9B
- GLM-4-9B results pool three target-matched seed pairs, yielding 3,957 question–pair observations per comparison; accuracy differences follow from (G − F)/n.Table 19 reports answer-change counts relative to the frozen model.Source: §E.4 · p. 29, Table 19 · p. 29
E.5 EXTERNAL SUBSPACE METHODS
- Against LoRA-Null/OPLoRA, changes between incorrect answers fell 25.41% under answer extraction and 21.95% under numeric normalization; the pooled reduction’s 95% interval was 3.03%–40.68%.Point estimates were positive in all three seeds: 9.09%, 31.25%, and 21.43%.Source: §E.5 · p. 29, Table 20 · p. 30
E.6 CSQA LABELS AND FORMAT TRANSITIONS
- On CSQA, conditional-choice label churn fell 38.42%; separate short-generation analyses track valid-format gains and losses across all questions.The choice analysis uses 12 target-matched pairs and all 1,221 questions.Source: §E.6 · p. 30, Table 21 · p. 30, Table 22 · p. 30
F.1 SHARED TARGET RANGES
- 53 of 54 seed comparisons were favorable under a shared interval across primary methods and seeds, with a geometric mean KL reduction of 60.92%.All 18 three-seed model–comparator summaries remained favorable.Source: §F.1 · p. 31
F.2 EQUAL CANDIDATE COUNTS
- 355 of 361 equal-candidate intersections favored ATLAS; Output-KL coefficient 1.0 achieved higher MBPP+ and lower retained KL than coefficient 0.3.One-problem target-count perturbations kept all 24 external seed directions favorable, with a smallest retained reduction of 0.333%.Source: §F.2 · pp. 31–32, Table 23 · p. 31, Table 24 · p. 31
F.3 DECODING BATCH AND QUESTION ORDER
- Reduced transition volume persisted under three decoding conditions; pooled ATLAS-minus-comparator accuracy differences were −1.398 points by answer extraction and −1.828 by numeric normalization.The comparisons used the same Qwen3.5 starting model and paired answers by question identity.Source: §F.3 · p. 32, Table 25 · p. 32
G TARGET-LEARNING ENGAGEMENT ACROSS MODELS AND TASKS
Target-learning engagement varies across the evaluated settings: Mistral-7B and Llama3-8B show no average coding gains, while the CSQA-target setting has only a small best gain. Resource measurements report modest decoding overhead and compact atlas storage.
G TARGET-LEARNING ENGAGEMENT ACROSS MODELS AND TASKS
- Across 75 checkpoints per backbone, mean MBPP+ changes from the starting model were −1.60 points for Mistral-7B and −0.70 for Llama3-8B; Gemma-2-9B had no sustained learning interval.Each backbone assessment covered five methods, three seeds, and five epochs.Source: G TARGET-LEARNING ENGAGEMENT ACROSS… · p. 32
- For CSQA-target Qwen3-8B, the best gains were 0.000 points for Plain and 0.104 for ATLAS; the 0.207-point shared range fell below the stated multi-seed thresholds.The protocol requires at least a one-point gain and a shared range of 0.5 points for multi-seed evaluation.Source: G TARGET-LEARNING ENGAGEMENT ACROSS… · p. 32
H COMPUTATIONAL RESOURCES
- On Qwen3-8B, ATLAS decoding took 2.52–3.13% longer than the frozen model, and prefill took 5.83–12.26% longer.The isolated local operator took 10.37–11.14 microseconds per token; measurements compared checkpoints at batch sizes 1, 8, and 24.Source: H COMPUTATIONAL RESOURCES · p. 33, Table 26 · p. 33
- Fitting the atlas to 4,096 activation vectors from 2,048 retained examples took 38.54 seconds, and the compressed atlas occupied 15.85 MiB.The fitting time excludes activation extraction.Source: H COMPUTATIONAL RESOURCES · p. 33