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

2 files changed, 8 insertions, 0 deletions

diff --git a/src/jlens_v3.py b/src/jlens_v3.py
index b5f59ed..69299dc 100644
--- a/src/jlens_v3.py
+++ b/src/jlens_v3.py
@@ -115,6 +115,9 @@ def compute_faithful_jlens(model, layer_idx, batches, device, chunk=32):
             torch.cuda.empty_cache()
 
         n_pairs += B * T * (T + 1) // 2      # valid (t, t' >= t) pairs
+        # NOTE: count = T(T+1)/2 per sequence (all source positions t and all
+        # futures t' >= t). Pairs with t' < t contribute zero gradient by
+        # causality, so summing over all t' and dividing by this count is exact.
         del logits, loss, h_l, h_final, h_final_sum
         torch.cuda.empty_cache()
 
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)