Microsoft/post-training

2 posts

microsoft

Engineering and algorithmic interventions for multimodal post-training at Microsoft scale (opens in new tab)

At production scale, post-training multimodal agents fail for reasons that standard reinforcement-learning literature often overlooks. Heterogeneous tasks, long tool-use trajectories, noisy reward sources, and strict latency and safety requirements can make aggregate reward look healthy while the policy gradient becomes uninformative and important capabilities regress. The post presents interventions designed to preserve useful advantage signals as scale, task diversity, and interaction horizons grow. ## Production-Scale Challenges - Copilot agents must simultaneously handle: - Tool orchestration - Enterprise documents and mixed-media inputs - Content moderation - Multi-step execution - Trajectories range from roughly 100 to more than 2,000 tokens and span 6 to 25+ interaction steps. - Rewards come from programmatic checks, human judgments, and implicit usage signals, each with different noise and latency. - A single scalar reward can hide regressions in robustness, long-horizon planning, or downstream task success. - Aggregate reward may rise while gradient updates increasingly depend on a small, unrepresentative subset of trajectories. ## Staged Objective Curriculum - The team separates: - **Verifiable objectives**, such as tool syntax and format compliance - **Preference objectives**, such as tool choice and response quality - Training uses only verifiable objectives during the first 30%. - Preference signals are then introduced linearly. - An entropy floor, implemented through a KL penalty activated below a threshold, prevents premature policy collapse. - Entropy bonuses were insufficient because the issue was not simply exploration; optimization was favoring behaviors that were easy to score. - A 30% warmup worked better than 10% or 50% across task families. - Early text-only supervision could also activate multimodal capabilities more reliably than noisy direct multimodal supervision, assuming adequate cross-modal alignment from pretraining. ## Adaptive Curriculum Based on Estimator Health - The team monitors effective sample size (ESS): `ESS = (Σ wᵢ)² / Σ wᵢ²` - ESS measures how many trajectories meaningfully contribute after importance weighting. - ESS falling below 20% of nominal batch size predicted learning stalls by about 35 epochs. - When ESS drops, the system: - Injects near-miss trajectories from a reservoir buffer - Temporarily increases the KL penalty to limit policy drift - Near-misses worked better than hard negatives because they preserve useful distinctions near the decision boundary instead of merely pushing the policy away from failure. - The intervention maintained ESS above 70%, with approximately 15% additional memory usage. ## Variance-Corrected Normalization - Per-task gradient normalization balances task magnitudes but ignores variance within each task. - Broad categories such as “coding” may contain trajectories ranging from 100 to 2,000 tokens, with very different variance. - Importance weighting can cause long trajectories to dominate the effective gradient even after task-level normalization. - The excerpt ends while introducing the team’s variance-correction approach, so its implementation and results are not included here. The central recommendation is to treat estimator health—not just reward and task metrics—as a first-class training signal. Monitoring ESS, controlling objective timing, and accounting for trajectory variance can help prevent healthy-looking dashboards from masking policy collapse and capability regressions.

microsoft

Diagnosing instability in production-scale agent reinforcement learning (opens in new tab)

Hugging Face has integrated its Post-Training Toolkit into TRL, bringing production-ready diagnostics to reinforcement learning and agent post-training pipelines. The work identifies a late-phase instability specific to tool-using, on-policy agents: variance can grow in post-tool contexts even while loss, reward, entropy, and global KL remain stable. Targeted tail, distributional, and effective-sample-size diagnostics can expose this failure before it becomes divergence. ## Production Monitoring for Long-Running Agents - Modern agent training runs over long horizons, uses external tools, and adapts continuously. - Failures often develop gradually rather than appearing as a single catastrophic event. - Standard aggregate metrics can hide rare but increasingly severe updates. - The proposed monitoring approach: - Computes diagnostics in-stream. - Separates text-only and post-tool interactions. - Aggregates statistics across workers. - Uses lightweight rolling windows and percentile tracking at fixed intervals. ## Tool-Conditioned Variance Amplification - Tool calls expand the state space through external transitions, exposing the policy to contexts it may rarely encounter in the reference distribution. - Training states can be modeled as: `d(s) = (1−α)·d_text(s) + α·d_tool(s)` - As the proportion of tool-conditioned states, `α`, increases, more updates occur where the reference policy assigns low probability to sampled actions. - This causes importance-weighted updates to develop increasingly large tails. - The mechanism is distinct from global entropy collapse or optimizer instability, though those factors may interact with it. ## Minimal Reproduction and Tail Diagnostics - A small on-policy experiment with an instruction-tuned open-weight model reproduced the pattern. - The 95th percentile of absolute per-token log-ratios, `|r|`, was tracked separately for text-only and post-tool contexts. - Findings included: - Text-only tail magnitudes remained stable or declined. - Post-tool tails grew steadily under fixed-policy baselines. - Drift-aware training substantially reduced tail growth. - Constraining tool outputs also suppressed the effect. - Aggregate loss, reward, and entropy remained stable while the tail was expanding. ## Distributional Shift in the Right Tail - Empirical CDFs across early, middle, and late training showed a change in distribution shape rather than a simple threshold crossing. - In tool-conditioned contexts: - The right tail flattened and stretched. - More probability mass moved toward high-magnitude updates. - Drift-aware baselines muted or reversed the shift. - This supports a distributional explanation rather than an artifact of choosing a particular percentile. ## Importance Ratios and Effective Sample Size - For ratio-based on-policy objectives, gradient variance is related to: `Var[ĝ] ∝ E[(π_θ(a|s) / π_ref(a|s))²]` - When `π_ref(a|s)` is small in tool-conditioned states, a small number of updates can dominate the estimator. - Larger batches and improved baselines may reduce noise but do not fix poor support overlap. - Effective sample size (ESS) provides a supporting signal: - ESS declines as importance weights become concentrated. - It is sensitive to window size and batch structure. - Its trends align with post-tool tail growth, but absolute values should not be over-interpreted. ## Delayed Failure and Misdiagnosis - Instability appears first in tool-conditioned contexts and may remain invisible in global metrics for a long time. - By the time aggregate metrics change, substantial variance amplification may already have accumulated. - The problem is often incorrectly attributed solely to optimizer behavior or inadequate global variance reduction. - Such interventions may delay failure without addressing the underlying support mismatch. - The mechanism is less pronounced when tool outputs are tightly constrained, policies are effectively frozen after tool calls, or interaction diversity plateaus early. The practical recommendation is to add slice-aware, tail-focused diagnostics to production TRL pipelines. Monitoring post-tool log-ratio percentiles, distributional changes, and supporting ESS trends can provide an early warning system for instability that global loss, reward, entropy, and KL metrics miss.