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

1 files changed, 3 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()