(prototype) Move the NT-Xent equation into lightly.functional - #2035
(prototype) Move the NT-Xent equation into lightly.functional#2035gabrielfruet wants to merge 1 commit into
Conversation
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 69200af70c
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| logits_11 = torch.einsum("nc,mc->nm", z1, z1_all) / temperature | ||
|
|
||
| logits_00 = logits_00[~diagonal].view(batch_size, -1) | ||
| logits_11 = logits_11[~diagonal].view(batch_size, -1) |
There was a problem hiding this comment.
When a caller supplies only one of z0_all or z1_all—which the independently optional arguments and documented defaults permit—the gathered tensor can have B * W rows while the other defaults to B. The diagonal mask is sized exclusively from z0_all, so applying it to logits_11 raises a shape-mismatch IndexError for either one-sided case. Either reject calls that do not provide both gathered views or build the self-similarity masks independently.
Useful? React with 👍 / 👎.
gabrielfruet
changed the title
1 participant
3 of 7 in a stack. Base: #2034.
NTXentLoss.forwardholds the memory bank, the distributed gather and the equation in one 75-line method. This moves the equation out tolightly/functional/ntxent.pyas a pure function, and leaves the module holding the state. No maths is duplicated: the module calls the function.The gather stays in the module, which is where it belongs, and reaches the function as
z0_all,z1_allandpositives_at. On one rank those default to the local tensors and the mask istorch.eye, which is what the old code built.DistributedKindis the other half, and it is a declaration rather than a mechanism today. DDP averages the per-rank loss and gradient, and whether that is correct depends on the loss.torch.distributed.all_reducehas no backward, so SIGReg shipped a correct forward with a gradient off by exactly1/world_size(#1920), and the same bug came back againstBarlowTwinsLoss(#1977). Both were fixed by hand and the fixes are comments.NTXentLossnow declaresGATHER_FOR_NEGATIVES. Nothing dispatches on the kind yet; the test selector belongs with its own PR.Testing:
pytest tests/loss/test_ntx_ent_loss.pyunchanged at 246 cases, plus golden values pinningntxentat seed 11.