Task Vector Descent: Learning from Non-IID Batches

Anton Baumann; Jonas Hübotter; Zeynep Akata; Andreas Krause

cs.LGOct 4, 2026

Task Vector Descent: Learning from Non-IID Batches

Authors: Anton Baumann, Jonas Hübotter, Zeynep Akata, Andreas Krause

Organizations: ETH Zurich · Technical University of Munich

Abstract

A central challenge in continual learning is to acquire new knowledge without forgetting what the model has already learned. This challenge appears in language model training when training data comes from various domain-, user-, or task-specific distributions that are encountered unevenly over time. In such settings, successive minibatches are temporally clustered by distribution instead of being sampled i.i.d. from the overall data mixture. Training on temporally clustered data induces a stability-plasticity tradeoff. Adapting the model to the active distribution can improve the model on the active distribution but may lead to a performance degradation on data it previously trained on. We find that this tradeoff intensifies with longer exposure to the same distribution. We therefore ask if the parameter displacement produced by such a sequence (the task vector) should be fully retained or applied partially. We compare applying the full displacement ( $λ=1$ ) with partial integration, which scales the task vector by $λ$ before applying it to the continuing model and scales the optimizer state by the same coefficient. Across continual pretraining, pretraining from random initialization, supervised post-training, and reinforcement post-training, we find that intermediate values of $λ$ often improve average continuing-model performance relative to full integration, particularly after longer same-distribution sequences. In continual-pretraining experiments with both controlled streams and naturally defined adaptation sequences, task-vector scaling outperforms full integration at the matched learning rate, showing that its benefits are not reproduced by learning-rate scaling alone.

Figures & tables

Panel of Figure 1 — Figure 1: Task-vector updates for non-i.i.d. minibatch streams. Minibatches arrive in task-clustered sequences (left). For each task, noisy minibatch updates produce a local endpoint; the continuing model applies a scaled task vector and begins the next task from the resulting checkpoint (right).

Panel of Figure 2 — Figure 2: Sequence-boundary evaluation. Plasticity uses the locally adapted endpoint; stability and AvgGain use the continuing model.

Panel of Figure 3 — Figure 3: Continual-pretraining gain. Solid curves show task-vector scaling; dashed curves show full integration at learning rate $\lambda\eta$ . Stars mark the best task-vector coefficient, error bars show SEM, and the black dashed line denotes the i.i.d. reference.

Panel of Figure 4 — Figure 4: Continuing-model gain for SDPO, SFT, and GRPO (left to right) as a function of $\lambda$ . Curves denote $L$ ; error bars show SEM across seeds, and stars mark the best $\lambda$ for each $L$ . Dashed lines are i.i.d. or joint-data references evaluated on the smallest- $L$ schedule in each setting.

Panel of Figure 5 — Figure 5: Stability–plasticity trade-off in continual pretraining. Solid curves show task-vector scaling and dashed curves show full integration at learning rate $\lambda\eta$ ; error bars show SEM.

Panel of Figure 6 — Figure 6: Non-targeted retention versus targeted adaptation for SDPO, SFT, and GRPO from left to right across task-vector scaling coefficients. The stability axis is the mean continuing-model gain over non-targeted distributions or contexts; the plasticity axis is targeted gain at the endpoint $\theta_{t}^{(L)}$ .

Panel of Figure 7 — Figure 7: Pre-training from random initialization. The left panel reports continuing-model macro BPT gain as a function of the task-vector scaling coefficient. Stars mark the best coefficient on each curve, and the dashed line is the shuffled i.i.d. reference. The right panel compares plasticity and stability.

Panel of Figure 8 — Figure 8: PG-19 BPT gain relative to initialization. Colors denote base learning rates $\eta$ ; points show $\lambda\in\{0.4,0.6,0.8,1\}$ . Gray dashed: full integration at matched effective rates. Error bars: $\pm 1$ SEM over three paired seeds.

Appendix figures & tables12 assets

Supplementary material from the paper’s appendix.

Appendix

Setting	Model	$K$	$L$	$\lambda$	# Seeds	Primary score
Continual PT	GPT-2	10	$1,32,64,128,256,512$	$.4,.6,.8,1$	3	BPT gain
Persona SDPO	Qwen3-4B	6	$4,8,16$	$0,.01,.05,.1,.2,.4,.6,.8,1$	3	Macro win rate
Reasoning SFT	Qwen2.5-3B	4	$32,64,128,256$	$0,.2,.4,.6,.8,1$	5	Best4
Reasoning GRPO	Qwen2.5-3B	4	$256,1024$	$0,.2,.4,.6,.8,1$	3	Best4
Scratch PT	GPT-2	10	$1,64,256,1024,2048$	$.2,.4,.6,.8,1$	3	BPT gain
Book-level PT	SmolLM-360M	200	$1$ – $158$	$.4,.6,.8,1$	3	BPT gain

Appendix

Table 1: Linear-integration experiments. In the fixed-distribution settings,

L

counts consecutive minibatches from one targeted distribution. In the on-policy extensions, it counts consecutive setting-specific collection units from one active context. For PG-19, sequence lengths vary with book length.

Panel of Figure 9 — Figure 9: Controlled streams. Each box contains $L$ batches from one distribution; i.i.d. batches draw from the complete mixture. Larger $L$ increases temporal clustering.

Asset / checkpoint	License / terms
Models and services
GPT-2 (124M) ( Radford et al., 2019 )	Modified MIT
Qwen3-4B , Qwen3-8B ( Yang et al., 2025 )	Apache 2.0
Qwen2.5-1.5B , Qwen2.5-7B ( Yang et al., 2024 )	Apache 2.0
Qwen2.5-3B , Qwen2.5-3B-Instruct ( Yang et al., 2024 )	Qwen Research License
Qwen2.5-Coder-1.5B , -7B ( Hui et al., 2024 )	Apache 2.0

Appendix

Table 2: Existing assets and their published licenses or terms. Asset names link to the source documentation; citations credit the original work.

Panel of Figure 10 — Figure 10: PG-19 book lengths and adaptation lengths for 200 books per seed. Outline histograms use common logarithmically spaced bins within each panel. Full-book token counts include the held-out portion; optimizer steps count only training updates before integration.

Seed	Full-book tokens (k)		Training targets (M)	Steps per book		Total steps
	Median	P90		Median	Range
0	80.9	221.4	18.49	9	1–91	2,353
1	81.1	219.0	18.98	9	1–57	2,410
2	83.2	242.1	20.41	9.5	1–158	2,585

Appendix

Table 3: Book lengths and training exposure per run. k and M denote thousands and millions of tokens. Training targets exclude held-out blocks, the first token of each block, and padding. P90 uses nearest rank.

Category	Seed 0	Seed 1	Seed 2
Novels	32.5	30.0	32.0
British Literature	20.0	18.0	19.5
American Literature	10.0	10.0	14.0
History - Modern (1750+)	10.5	10.5	13.0
Adventure	11.5	12.0	9.5
Biographies	9.0	8.5	12.5

Appendix

Table 4: Most frequent Gutenberg categories. Entries are percentages of the 200 books in each seed; overlapping labels need not sum to 100%.

Setting	GPUs/run	Timed runs	Hours/run	GPU-hours
Continual pretraining	4	72	0.11–0.55	79.85
Persona SDPO	2	84	0.91–2.94	292.57
Reasoning SFT	4	125	0.29–3.58	540.36
Reasoning GRPO	4	39	4.53–8.51	896.76
Pretraining from scratch	4	78	1.94–2.59	663.71
PG-19 continual pretraining	4	72	0.49–1.09	202.87

Appendix

Table 5: Recorded compute for the six training settings. Hours/run gives the runtime range; GPU-hours sums runtime multiplied by allocated GPUs. Timings include evaluation.

Panel of Figure 11 — Figure 11: Immediate effect of task-vector scaling for SFT (left) and GRPO (right). The dotted diagonal marks zero change in average performance relative to full integration.

Panel of Figure 12 — Figure 12: Task-count-normalized integration trade-off for continual pretraining and persona SDPO. The horizontal axis is the combined non-targeted recovery contribution and the vertical axis is the targeted-loss contribution. The dotted diagonal is the zero-change boundary. Circles mark partial integration ( $0<\lambda<1$ ), squares mark $\lambda=0$ where evaluated, and triangles mark $\lambda=1$ .

Panel of Figure 13 — Figure 13: Immediate effect of task-vector scaling during pretraining from random initialization. Non-targeted recovery and targeted loss are defined in Eq. 90 ; their difference gives the change in average performance. The dotted diagonal marks zero change. Colors denote sequence lengths as in Fig. 7 ; circles mark partial integration and triangles mark $\lambda=1$ . Error bars show SEM across seeds.

Panel of Figure 14 — Figure 14: GSM8K accuracy along interpolations between the initial model and checkpoints obtained after different numbers of MetaMathQA SFT steps.

Panel of Figure 15 — Figure 15: Direct interpolation between released Qwen2.5 base checkpoints and their Qwen2.5-Coder descendants. Columns show parameter scale. The top row reports targeted MBPP pass@1, while the bottom row reports the unweighted mean of full MMLU, normalized ARC-Challenge, normalized HellaSwag, and WinoGrande accuracy.

Explore similar work

CLARE: Scalable Class-Incremental Continual Learning via a Sparsity-Based Framework
Sep 15, 2026Yunxiang Fu, Meng Lou, Zicheng Liao +1Replay-Based Continual Learning Class-Incremental Learning
School of Computing and Data Science, The University of Hong Kong · Hong Kong Generative AI Research and Development Center
To Retain or to Adapt? Generalizing Continual Learning
Jul 6, 2026Giulia Lanzillotta, Mandana Samiei, Doina Precup +2Replay-Based Continual Learning Continual Learning
ETH AI Center, Department of Computer Science, ETH Zürich · Mila - Quebec AI Institute, School of Computer Science, McGill University · Mila - Quebec AI Institute +1
Latent-LoRA: Compact Latent-Space Adapters with Gradient-Free Routing for Continual Learning
Jul 26, 2026Reza Rahimi Azghan, Gautham Krishna Gudur, Giulia Pedrielli +2Replay-Based Continual Learning Catastrophic Forgetting
College of Health Solutions, Arizona State University, Phoenix, Arizona, USA · Department of Electrical and Computer Engineering, The University of Texas at Austin, Austin, Texas, USA · School of Computing and Augmented Intelligence, Arizona State University, Tempe, Arizona, USA +1

Sep 15, 2026, cs.LGJ/K move · Enter open · S save

CLARE: Scalable Class-Incremental Continual Learning via a Sparsity-Based Framework

Yunxiang Fu, Meng Lou, Zicheng Liao, Yizhou Yu

School of Computing and Data Science, The University of Hong Kong · Hong Kong Generative AI Research and Development Center

Replay-Based Continual Learning Class-Incremental Learning Continual Learning Catastrophic Forgetting Model Fine-Tuning Sparsity

Open paper AskPDF arXiv

Continual learning must balance the learning of new knowledge with the retention of previously learned knowledge to incrementally learn tasks from a data stream without catastrophic forgetting. While leveraging pretrained models has significantly advanced continual learning, existing methods exhibit a scalability bottleneck when trained sequentially on many tasks, suffering from performance degradation due to inter-task interference and loss of plasticity. Inspired by evidence that sparse fine-tuning achieves performance comparable to full fine-tuning, this paper presents a novel sparsity-driven continual learning framework. Our continual learning method, termed CLARE, operates in two stages: it first identifies a sparse, task-critical parameter mask via a sparsity-inducing objective, then performs mask-constrained fine-tuning by only optimizing parameters selected by the mask. This two-stage sparse adapter mechanism enables all tasks to be accumulated within a shared adapter space while reducing destructive interference across tasks. Extensive experiments demonstrate the scalability of CLARE. On the long task-sequence benchmark Omnibenchmark-1k, CLARE outperforms strong baselines in final accuracy by a large margin, e.g, improving EASE by 4.64% and 13.34% after learning 100 tasks, respectively.

Abstract

Explore similar work

CLARE: Scalable Class-Incremental Continual Learning via a Sparsity-Based Framework

To Retain or to Adapt? Generalizing Continual Learning

Latent-LoRA: Compact Latent-Space Adapters with Gradient-Free Routing for Continual Learning