Document n_pairs denominator (exact count) and ctrl_random gradient-mass caveat (comment-only)

1 files changed, 5 insertions, 0 deletions

diff --git a/src/loss_reweight.py b/src/loss_reweight.py
index 1bf3124..f01d88a 100644
--- a/src/loss_reweight.py
+++ b/src/loss_reweight.py
@@ -53,6 +53,11 @@ def _weighted_loss(logits, y, mode, q_id, batch_k, V, device):
     if mode == 'q':
         w[y.view(-1) == q_id] = WEIGHT
     elif mode == 'ctrl_random':
+        # Same number of upweighted positions as the q-mode model, but on
+        # random non-q targets. NOTE: gradient-magnitude distribution differs
+        # from q-mode (random positions spread across the batch vs rare 'q'
+        # positions); the design controls for "any 2x reweighting changes the
+        # model", not for exact gradient-mass matching.
         g = torch.Generator().manual_seed(1000 + batch_k)  # CPU generator (randperm)
         n_q = int((y == q_id).sum().item())
         flat = torch.arange(y.numel(), device=device)