diff options
Diffstat (limited to 'src')
| -rw-r--r-- | src/jlens_v3.py | 3 | ||||
| -rw-r--r-- | src/loss_reweight.py | 5 |
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) |