PaperlayerView exact revision on arXiv

Source-linked AI summary

One Step, One Lead: Mitigating Higher-Order Interference in Multi-Domain Reinforcement Learning via Cross-Step Control

Zihan Lin, Xiaohan Wang, Jie Cao, Jiajun Chai, Guojun Yin, Wei Lin, Ran He

arXiv:2609.06469v1cs.LGcs.CL

TL;DR

Joint training across domains can degrade individual-domain performance, while single-step gradient and curvature diagnostics may miss sequential interference. OSOL uses adjacent token footprints to guide focus-domain corrections within GRPO. On Qwen3-30B-A3B, it reaches a 0.4822 domain-macro average, 5.7% above MGS.

  • Problem

    Joint RL across heterogeneous domains can produce uneven domain-wise gains, degradation, and training instability, while same-step gradient or curvature diagnostics do not directly capture interactions between consecutive updates.

    Source: §1 · pp. 1–2

  • Method

    OSOL designates a focus domain and uses the preceding token footprint to rank rebound risk and add a drift-ranked, adaptively scaled history correction within standard GRPO.

    Source: §1 · p. 2, §4.1 · p. 6

  • Results

    On Qwen3-30B-A3B, OSOL reaches a 0.4822 domain-macro average, 5.7% above MGS; cross-step backtracking correlates positively with subsequent task damage (ρs = 0.40).

    Source: §1 · p. 2, §5.2 · p. 7

  • Takeaways & Limitations

    Adjacent checkpoint footprints expose the cross-step output interaction, and OSOL's history correction targets the negative-to-positive backtracking component without explicit higher-order differentiation.

    Source: §1 · p. 2, §3.2 · p. 5

    The sign guarantee requires each selected token’s own local response to dominate aggregate cross-token effects.

    Source: §3.2 · p. 5

Abstract

from arXiv · show

Reinforcement learning (RL) across multiple domains can broaden the reasoning capabilities of large language models (LLMs), yet joint training often degrades individual-domain performance and can destabilize optimization. Existing work typically diagnoses such interference from a single-step view using first-order gradient alignment or curvature-based proxies. We show that this view can miss a critical form of sequential interference: same-point domain gradients may remain nearly orthogonal even when consecutive realized updates partially reverse one another in output space. We further show that consecutive token log-probability footprints recover this interaction directly from adjacent checkpoints as a local second-order interaction in output space, without explicitly reconstructing same-step curvature. Building on this insight, we propose OSOL, which designates a focus domain at each iteration, uses the preceding checkpoint footprint to rank token-level rebound risk, and applies a drift-ranked, adaptively scaled correction within the standard GRPO update. Our analysis shows that this correction suppresses the targeted cross-step output backtracking component. Controlled studies further show that cross-step backtracking is more strongly associated with subsequent task damage than same-point gradient diagnostics, while the preceding footprint ranks future rebound risk more accurately than Hessian-based proxies. On Qwen3-30B-A3B, OSOL reaches a domain-macro average of 0.4822, improving by 5.7% over the strongest compared baseline, without explicit higher-order differentiation.

1 Introduction

OSOL treats interference as an interaction between consecutive updates in task-conditioned output space, which same-point gradient and curvature diagnostics can miss. It uses preceding token footprints to guide cross-step corrections and improves the domain-macro average on Qwen3-30B-A3B.

1 Introduction

  • Joint multi-domain RL can yield uneven domain-wise gains, degrade individual-domain performance, and destabilize training, limiting broad capability improvement in one foundation model.The reported degradation varies by domain, with gaps between mixed-domain and single-domain GRPO.Source: §1 · p. 1, Figure 1 · p. 2
  • Positive same-point gradient cosine does not reliably rule out subsequent task damage, while broader Hessian-proxy coverage brings modest rebound-identification gains at substantially higher peak GPU memory.Nine of 16 replicates have positive gradient cosine, yet aligned cases can still incur substantial damage.Source: §1 · pp. 1–2, Figure 2 · p. 3
  • 0.4822 domain-macro average on Qwen3-30B-A3B, 5.7% above MGS, makes OSOL the strongest compared method in the reported comparison.OSOL also outperforms curvature-guided CGPO.Source: §1 · p. 2
  • Adjacent token footprints expose the mixed interaction between consecutive updates in task-conditioned output space and interpret opposing token movements as cross-step output backtracking.OSOL designates a focus domain and applies history corrections within standard GRPO without auxiliary supervision or explicit higher-order derivatives.Source: §1 · p. 2
  • Equalizing focus and non-focus weights or replacing drift-ranked token weights with random weights reduces the domain-macro average by 51.4% and 50.8%, respectively.These ablations indicate that both asymmetric weighting and drift-based token ranking matter in the reported design.Source: §1 · p. 3

2 Related Work

Prior methods commonly balance task contributions or coordinate gradients evaluated at the same checkpoint, while curvature approaches estimate isolated-update sensitivity. This paper instead focuses on the observable interaction between consecutive updates in task-conditioned output space.

2 Related Work

  • Existing multi-domain optimization methods balance task contributions or modify same-step gradient interactions, but do not directly model interference between consecutive mixed-domain updates.The surveyed approaches include adaptive weighting, multi-objective methods, gradient surgery, and LLM-specific weighting, sampling, and curriculum methods.Source: §2 · p. 3
  • Curvature-based analyses characterize isolated-update sensitivity and can be costly online at LLM scale, whereas this work examines consecutive-update interactions in task-conditioned output space.The related work also connects this perspective to policy geometry and checkpoint-based observability.Source: §2 · p. 3
  • The paper distinguishes its empirical-Fisher use from treating the empirical Fisher as a generic substitute for an arbitrary objective Hessian.Its focus is a directly observable output-space interaction, not an identification of the complete objective Hessian with the Fisher.Source: §2 · p. 3

3 Cross-Step Dynamics

The section formalizes cross-step interference through token log-probability changes on realized response coordinates. It connects adjacent-checkpoint footprints to output-space interaction and shows how a history-conditioned correction targets token rebound, subject to a local response condition.

3 Cross-Step Dynamics

  • A token log-probability footprint is the vector of log-probability changes induced by an update on fixed realized response-token coordinates.The task-conditioned output space is the vector of token log-probabilities on coordinates from the domain being analyzed.Source: §3 · p. 3

3.1 Output Geometry

  • The inner product of preceding and current base footprints on shared realized-token coordinates exposes their mixed output interaction under the empirical output Fisher.The controller needs only the already available preceding footprint, making the interaction observable from adjacent checkpoints without constructing the Fisher matrix.Source: §3.1 · p. 4
  • Cross-step backtracking describes output motion canceled in the net displacement; controlled backtracking mass targets tokens whose log-probability fell and then rose.The interaction is given a trajectory interpretation through opposing token movements.Source: §3.1 · p. 4

3.2 Backtracking Control

  • OSOL applies a nonpositive residual to eligible tokens with negative preceding drift, and the analysis shows contraction of the targeted negative-to-positive backtracking component.The history-conditioned penalty acts against increasing selected likelihood ratios; the local response condition is needed to transfer coefficient correction into nonpositive output change on selected tokens.Source: §3.2 · p. 5

4 Method Design

OSOL focuses each mixed-domain update on one domain while retaining smaller-weight contributions from the others. It uses the preceding checkpoint to identify negative-drift tokens, rank their correction strength, and add an adaptively scaled residual to the standard clipped GRPO objective.

4 Method Design

  • Each mixed-domain update designates one focus domain, gives it a larger base task weight, and applies history correction only to its eligible response tokens while retaining the other domains at smaller weights.A schedule selects the focus domain at each iteration.Source: §4 · p. 5
  • OSOL scores the same realized eligible focus-domain tokens using the frozen preceding checkpoint, then adds the resulting residual to the GRPO coefficient and optimizes with the unchanged clipped policy objective.History-dependent quantities are stop-gradient, and focus assignment does not change the group-relative reward signal.Source: §4 · p. 5

4.1 Update History

  • The preceding footprint is the log-probability difference across adjacent checkpoints on current rollout coordinates, and negative drift identifies tokens whose realized log-probability was reduced by the previous step.This is an observed preceding footprint, not a prediction of the next update, and requires no token matching or auxiliary model.Source: §4.1 · p. 6

4.2 Residual Construction

  • Within the negative-drift set, drift rank assigns stronger penalties to more negative tokens, while rank-based allocation reduces sensitivity to batch-scale changes and isolated extreme shifts.The stated rank profile uses (plo, phi, rmin, rmax) = (0.1, 1, −1, 0).Source: §4.2 · p. 6
  • Adaptive scaling sets the residual relative to the current base-coefficient scale, and the tokenwise history correction is optimized within the standard GRPO objective.The standard setting uses τ = 0.03.Source: §4.2 · p. 6

5 Experiments

OSOL improves the domain-macro average over the strongest compared baseline on both evaluated backbones, while ablations and diagnostic comparisons examine its components and the preceding footprint’s value. Experiments use a balanced four-domain training mixture and evaluate domain-level performance alongside token-level rebound risk.

5.1 Setup

  • The training mixture contains 7.5K examples each for mathematics, code generation, instruction following, and chat; the primary aggregate equally weights the four domains.Evaluation covers AIME 2024/2025 and AMC 2023, LiveCodeBench, IFBench, WildBench-Hard, and CreativeWriting v3.Source: §5.1 · p. 7

5.2 Main Results

  • 0.4822 domain-macro average on Qwen3-30B-A3B is 5.7% above MGS; on Qwen3-8B-Base, 0.3735 is 0.9% above GRPO+KL0.001.The Qwen3-30B-A3B result is 0.0262 points above MGS; Qwen3-8B-Base improves by 0.0224 points over standard GRPO.Source: §5.2 · pp. 7–8, Table 1 · p. 7, Table 2 · p. 8

5.3 Ablation Analysis

  • Removing the checkpoint-history residual lowers the dense-backbone domain-macro average to 0.3600 from 0.3735; equal task weights and random token weights reduce it by 51.4% and 50.8%.Response-level allocation reaches 0.3640, while token-level allocation adds 0.0095 points (2.6%), mainly through code and chat gains.Source: §5.3 · p. 8, Table 3 · p. 8

5.4 Mechanism Analysis

  • Preceding-footprint ranking reaches 0.567 ROC-AUC and 0.627 average precision, versus 0.498 and 0.579 for the Hessian proxy; opposing module contributions can cancel in global aggregation.On the three-seed control subset, the realized-update Hessian proxy reaches 0.549 ROC-AUC, still below checkpoint ranking at 0.557.Source: §5.4 · p. 9, Figure 4 · p. 9, Figure 3 · p. 9

6 Conclusion

The conclusion frames sequential interference as a multi-domain RL problem that may remain hidden in same-point gradient comparisons, and presents OSOL as a cross-step output-space controller. The controlled diagnostic studies replay adjacent updates in a matched setting, using an output-facing subspace to compare realized trajectory signals.

6 Conclusion

  • Adjacent checkpoint footprints expose cross-step output interactions; OSOL uses focus-domain weighting and a drift-ranked token-level residual, without auxiliary supervision or explicit higher-order differentiation.On Qwen3-30B-A3B and Qwen3-8B-Base, domain-macro averages are 0.4822 and 0.3735, with gains of 5.7% and 0.9% over the strongest compared baselines.Source: §6 · p. 10

A Controlled Diagnostic Study Protocols

  • The controlled studies instantiate adjacent updates, preserve optimizer state, and evaluate diagnostics against quantities measured from the realized checkpoint trajectory.One study tests replicate-level signals against downstream cross-domain interference; the other tests token-level signals for identifying rebound.Source: A Controlled Diagnostic Study Protocols · p. 13

A.1 Common model, data, and one-step optimization

  • Both studies use Qwen3-4B and perform diagnostic replay in the final transformer block plus language-model head, a subspace covering 2.51% of parameters.Each replicate uses two consecutive AdamW steps; matched trajectories, rewards, and evaluation targets support comparison among diagnostics.Source: §A.1 · p. 13

A.2 Study I: same-point diagnostics and realized task interference

The controlled studies compare realized two-step output backtracking with same-point diagnostics and test whether checkpoint footprints can identify token rebound before a subsequent update. They also assess finite-difference Hessian proxies under changes in optimizer direction, parameter coverage, and probe scale.

A.2 Study I: same-point diagnostics and realized task interference

  • The 32 replicates use disjoint update, motion-probe, and damage-probe prompts, with the damage target measured as the chat update’s change in fixed-response math-probe NLL.The motion and damage probes are disjoint from the update prompts.Source: §A.2 · pp. 13–14
  • ρs = 0.40 (95% bootstrap interval [0.05, 0.69]) for backtracking mass contrasts with −0.19 for the global diagnostic and −0.22 for the module diagnostic.The ordering persists under every leave-one-replicate-out recomputation.Source: §A.2 · p. 14

A.3 Study II: token-level rebound identification

  • Within negative-drift candidates, 22.6% of eligible tokens enter the candidate set, where mean rebound prevalence is 57.8%.The comparison asks which negative-drift tokens rebound under the subsequent unregulated update.Source: §A.3 · p. 15
  • 0.567 ROC-AUC, 0.627 average precision, and 0.661 top-5% precision for checkpoint-footprint ranking exceed the Hessian proxy’s 0.498, 0.579, and 0.551.The comparison is within negative-drift candidates.Source: §A.3 · p. 15

A.4 Hessian controls and computational measurements

  • Replacing the Hessian proxy’s virtual direction with the realized optimizer displacement raises ROC-AUC from 0.508 to 0.549 and average precision from 0.617 to 0.625, while checkpoint footprint remains highest at 0.557 and 0.636.Source: §A.4 · p. 16
  • Expanding Hessian coverage from 2.51% to 20.07% of parameters raises peak memory from 26.6 to 43.5 GB and ROC-AUC from 0.508 to 0.534; matched checkpoint-footprint ROC-AUC is 0.557 to 0.583.Source: §A.4 · p. 16
  • The finite-difference probe scale changes the output-space displacement evaluated by the Hessian proxy, including the induced policy KL.Source: §A.4 · p. 16

B Additional Diagnostic Evidence

The additional diagnostics compare realized cross-step backtracking with same-point gradient scores, test the stability of that ordering, and examine curvature-proxy sensitivity. The theoretical discussion distinguishes this output-space interaction from objective-Hessian perturbation analyses.

B Additional Diagnostic Evidence

  • Backtracking mass correlates with held-out task damage at ρs = 0.40 (95% bootstrap CI [0.05, 0.69]), while global and module scores are −0.19 and −0.22.The motion and damage measurements use disjoint probes.Source: B Additional Diagnostic Evidence · p. 17
  • Paired bootstrap gains for backtracking mass are +0.59 over the global score (95% CI [0.05, 1.06]) and +0.63 over the module score (95% CI [0.03, 1.15]).The paired comparisons use identical resampled replicate indices.Source: B Additional Diagnostic Evidence · p. 17
  • Across all 32 leave-one-replicate-out subsets, backtracking correlations remain positive while global and module correlations remain negative, preserving their ordering.The analysis does not substitute for an independent replication study.Source: B Additional Diagnostic Evidence · pp. 17–18
  • Changing finite-difference scale ϵ changes both the estimated HVP and the policy-space displacement used for token ranking; checkpoint footprints instead come from realized adjacent checkpoints.Source: B Additional Diagnostic Evidence · p. 18

C.1 Notation and local regularity

  • The local analysis assumes fixed policy support and three-times continuous differentiability in a neighborhood, bounded derivatives, and relevant update segments remaining within that neighborhood.Source: §C.1 · p. 20

C.2 Relation to objective-Hessian perturbation analyses

  • Unlike prior objective-Hessian analyses of a later displacement’s self-quadratic sensitivity, this analysis studies a mixed output-KL interaction under the Fisher geometry, observable through checkpoint log-probabilities.The paper specifies that this is not an identification of the complete objective Hessian with the Fisher.Source: §C.2 · p. 20

C.3 Proof that task-conditioned output KL has Fisher Hessian

The proof sections establish the Fisher–KL identity, derive a local expansion with a cubic-order remainder, and relate adjacent checkpoint footprints to a mixed output interaction.

C.3 Proof that task-conditioned output KL has Fisher Hessian

  • The task-conditioned output KL has the Fisher matrix as its Hessian at the anchor, by the standard Fisher–KL identity.Source: §C.3 · p. 21

C.4 Explicit KL Taylor remainder and mixed derivative

  • The local KL expansion has a cubic-order remainder under a locally bounded third derivative or Hessian Lipschitzness.Source: §C.4 · p. 21

D Checkpoint Observability of Mixed Output Interaction

  • The preceding and current base footprints on the same realized token coordinates multiply to recover the mixed output interaction under the empirical output Fisher, up to cubic local error.Source: D Checkpoint Observability of Mixed… · p. 22

D.3 Proof of checkpoint token–Fisher duality

The appendix links checkpoint token footprints to output-space backtracking and establishes how OSOL’s residual correction controls that component under stated local conditions.

D.3 Proof of checkpoint token–Fisher duality

  • The finite-step approximation proves Theorem 1 with leading error terms of O(η^3) and O(η^4).Equation (69) also states the finite-step form explicitly.Source: §D.3 · p. 23

D.4 Population and empirical Fisher

  • Under anchor-policy conditional sampling, the empirical output Fisher averages to the population Fisher; the deterministic identity remains unchanged for autoregressive trajectories.Source: §D.4 · p. 23

D.5 Path cancellation and token-level backtracking

  • Equation (7) measures net cancellation between two stepwise output motions, while token-level backtracking separates opposing from aligned coordinates.The negative-to-positive rebound component is the one OSOL controls.Source: §D.5 · p. 23

E From Residual Coefficients to Backtracking Control

E.1 Effective coefficient correction after clipping

  • Final coefficient clipping does not reverse the penalty direction; it can only attenuate its magnitude.Source: §E.1 · p. 24

E.2 Proof of the history-conditioned PPO penalty

The section traces OSOL’s residual from a history-conditioned PPO penalty to a local output correction, then explains its support, ranking, and scaling choices.

E.2 Proof of the history-conditioned PPO penalty

  • Gradient descent adds a penalty against increasing the selected token log-ratio, independent of the sign of the base advantage.Source: §E.2 · p. 24

E.3 Base update and correction footprint

  • The base footprint retains focus- and non-focus-domain contributions, while eligible-focus-token corrections use a local model rather than the full clipped PPO–AdamW trajectory.Source: §E.3 · p. 25

E.4 Residual-to-output sign transfer

  • Residual-to-output sign transfer holds automatically for diagonal or weighted-diagonally-dominant selected coordinates; aggregate contraction can tolerate some positive cross-token responses.Source: §E.4 · p. 25

E.5 Proof of first-order backtracking contraction

  • The control gain is nonnegative and strictly positive when a selected coordinate has both negative preceding drift and negative correction footprint.Source: §E.5 · p. 26

E.6 Finite-step control gain

  • The finite-step expression retains the controller’s first-order reduction explicitly and separates it from the cubic deviation due to observing a finite checkpoint step.Source: §E.6 · p. 26

F.1 Why OSOL uses the penalty-only support

  • Penalty-only support avoids history-only positive incentives while retaining a monotone backtracking-reduction guarantee on negative-drift coordinates.Source: §F.1 · p. 27

F.2 Rank-optimal local allocation

  • Rank-optimal allocation assigns larger penalties to more negative drift; quantile ranking avoids raw drift magnitudes and is robust to isolated extremes.Source: §F.2 · p. 28

F.3 Code-faithful residual construction; F.4 Adaptive scale identity and degenerate cases

  • OSOL changes only the existing policy-update coefficient, normalizes residual dispersion to base-advantage dispersion, and falls back to the base objective for empty or degenerate candidate sets.Source: §F.3 · p. 29, §F.4 · p. 29

F.5 Main-Training Reproducibility Details

The training setup rotates a focus domain while retaining all domains in each mixed-domain batch, and uses the preceding checkpoint only to rescore current focus-token coordinates.

F.5 Main-Training Reproducibility Details

  • At each trainer step, the focus domain is sampled uniformly from the four training domains other than the preceding focus, preventing consecutive repeats without fixing an order.Source: §F.5 · p. 29
  • All four domains remain in the mixed-domain batch; the focus domain determines asymmetric base weighting and which response tokens receive history correction.Source: §F.5 · p. 29
  • Each step uses 128 prompts and eight responses per prompt, while checkpoint-derived quantities are fixed coefficients with no gradient through the preceding parameters.This yields 1,024 rollout trajectories before optimization.Source: §F.5 · p. 29

Related papers

Loading 2609.06469v1…