feat: add Geo-Mask, Seq-TIS, Seq-MIS and MIS-PO loss types (sequence-level IS ablations)

[faresobeid]

Adds four loss types to the algo-ablations set (#3381), drawing on RL collapse, part 3 (Li) and Is Frontier Asynchronous RL Solved? (Huang). seq_tis and seq_mis are the truncate-vs-mask pair for sequence-level IS; geo_mask is Li's pure-filter estimator.

All three compute the rollout's summed log-ratio; they differ in what they do with it. Weights are always detached (score-function surrogates), so clipping/masking never silently zeroes gradients through the ratio.

seq_tis — truncated sequence-level IS (Huang)

$$ w = \prod_t \rho_t, \qquad \mathcal{L} = -\sum_{j,t} \min\big(w^{(j)}, C\big), \hat{A}_t^{(j)} \log \pi(y_t^{(j)}) $$

Every rollout keeps a (possibly damped) gradient contribution; bias is one-sided. The estimator Huang's K=12 lag experiments found most durable.

Knob	Default
seq_clip	2.0
adv_tau	1.0

Metrics: mismatch_kl, is_clipped, seq_log_ratio, geo_log_ratio, seq_weight.

seq_mis — masked sequence-level IS, geometric criterion (Li × Huang)

$$ \mathcal{L} = -\sum_{j,t} \mathbb{I}\big(C_{\text{low}} \le (w^{(j)})^{1/T_j} \le C_{\text{high}}\big), w^{(j)} \hat{A}_t^{(j)} \log \pi(y_t^{(j)}) $$

Drops the rollout wholesale unless the geometric mean of its per-token ratios is in the trust region; accepted rollouts keep the untempered product weight, so the estimator is exactly unbiased on the accepted region. The criterion is deliberately geometric rather than a product threshold (which rejects by cumulative divergence and kills all long rollouts); the weight is deliberately untempered (tempering to $w^{1/T}$ is the bias Huang shows makes GeoMean-IS track token-IS). Numerics guard caps the weight at $e^{40}$. The mask-vs-truncate contrast with seq_tis: no mis-weighted gradients from the divergent tail, but effective batch shrinks as policies drift (watch is_masked).

Knob	Default
geo_mask_low / geo_mask_high	0.5 / 2.0
adv_tau	1.0

Metrics: masked/unmasked_mismatch_kl, is_masked, seq_log_ratio, geo_log_ratio, seq_weight.

geo_mask — geometric trust region, pure filter (Li)

$$ \mathcal{L} = -\sum_{j,t} \mathbb{I}\big(C_{\text{low}} \le \rho_{\text{geo}}^{(j)} \le C_{\text{high}}\big), \hat{A}_t^{(j)} \log \pi(y_t^{(j)}), \qquad \rho_{\text{geo}} = \Big(\prod_t \rho_t\Big)^{1/T} $$

Same length-invariant acceptance rule as seq_mis, but accepted rollouts use the plain score function with no importance weighting (the post's pure-filter design: inside the trust region the ratio is ~1). token_clip enables the post's Geo-Mask-Token-TIS hybrid (score function weighted by the detached, upper-clipped per-token ratio) for full per-token IS correction.

Knob	Default
geo_mask_low / geo_mask_high	0.5 / 2.0
adv_tau	1.0
token_clip	None

Metrics: masked/unmasked_mismatch_kl, is_masked, geo_log_ratio.

mis_po — dual-level binary masking, pure score function (Step 3.5 Flash)

From Step 3.5 Flash (eq. 2), "Metropolis Independence Sampling-filtered Policy Optimization": treat the inference policy as a proposal distribution and keep only samples close enough to the trainer, as effectively on-policy — no importance weighting anywhere:

$$ \mathcal{L} = -\sum_{j,t} \mathbb{I}\big(\rho_{\text{tok,low}} \le \rho_t^{(j)} \le \rho_{\text{tok,high}}\big), \mathbb{I}\big(\rho_{\text{geo,low}} \le \bar\rho^{(j)} \le \rho_{\text{geo,high}}\big), \hat{A}_t^{(j)} \log \pi(y_t^{(j)}) $$

Token-level indicator drops individually mismatched tokens; trajectory-level indicator drops the rollout when its geometric-mean ratio drifts. Structurally it is geo_mask's trajectory filter + a token-level mismatch filter, with the paper's much tighter trajectory band (their ratios measure pre-update-snapshot/vLLM mismatch, which concentrates near 1 at 128k context — under real policy lag the band should be widened; docs note this).

Knob	Default
token_mask_low / token_mask_high	0.5 / 2.0 (paper)
geo_mask_low / geo_mask_high	0.996 / 1.001 (paper)
adv_tau	1.0

Metrics: masked/unmasked_mismatch_kl, is_masked, is_token_masked, is_traj_masked, geo_log_ratio.

Test plan

Smoke-tested all four loss fns directly (no test files added, per repo convention): exact gradients in every regime (score function with per-token drops for mis_po; score function -adv; clipped weight -2·adv; untempered in-region weight -e^{0.8}·adv; token_clip hybrid's detached weight), both-sided geo rejection with zero gradient while staying backward-able, masks computed only over loss-masked tokens, the $e^{40}$ overflow guard (finite loss/grads at summed log-ratio 800), and clean zero loss on empty masks. ruff check / ruff format clean; all configs parse through the LossConfig discriminated union.

🤖 Generated with Claude Code