VonMisesFisher loss: log_cmk derivative is discontinuous at κ=100, capping learnable concentration

[sevmag]

VonMisesFisherLoss.log_cmk switches from the exact Bessel branch to log_cmk_approx at kappa_switch=100 (src/graphnet/training/loss_functions.py). The continuity offset makes log C₃ continuous in value but not in derivative.

The loss gradient w.r.t. κ is A₃(κ) = 1/tanh(κ) − 1/κ (the vMF mean resultant length). On the approx branch (log_cmk_approx = −√(4+κ²) for m=3) this becomes κ/√(4+κ²), so A₃ jumps at the switch:

• A₃(100⁻) ≈ 0.9900 (exact) → A₃(100⁺) ≈ 0.9998 (approx)

A₃ therefore never takes values in (0.9900, 0.9998). Any event whose target cosθ falls in that band — true κ from ~100 up to a few thousand — has no stationary point except the discontinuity, so gradient descent pins κ ≈ 100. It is a global attractor: κ initialized above 100 is dragged back down. Net effect: the learned concentration effectively caps at ~100.

Minimal reproduction

We pick a true concentration kappa_star, set the prediction/truth alignment to the value that makes kappa_star the maximum-likelihood optimum (cos = A₃(kappa_star)), then descend on the concentration alone and check whether it converges back to kappa_star:

import torch
from graphnet.training.loss_functions import VonMisesFisher3DLoss

loss = VonMisesFisher3DLoss()
mu = torch.tensor([[0., 0., 1.]], dtype=torch.float64)   # fixed predicted mean direction

def recover(kappa_star, init=10.0, steps=8000):
    k_star = torch.tensor(kappa_star, dtype=torch.float64)
    cos = float(1/torch.tanh(k_star) - 1/k_star)         # A3(kappa*): the optimum has mu·t = cos

    # t: truth unit vector, placed so that mu·t = cos (so kappa_star is the MLE concentration)
    t = torch.tensor([[(1-cos**2)**0.5, 0., cos]], dtype=torch.float64)

    # k: the trainable vMF concentration kappa — the single parameter we descend on
    k = torch.tensor([init], dtype=torch.float64, requires_grad=True)

    opt = torch.optim.Adam([k], lr=1.0)
    for _ in range(steps):
        opt.zero_grad()
        # prediction = [direction (=mu), kappa (=k)]; loss internally forms p = k * mu
        loss._forward(torch.cat([mu, k.view(1, 1)], 1), t).sum().backward()
        opt.step()
        with torch.no_grad():
            k.clamp_(min=1e-3)                            # keep kappa physically positive
    return k.item()

for ks in [95, 150, 300, 1000]:
    print(ks, round(recover(ks), 1))
# 95 -> 95.0,  150 -> 100.0,  300 -> 99.9,  1000 -> 100.4

mu is the (fixed) predicted mean direction, t is the truth direction positioned so the true optimum concentration is exactly kappa_star, and k is the only trainable parameter — the vMF concentration κ. κ*=95 is recovered; κ*=150/300/1000 all collapse to ~100 (same from init=500).

Seen in trained models

The same cap is visible in the κ distributions of trained vMF models. Below is the per-event vMF concentration from the published NuBench DeepIce direction predictions, split by energy: a pile-up spike forms at κ≈100, growing with energy (the high-energy bins, whose events are best reconstructed and should have the largest κ, are exactly those stacked against the cap).

Origin

The continuous exact↔approx switch was introduced in #82; the κ-validity masking was later refined in #123. The discontinuity is in the derivative — the offset only restores continuity of the value.