PaperlayerView exact revision on arXiv

Source-linked AI summary

Dynamic Latent Routing

Fangyuan Yu, Xin Su, Amir Abdullah

arXiv:2605.14323v1cs.LGcs.AIcs.CL

TL;DR

The paper asks whether temporal composition of sub-policies can be understood and exploited when rewards vary over time. It introduces GDS for theoretically guaranteed policy composition and DLR for single-stage language-model post-training with searched discrete codes. In low-data fine-tuning, DLR outperforms SFT by a mean of +6.6 percentage points and learns causally specialized routing codes.

  • Problem

    Hierarchical reinforcement learning uses sub-policies, but the theoretical justification for temporal composition under time-varying rewards remains limited.

  • Method

    The paper introduces GDS and DLR, which applies search, selection, and joint updating to learn discrete codes, routing, and model parameters while steering residual-stream representations.

  • Results

    +6.6 percentage points is DLR’s mean gain over SFT across 24 model–dataset settings in low-data fine-tuning.

  • Takeaways & Limitations

    DLR learns discrete codes that specialize to distinct subtasks and have causal effects without subtask labels.

  • Takeaways & Limitations

    The theoretical formulation assumes nonpositive time-varying rewards, with rewards interpreted as negated cost after a shift.

Abstract

from arXiv · show

We investigate the temporal concatenation of sub-policies in Markov Decision Processes (MDP) with time-varying reward functions. We introduce General Dijkstra Search (GDS), and prove that globally optimal goal-reaching policies can be recovered through temporal composition of intermediate optimal sub-policies. Motivated by the "search, select, update" principle underlying GDS, we propose Dynamic Latent Routing (DLR), a language-model post-training method that jointly learns discrete latent codes, routing policies, and model parameters through dynamic search in a single training stage. In low-data fine-tuning settings, DLR matches or outperforms supervised fine-tuning across four datasets and six models, achieving a mean gain of +6.6 percentage points, while prior discrete-latent baselines consistently underperform SFT. Mechanistic analyses and targeted code ablations show that DLR learns structured routing behaviors with distinct causal roles.

1 Introduction

The paper develops a theoretical and language-modeling framework for composing discrete sub-policies over time. GDS provides the guarantee, while DLR adapts its search–select–update principle to jointly learn latent codes and routing.

  • Hierarchical reinforcement learning expands action spaces with sub-policies, but lacks a theoretical account of why temporal composition helps.
  • Time-varying rewards make static policies generally suboptimal, motivating GDS to search over concatenated sub-policies with optimal goal-reaching guarantees.
  • DLR treats unsupervised discrete codes as an action space whose sequences condition language-model likelihood and must therefore be explored.
  • DLR injects codes as residual-stream steering vectors rather than extra tokens, avoiding disruption to pretrained sequence structure in low-data fine-tuning.
  • DLR jointly learns code use, code effects, and composition in one stage through guided search, selection by conditional likelihood, and parameter updates.
  • +6.6 percentage points is DLR’s mean gain over SFT across 24 model–dataset settings, while learned codes specialize to subtasks and have causal effects without subtask labels.

2 Related Work

Related work covers discrete latent tokens, hidden-state steering, hierarchical reinforcement learning, and mechanistic interpretability. The paper distinguishes DLR through dynamic temporal composition, single-stage learning, and explicit discrete routing.

  • Prior discrete-code methods commonly inject extra tokens and use warm-up, pre-labeling, or fixed-code stages for post-training.
  • DLR instead learns dynamic steering codes jointly with the language model, contrasting with methods that apply a single fixed steering vector.
  • Hierarchical reinforcement learning introduces high-level actions, but prior frameworks lack a theoretical justification for temporal sub-policy composition.
  • Unlike related compositional methods that assume stationary rewards or compose non-temporally, GDS proves optimal goal-reaching policies arise from temporal concatenation in dynamic MDPs.
  • DLR’s explicit routing symbols expose intermediate computation for causal interventions and mechanistic analysis rather than only post-training probing.

3 Theoretical Foundation

The theoretical foundation introduces dynamic MDPs with time-varying rewards and formalizes time-varying policies, concatenation, goal-reaching, and GDS. It establishes value decomposition and GDS optimality under stated assumptions.

  • A dynamic MDP uses a finite horizon, time-varying rewards, time-homogeneous transitions, discounting, and a nonpositive-reward convention interpretable as negated cost.
  • Dynamic policies are sequences of decision rules indexed by time, so the action selected from a state may vary across the horizon.
  • Policy concatenation joins time-varying policies with successive horizons, enabling temporal composition of intermediate sub-policies.
  • The value of a concatenated policy is formally decomposed under the paper’s assumptions.
  • GDS searches policy concatenations with a priority queue, pruning dominated candidates and stopping when a feasible goal-reaching policy is selected.
  • Under Assumptions A1–A3, GDS finds an optimal goal-reaching policy for every reachable goal set.

4 Methodology

DLR translates GDS into a language-model training procedure: discrete codes steer token chunks through the residual stream, while search, selection, and joint updating optimize routing and model behavior.

  • DLR replaces GDS’s priority queue with a learned policy head while retaining the search–select–update structure for latent-code optimization.
  • Each discrete code steers a chunk of K consecutive tokens at a chosen layer, yielding M = ⌈T/K⌉ codes for a sequence of length T.
  • A routing head assigns each chunk a codebook index, whose learned vector is added to every hidden state in that chunk at the injection layer.
  • The objective combines unsteered language-model likelihood, information gain from steering, and entropy terms that encourage diverse codes with lower local uncertainty.
  • At each training step, DLR samples code sequences, selects the sequence maximizing pθ(x | a), and optimizes the resulting objective.
  • Joint updates are necessary because the code-sequence reward changes as the model changes, coupling routing quality and model likelihood.

5 Experiments

DLR is evaluated against supervised fine-tuning and discrete-latent baselines in matched low-data settings, where it improves accuracy across models and benchmarks. Ablations support the importance of dynamic search and codebook diversity.

  • Evaluation setting: DLR adds steering vectors to hidden states, preserving the input sequence during single-epoch, low-data fine-tuning.This avoids the sequence disruption associated with inserting new tokens.
  • Main results: DLR wins every model-dataset cell, while TokenAssorted collapses on math and reasoning tasks at small scale.For example, TokenAssorted scores 15.7 versus SFT’s 46.0 on GSM8K-Qwen3-0.6B.
  • Ablations: Removing any DLR objective term degrades accuracy or triggers codebook collapse, validating the search-select-update design.Accuracy worsens when rollouts are non-diverse, search is removed, or the policy head is untrained.

6 Analysis

DLR learns routing that is diverse, context-dependent, and causally important. The arithmetic case study further shows that discrete codes can expose and support intervention on task-relevant computation.

  • ScienceQA routing: DLR learns a diverse codebook: steering-vector cosine similarity is at most 0.28, while code utilization ranges from 31% to 100%.These measurements indicate non-collapsed codebooks across the evaluated models.
  • ScienceQA routing: Code n-grams are more topic-pure than the 0.17 random-router baseline, with purity increasing for longer n-grams.Topic labels were unavailable during training, so the structure emerged from the DLR objective.
  • Causal ablations: Global ablations reduce SciQA accuracy by 6.2–17.4 pp when steering is zeroed and by 4.8–11.7 pp when routed codes are randomized.Single-code ablations also produce mean drops of 0.71–2.74 pp across four tested codes per model.
  • Causal ablations: Single-code effects are topic-specific: ablating code 0 improves biology by +3.6 pp while reducing physics by −4.0 pp on Qwen3-0.6B.Other improvements include writing-strategies by +4.8 pp and chemistry by +9.8 pp.
  • Six-digit arithmetic: In six-digit arithmetic, DLR assigns named codes to carry, sum-9 boundary, and trivial positions, making the carry structure readable from the code sequence.The example uses t2 for carry positions and t6 at the sum-9 boundary.
  • Six-digit arithmetic: The arithmetic codebook specializes without subtask labels: most active codes have a dominant subtask accounting for at least 70% of occurrences.Removing all codes collapses accuracy from 95.5% to 0.1%.
  • Six-digit arithmetic: Swapping one code at one answer position fixes wrong predictions on carry-heavy examples at a 27–31% rate.Cross-operation transplantation reaches 93.5% versus a 75.5% random baseline.
  • Six-digit arithmetic: DLR outperforms SFT in 12 of 13 tested configurations and in all 13 hardest six-deep carry-cascade configurations.The largest reported gains reach +50 pp.

7 Conclusion

The paper combines a proof of optimal temporal policy composition with DLR, a single-stage method for learning and routing discrete latent codes. Experiments show gains over SFT and structured, causally load-bearing codes.

  • Theory: GDS proves that optimal goal-reaching policies in dynamic MDPs can be recovered by concatenating sub-policies.The result concerns MDPs with time-varying rewards.
  • Method: DLR applies GDS’s search-select-update principle to jointly learn, search, and concatenate latent codes in one training phase.The codes steer the model through the residual stream.
  • Empirical conclusion: +6.6 pp is DLR’s reported advantage over SFT in low-data fine-tuning versus other discrete-latent baselines.The conclusion states that DLR uniquely outperforms SFT in this regime.
  • Empirical conclusion: Analyses show that DLR codes specialize to distinct subtasks and are causally load-bearing without subtask labels during training.The paper discusses limitations in Appendix I.8.

A Theory for Dynamic MDP and Existence of Optimal Policy

The paper develops dynamic MDP theory for time-varying policies, establishing Bellman equations, policy improvement, convergence, and existence of optimal policies. It also shows that static policies can be suboptimal and that GDS recovers optimal goal-reaching policies.

  • Dynamic MDP foundations: Dynamic value and Q-functions are defined for time-indexed policies in finite-horizon MDPs.The corresponding Bellman equations characterize these functions and also apply to truncated policies.
  • Policy improvement and convergence: The time-varying Bellman policy operator improves values and converges from any initial policy to an optimal policy.The improvement is established across states and times, with convergence obtained through repeated policy updates.
  • Existence of optimal policies: An optimal time-varying policy exists under the paper’s stated compactness, finite-action, bounded-reward, and continuity assumptions.The theory also establishes existence and uniqueness of the optimal value function.
  • Dynamic versus static policies: Time-varying policies can achieve strictly higher value than every static policy in a finite-horizon dynamic MDP.Thus, allowing policies to depend on time can strictly improve the optimal value in this setting.
  • General Dijkstra Search: GDS searches over temporal concatenations of sub-policies and is proved to find an optimal goal-reaching policy.This search-based guarantee differs from Bellman iteration’s per-state policy refinement.

B Policy Composition

The policy-composition analysis formalizes how General Dijkstra Search prioritizes, extends, and prunes temporal policy prefixes. Its queue and domination arguments show that the resulting popped policies are optimal for the goal sets they reach.

  • Policy representation and search: GDS represents candidate solutions as temporal policy prefixes paired with values and reachable goal sets.Queue elements are extended one step at a time, while policies are classified by whether they are popped, queued, added later, or never queued.
  • Queue ordering: A popped policy has value at least as large as policies currently in the queue or to be added later.This follows from priority-queue ordering and continuation-value bounds for policy extensions.
  • Pruning and domination: Policies omitted from the queue are dominated by popped, unskipped policies with no worse value up to an epsilon gap.For optimal reach, the dominating policy reaches a subset of the omitted policy’s goals.
  • Optimality of composition: For optimal reach, a popped and unskipped policy is optimal among all policies it dominates.The proof compares policies across queue states, future additions, and policies that never appear in the queue.
  • Dual coverage analysis: The same search logic supports optimal coverage, with the direction of goal-set inclusion reversed.A popped and unskipped policy is optimal among policies that dominate it, and policies omitted from the queue are dominated by retained policies.

C Symmetric goal-covering result

The paper gives a coverage-side counterpart to goal reaching: GDS searches for policies whose reachable goal sets contain the target and proves optimality for every coverable goal set.

  • Coverage formulation: Covering policies are defined by reachable goal sets that contain the target goal set.This reverses the inclusion direction used in the goal-reaching formulation.
  • Coverage search: The coverage algorithm initializes a priority queue, repeatedly pops the highest-value candidate, records feasible reachable sets, prunes dominated candidates, and extends retained policies.The algorithm continues until the queue is empty or the target is covered.
  • Coverage optimality: For every coverable goal set, General Dijkstra Search for Optimal Coverage finds an optimal goal-covering policy.The result applies under Assumptions A1–A3.
  • Proof symmetry: The coverage proof mirrors the reaching proof using a coverage-side queue invariant, pruning guarantee, and reversed inclusion direction.The first popped feasible covering policy is optimal among policies covering the same goal set.

E Loss-term ablations

The ablations show that each major component of DLR’s objective contributes to performance, with policy optimization and codebook regularization especially important.

  • Policy optimization: −9.8 pp on GSM8K and −9.0 pp on ScienceQA are the average degradations after removing the policy-optimization term.The largest individual drops range from −5.4 to −23.7 pp on GSM8K and −6.5 to −13.1 pp on ScienceQA.
  • Generalist and information gain: −1.8 pp on CSQA, −3.96 pp on GSM8K, −4.78 pp on ScienceQA, and −3.66 pp on StrategyQA are the average drops from the specialist-loss ablation.No single cell exceeds −6.4 pp; the ablation removes both the unsteered LM term and the information-gain contrast.
  • Codebook regularization: ∼0–8 pp are lost across the four benchmarks when the bigram-Zipfian regularizer is removed.The largest hits occur on ScienceQA and StrategyQA.
  • Codebook regularization: Near-1/32 code utilization after removing the regularizer indicates collapse onto essentially one routing code.The prior maintains codebook diversity and discourages a single steering direction from carrying nearly all routing mass.

F Hyperparameter sweeps

The sweeps identify search and routing settings that balance candidate diversity, codebook capacity, abstraction, and regularization across model families.

  • Search budget: −5.0 pp is the average cost of reducing the rollout budget from N=4 to N=1, while N=8 changes performance by −0.5 pp on average.N=4 is therefore the lowest-compute setting reported to capture the search benefit.
  • Sampling temperature: −6.9 pp at τ=0 and −4.1 pp at τ=2 show that both zero and excessive sampling temperature underperform τ=1.The results support using some randomness without flattening rollouts toward uniform sampling.
  • Codebook size: C=32 is the codebook-size plateau: C=1 trails it by −3.6 pp on average, while C=64 is within 0.5 pp on every cell.Accuracy increases monotonically from C=1 through C=32 before plateauing.
  • Abstraction ratio: K=4 is the abstraction-ratio sweet spot, while K=1 loses −3.7 pp and K=2 loses −1.5 pp on average.K controls how many consecutive natural-language tokens share a routed code.
  • Regularization weight: −4.1 pp is the average cost of setting αreg=1.0, whereas αreg=0.01 differs from the default by only −0.7 pp.An over-strong prior pushes routes toward the Zipfian reference rather than task-relevant partitions.
  • Policy weight: Qwen backbones prefer αpolicy=0.5, while Llama-3.2-1B prefers αpolicy=0.1.Under-weighting costs −2.3 pp on average, and over-weighting costs −3.6 pp on average.

G Global ablation: full results

Global perturbations across all six checkpoints show that learned routing contributes directly to ScienceQA accuracy across model families.

  • Global perturbations: Zeroing the steering scale or replacing routed codes with random alternatives reduces ScienceQA accuracy across all six checkpoints.The result extends the global perturbation analysis across both Qwen and Llama checkpoints.

H Single-code ablation case studies (full traces)

Single-code ablations produce targeted answer changes across genotype, writing, and chemistry examples, illustrating distinct routing effects on model behavior.

  • Genotype reasoning: Code 0 ablation changes genotype responses from phenotype descriptions to symbolic answers such as WW or ff.The case studies show the altered answer can become correct after the code is removed.
  • Trace comparison: 0–65% of natural-language output can diverge while code-sequence divergence remains 0–8% in the reported traces.The examples therefore show substantial response changes from sparse routing-code changes.
  • Writing strategies: Code 0 ablation switches writing responses toward identifying the flawed sentence or vague pronoun reference rather than selecting a clean alternative.The traces include examples involving unclear pronoun antecedents and ambiguous references.
  • Chemical reasoning: Code 0 ablation changes a chemistry classification from appearance-based physical reasoning toward the underlying chemical process.The chemistry example distinguishes shell growth as calcium-carbonate deposition from merely apparent growth.
  • Overall performance: DLR wins consistently across architectures and data sizes, with larger margins on harder data splits.This full-results summary accompanies the global ablation case studies.

I.2 Causal ablations

DLR’s causal ablations show that routing codes are necessary for performance and encode structured, position-specific computation. The learned codebook supports both interpretable specialization and targeted correction, within a study whose theorem and empirical scope remain bounded.

  • Causal necessity: 95.5% to 0.1% overall accuracy after knockout shows that DLR’s abstraction codes are causally necessary.Hardest borrow-cascade splits also fall to ≤3% under knockout.
  • Structured routing: 23 of 30 codebook codes specialize in one or two Quirke subtasks and are tied predominantly to one or two answer positions.Most active codes have dominant-subtask purity of at least 70%.
  • Targeted intervention: 27–31% of mispredicted carry-heavy examples are fixed by replacing a single abstraction code, without weight updates or activation access.The best swap fixes 10 wrong predictions while breaking 5 correct ones.
  • Mechanistic interpretation: DLR recovers carry-state regimes as readable code clusters without ground-truth circuit labels or activation-level analysis.The clusters correspond to uncertain, generated-carry, and simple-digit regimes across addition and subtraction.
  • Codebook heterogeneity: Polysemantic fallback codes coexist with high-purity specialists, with t21 reaching 94% purity for one subtask while t1 reaches only 24%.The fallback code spans all five answer positions, unlike the specialist code’s single-position behavior.
  • Limitations: DLR’s GDS connection motivates the empirical method but does not certify it because DLR’s reward depends on model parameters.Experiments cover 24 configurations across two open-weight families and a low-data post-training setting.
Loading 2605.14323v1…