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