Shared-prefix training accelerates reinforcement learning for hybrid-attention agents

What they did

The authors designed HARTS for reinforcement-learning workloads in which many agent trajectories share an interaction history before branching. Rather than train every root-to-leaf trajectory separately, HARTS compresses common prefixes, measures work using unique non-replayed token rows, and jointly plans microbatches, replica assignments, and execution slots.

For hybrid-attention models, the system packs all branches in a round into one call and passes differentiable states between segments. Its chunkwise linear-attention algorithm recovers and selectively replays boundary states, supports activation recomputation, and restores per-token log-probabilities. It also corrects token weights and load statistics for deterministic, no-token-drop top-k mixture-of-experts routing.

Key findings

Across multiple parallel configurations on an agentic RL workload generated from SWE-bench tasks, HARTS achieved 4.81–4.87× speedups for forward, backward, and gradient computation with activation recomputation enabled.
Under the paper’s packed chunkwise execution model, the scheduling algorithm runs in linear time and minimizes the number of sequential linear-attention calls.
Prefix sharing avoids repeating projection, MLP or mixture-of-experts computation, and final outputs; additional work is limited to bounded state replay needed for numerical alignment.
Numerical differences were comparable to variation from rerunning the baseline, while reward trends remained similar during the first 120 steps of τ³-Bench training.

Why it matters

Tree-based execution has mainly been developed for full-attention models, leaving hybrid architectures unable to exploit shared agent histories during training. HARTS shows that prefix compression can extend to chunkwise linear attention without breaking gradient flow, activation recomputation, or existing RL-objective interfaces, potentially lowering the cost of training agents on branching, long-horizon interactions.

Caveats

The reported evidence covers a real hybrid-attention model but only the workloads, parallel configurations, and early 120-step training window described by the authors. Longer-run learning equivalence, gains on other model families and rollout distributions, and behavior under routing schemes beyond deterministic no-token-drop top-k MoE remain open; the minimum-call guarantee also applies specifically to HARTS’s packed execution model.