fix(training): align contrastive pretraining anchors with their pairs

[ChrisW09]

Fixes #414

Problem

In ContrastivePretrainer.contrastive_loss (both the pooled and unpooled branches) the pair tensor is built as

all_pairs = torch.cat(pairs, dim=0)   # [positives of a0..aN | negatives of a0..aN]

but the anchor tensor is

embeddings_s = embs.repeat_interleave(k_neighbors * len(pairs), dim=0)   # [a0 x 2k, a1 x 2k, ...]

These layouts do not line up. Row j of the pairs belongs to anchor j // k, while the anchor at row j is e_{j // (2k)}. With the defaults (use_positive and use_negative both True), 12 of 16 triples reference the wrong anchor for N=4, k=2 — CosineEmbeddingLoss ends up pushing an anchor away from an embedding that was actually a positive neighbour of a different anchor. The pretraining objective is essentially scrambled noise; it only lines up when exactly one pair type is enabled.

Fix

embeddings.repeat_interleave(k_neighbors, dim=0).repeat(len(pairs), 1) — anchor-major blocks of k, tiled once per pair type — in both branches.

Tests

Two tests appended to tests/test_training_pretraining.py:

• With orthonormal embeddings, self-positives (cos 1) and orthogonal negatives (cos 0, margin 0) must give exactly zero loss — only true when every pair is attributed to its own anchor. Verified this fails on main and passes with the fix.
• A layout test pinning the anchor index sequence against the pair index sequence.

🤖 Generated with Claude Code