Fix inverted sign in KDE loss

[anishdulal]

Summary

KDELoss.forward returns -log(density), an entropy estimate that is high when
embeddings are spread out and low when they collapse. ssl_meta_arch.py:332
adds that value to the minimized total loss with a positive weight, so gradient
descent minimizes entropy — it rewards collapse.

KoLeoLoss, which this term replaced, has the opposite (correct) orientation.
Since koleo_loss_weight: 0 in every training config, KDE is currently the
only entropy regularizer in the objective, and it points the wrong way.

Evidence

1. Value at the two extremes (n=24, D=1536, both classes instantiated):

	collapsed (cos=1.0)	spread (orthonormal)	a minimizer picks
KDELoss (current)	−8.1781	−5.1441	collapsed
KoLeoLoss	+13.4534	−0.3466	spread

2. Gradient descent on each term alone, from a random init (n=24, D=128):

loss += KDE   (current)      mean pairwise cos  +0.0002 -> +1.0000   COLLAPSE (<50 steps)
loss -= KDE   (sign fixed)   mean pairwise cos  +0.0002 -> -0.0435   SPREAD
loss += KoLeo (dinov2 ref)   mean pairwise cos  +0.0002 -> -0.0433   SPREAD

Sign-fixed KDE (−0.0435) and the KoLeo it replaced (−0.0433) converge to the same
configuration, which is what you would expect from two implementations of the same idea.

3. It has sat at its pro-collapse optimum in every run. Observed kde_loss from
training_metrics.json against the analytic collapse floor -0.05 * (ln n + kappa):

run	n/GPU	observed	analytic floor
ViT-G high-res post-training	24	−0.408936	−0.408903	constant, iters 10 → 96,060
ViT-G @224 continuation	48	−0.443359	−0.443563	constant
ViT-S dry run (random init)	8	−0.354004	−0.353972	at floor by iter 10

Three different batch sizes, three matches to within one fp16 ULP. The term
contributes no anti-collapse pressure at any point in training.

Why this is the right direction

• The cited source. Karasikov et al., Training state-of-the-art pathology
  foundation models with orders of magnitude less data (2025) — the paper this
  regularizer was adopted from — introduces it in §"Self-supervised training with
  DINOv2" as "a more stable KDE regularizer [28]". Its ref [28] is Wang & Isola
  (2020), Understanding contrastive representation learning through alignment and
  uniformity on the hypersphere, whose uniformity loss is
  log E[exp(-t||x-y||^2)], minimized. For unit vectors that reduces to
  logsumexp(2t·cos) - log N + const, i.e. +log(density). This implementation
  returns its negation.
• The stated purpose. The same sentence adopts KDE over KoLeo "to ensure the
  diversity of tile embeddings generated by the FM." Diversity means spread.
• MedARC's nanopath implements the corrected form (train.py, kde_loss).

On convergence

Karasikov et al. report that "Without replacing the default KoLeo regularizer
with KDE and without the HSV filter, our training did not converge". That is
a joint ablation: both changes were absent together, so it does not isolate the
contribution of either. OpenMidnight's training converged with the term in its
present form.

Notes

• The logged kde_loss metric changes sign, so it is not comparable with
  historical run logs.
• The released 250k checkpoint was trained with this bug. vitg14_reg4.yaml
  has carried koleo_loss_weight: 0 + kde_loss_weight: .05 since its first
  commit (016786c).
• Adjacent trap, not fixed here: do_kde / do_koleo in the YAML are never
  read — the gates are cfg.dino.kde_loss_weight > 0 and
  cfg.dino.koleo_loss_weight > 0. So kde_loss_weight: -.05 silently disables
  the term rather than flipping it, and the sign cannot be corrected from config.

Reproduction

import torch, torch.nn.functional as F
from dinov2.loss import KDELoss

kde, N, D = KDELoss(), 24, 128
torch.manual_seed(1)
x = F.normalize(torch.randn(N, D), dim=-1).clone().requires_grad_(True)
opt = torch.optim.SGD([x], lr=1.0)
for _ in range(100):
    opt.zero_grad(); kde(F.normalize(x, dim=-1)).backward(); opt.step()

xn = F.normalize(x.detach(), dim=-1); s = xn @ xn.T
print("mean pairwise cosine:", ((s.sum() - s.diag().sum()) / (N * (N - 1))).item())
# on main: ~1.0 (collapsed)      with this PR: ~ -0.04 (spread)

Refs #4.

[CLAassistant]

All committers have signed the CLA.