Teacher-forced (parallel) TB loss recompute for recurrent estimators

[josephdviviano]

Summary

Stacks on #525 (fixed-length KV-cache). With use_kv_cache=True, sampling runs under no_grad, so gradients must come from a loss-time recompute. torchgfn's existing recompute for recurrent estimators (is_vectorized=False) re-runs the estimator per step, re-paying the sequential rollout cost inside the loss.

This PR adds an opt-in teacher-forced (parallel) recompute: one causal forward over the whole committed trajectory [BOS, w0, …, w_{L-1}], gathering per-step action log-probs. It reuses the estimator's own full-prefix forward, so the loss-time input is built identically to sampling; the result matches the per-step path to tolerance (atol=1e-5) with a single grad-bearing forward and a shallow autograd graph.

Changes

• RecurrentDiscretePolicyEstimator(teacher_forced_loss=…) — recompute PF log-probs with one parallel forward per trajectory instead of the per-step loop. Auto-enabled by use_kv_cache, so a cached policy stays trainable out of the box (you only pass recalculate_all_logprobs=True at loss time).
• get_trajectory_pfs takes a teacher-forced branch for recurrent estimators; the per-step / vectorized paths are unchanged (their shared mask + scatter bookkeeping is now factored into small helpers). Fixed-length trajectories only (guarded; see the Gap-2 TODO).
• Fail-fast guard: get_trajectory_pfs now raises an actionable error when asked to reuse detached (no_grad-sampled) log-probs under an active autograd graph, instead of the previous opaque backward() failure.

Result

Together with use_kv_cache=True (recalculate_all_logprobs=True), this turns the sampling win into an end-to-end training speedup.

Measured on this branch — single CPU run (Apple Silicon), batch 32, dim 64 / 4 heads / 2 layers, word_size=1:

seq_len	sample base	sample cached	sampling ×	train base	train cached	training ×	|Δ log_pf|
8	15.36 ms	10.62 ms	1.4×	38.92 ms	17.53 ms	2.2×	2.4e-07
16	38.27 ms	19.99 ms	1.9×	93.00 ms	33.02 ms	2.8×	3.6e-07
32	120.12 ms	44.14 ms	2.7×	268.07 ms	59.86 ms	4.5×	3.6e-07
64	424.81 ms	102.28 ms	4.2×	899.33 ms	131.55 ms	6.8×	3.6e-07

The |Δ log_pf| residual is the equivalence gate: both configurations are the same model, so the wall-clock gap is the optimization. Absolute figures are CPU- and machine-dependent and vary run to run; the robust claim is the trend — the training-step gap widens with sequence length, because the cached rollout is O(L²) instead of O(L³) and the loss recompute collapses to one parallel forward.

Tests

testing/test_kv_cache.py: teacher-forced log_pf == per-step recompute; the teacher-forced path carries gradients; use_kv_cache auto-enables teacher_forced_loss; and the detached-log-probs guard fires with an actionable message.

Unaffected

Non-recurrent estimators and the default (teacher_forced_loss=False) path are unchanged. Stateless policies and the Sampler API are unaffected. Full testing/ suite green on this branch.

Review note: this PR is stacked on #525, so its diff is shown against feat/kv-cache. Please merge #525 first (GitHub will then retarget this PR to master).

🤖 Generated with Claude Code

[codecov]

Codecov Report

❌ Patch coverage is 91.42857% with 3 lines in your changes missing coverage. Please review.
✅ Project coverage is 74.54%. Comparing base (d9cfd34) to head (25c24c2).

Files with missing lines	Patch %	Lines
src/gfn/utils/prob_calculations.py	91.17%	2 Missing and 1 partial ⚠️

Additional details and impacted files

@@                Coverage Diff                @@
##           feat/kv-cache     #526      +/-   ##
=================================================
+ Coverage          74.49%   74.54%   +0.04%     
=================================================
  Files                 59       59              
  Lines              10692    10720      +28     
  Branches            1558     1561       +3     
=================================================
+ Hits                7965     7991      +26     
- Misses              2222     2223       +1     
- Partials             505      506       +1     

Files with missing lines	Coverage Δ
src/gfn/estimators.py	86.14% <100.00%> (+0.02%)	⬆️
src/gfn/utils/prob_calculations.py	83.15% <91.17%> (+1.56%)	⬆️

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.