Pytorch port: Finetuning on AligNet triplets always lands at chance level

Thank you for open-sourcing the AligNet dataset and related code.

I ported the pipeline to PyTorch with the goal to finetune models on the AligNet triplets but so far I fail to align models above chance-level.

I tried several backbones (ViT-B/16, DINOv2, SigLIP), several dataset variants (`between_class`, `cluster_500`), and all three losses (contrastive, KLD, JSD), with a wide hyperparameter sweep (on top of the default ones suggested here). In every case, alignment on Levels never goes above chance, which is worse than their vanilla baseline.

Before raising this, I tried to rule out a bug on my side:

- The npz files and `imagenet2012_filenames.npy` are checksum-identical to the official release bucket.
- I cross-checked every triplet's `indices`-based image resolution against the npz's own `filenames` field and found zero mismatches, so I am confident that I am correctly loading the images.
- Training does work in principle: I can overfit a small fixed subset of AligNet triplets (a few hundred examples) to ~99% accuracy, so the training loop itself is not broken.
- Actually, pointing the exact same training code at the THINGS triplet dataset (a fair alternative to AligNet) instead works correctly (chance to well above chance), so the shared training/eval code does not seem to be the problem.
- I tested the paper recipe (`tau=100, normalize=False`), and others, and got the same chance-level result.
- I tried the `contrastive` loss, which ignores `target_sims` entirely and only uses the hard "image 2 is the odd one out" label, and it degrades just as much as KLD/JSD.
- I tried both the `untransformed` and `uncertainty_distillation` releases and got the same chance-level result on both.

With those sanity checks, I am fairly confident it is not a basic loading or infra bug on my end, and I feel like I am missing some convention about how the triplets or `similarities` are meant to be consumed. The part of my code most directly responsible for that is the dataset class:

```python
class AligNetDataset(Dataset):
"""Loads AligNet triplets and the corresponding ImageNet images."""

def __init__(
self,
npz_path: str | Path,
imagenet_root: str | Path,
transform=None,
use_indices: bool = False,
filename_map_path: str | Path | None = None,
):
"""
Args:
npz_path: path to an AligNet .npz file
imagenet_root: root of the ImageNet directory; images are expected at
{imagenet_root}/train/{synset}/{filename}
transform: torchvision transform applied to each PIL image independently
before stacking into a (3, C, H, W) tensor
use_indices: resolve images from indices via imagenet2012_filenames.npy,
mirroring the JAX/TFDS access path. If False, use the
filenames stored directly in the npz.
filename_map_path: path to imagenet2012_filenames.npy. Required when
use_indices=True.
"""
self.imagenet_root = Path(imagenet_root) / "train"
self.transform = transform
self.use_indices = use_indices

data = np.load(npz_path)
self.filenames = data["filenames"]
self.similarities = data["similarities"]
self.indices = data["indices"]
self.filename_map = None
if use_indices:
if filename_map_path is None:
filename_map_path = Path(npz_path).parent / "imagenet2012_filenames.npy"
self.filename_map = np.load(filename_map_path)

def __len__(self) -> int:
return len(self.filenames)

def __getitem__(self, idx: int) -> dict[str, torch.Tensor]:
if self.use_indices:
fnames = self.filename_map[self.indices[idx]]
else:
fnames = self.filenames[idx]
sims = self.similarities[idx]

images = []
for fname in fnames:
fname = str(fname)
synset = fname.split("_")[0]
path = self.imagenet_root / synset / fname
img = Image.open(path).convert("RGB")
if self.transform is not None:
img = self.transform(img)
images.append(img)

return {
"images": torch.stack(images, dim=0), # (3, C, H, W)
"similarities": torch.from_numpy(sims), # (3,)
}

```

1. One difficulty I have in debugging this is that I am not sure what numbers I should expect to reproduce. The JAX implementation uses specific BigVision weights while my PyTorch port uses HuggingFace/timm equivalents, so even with the same architecture (e.g. DINOv2), I would not necessarily expect the same odd-one-out accuracy on the same split. Is there any reference output you would suggest trying to match — even roughly — that would confirm the dataset is being used correctly?

2. Is there a transform or scale expected on `similarities` (e.g. log-space, a fixed reference point) before it is used as a softmax target, beyond what is in the released `losses.py`?

3. Is `image 2 = odd-one-out`, `(0,1) = similar pair` a safe assumption for every sampling (`between_class`, `within_class`, `class_border`, `cluster_border_500`), or does that convention differ by split?

Perhaps I am misunderstanding something about how the dataset is meant to be used. Happy to share a minimal repro script or more diagnostics if that would help. Thank you for taking the time to look at this!

Start with the released losses.py and the AligNetDataset implementation shown in the report, then compare their handling of similarities, indices, and triplet ordering with the JAX implementation. Run the PyTorch port on both AligNet and THINGS using the reported dataset variants and losses. Done means identifying the expected convention or reference output that explains the chance-level AligNet results.

Written by the indexing model from the issue text.