Hopefully final draft

This commit is contained in:
2026-09-10 02:46:31 +03:00
parent f4de5a9e8a
commit 1d1b783090
10 changed files with 414 additions and 123 deletions
+111 -59
View File
@@ -1,7 +1,8 @@
This appendix contains the full PyTorch implementation of the ROLL loss
functions used in the experiments. The code is reproduced verbatim from
\texttt{src/roll.py} in the accompanying implementation repository, with
unrelated losses (baselines, Beta-ROLL) omitted.
functions used in the experiments. The code is reproduced from
\texttt{src/roll.py} in the accompanying implementation repository,
reformatted for print width and with unrelated losses (baselines,
Beta-ROLL) omitted; the logic is unchanged.
\section{Gaussian ROLL}
@@ -17,8 +18,9 @@ import numpy as np
# -- helpers --
def _torch_mean_std(yh):
# MLE variance (n denominator), not torch.std's Bessel-corrected default.
yhmean = torch.mean(yh)
yhstd = torch.std(yh)
yhstd = torch.std(yh, correction=0)
return yhmean, yhstd
def _torch_normal_fit(yh):
@@ -123,6 +125,20 @@ class KernelizedROLLoss(Function):
x = torch.abs(x)
return (v * torch.exp(-v * x)) / ((1 + torch.exp(-v * x)) ** 2)
@staticmethod
def _log_sigma_prime(v, x):
"""log sigma'(x; v) = -v|x| - 2*softplus(-v|x|) [stable log-derivative].
Used wherever sigma' only ever feeds a normalised ratio (softmax):
computing sigma*(1-sigma) directly saturates to exactly 0 once
v|x| exceeds ~17 in fp32 (sigma itself rounds to exactly 1 there),
while the true derivative stays representable to ~88. This form
never materialises the raw derivative, so it has no such failure
point.
"""
vx = v * torch.abs(x)
return -vx - 2 * torch.nn.functional.softplus(-vx)
# -- KDE CDF and ICDF --
@staticmethod
@@ -131,69 +147,100 @@ class KernelizedROLLoss(Function):
return torch.mean(KernelizedROLLoss._sigma(v, tau - scores))
@staticmethod
def _icdf(v, scores, alpha, tol=1e-4, max_iters=60):
"""Find tau s.t. F_hat(tau; scores) = alpha via bisection.
def _icdf(v0, scores_neg, alpha, v1, scores_pos,
loss_tol=1e-4, max_iters=60):
"""Find tau s.t. F_hat_0(tau; scores_neg) = alpha via bisection.
Bracket: [min(scores) - 10/v, max(scores) + 10/v] always
satisfies F_hat(lo) < alpha < F_hat(hi). Converges in <=16
steps for normalised scores; 60 is a generous cap.
Bracket: lo/hi = min/max(scores_neg) + h0*log(alpha/(1-alpha)),
h0 = 1/v0. Gives exact containment F_hat_0(lo) < alpha <
F_hat_0(hi) for every alpha in (0,1), not just an asymptotic
range like a fixed +-10/v0 margin.
Stopping criterion certifies the LOSS, not just tau: F_hat_1 is
itself a CDF (monotonic), and the bracket invariant keeps the
true root tau* in [lo, hi] throughout, so F_hat_1(tau*) lies in
[F_hat_1(lo), F_hat_1(hi)] regardless of where tau* falls in the
bracket. Stop once that interval is narrower than loss_tol.
"""
if v > 10000.:
if v0 > 10000.:
logging.warning(
f'Bandwidth v={v:.1f} exceeds maximum 10000, capping')
v = 10000.
f'Bandwidth v0={v0:.1f} exceeds maximum 10000, capping')
v0 = 10000.
margin = 10.0 / v
lo = scores.min() - margin
hi = scores.max() + margin
h0 = 1.0 / v0
offset = (h0 * torch.log(alpha / (1 - alpha))
if isinstance(alpha, torch.Tensor)
else h0 * np.log(alpha / (1 - alpha)))
lo = scores_neg.min() + offset
hi = scores_neg.max() + offset
for _ in range(max_iters):
mid = (lo + hi) / 2
F_mid = KernelizedROLLoss._cdf(v, mid, scores)
if torch.abs(F_mid - alpha) < tol:
break
if F_mid < alpha:
mid = (lo + hi) / 2
F0_mid = KernelizedROLLoss._cdf(v0, mid, scores_neg)
if F0_mid < alpha:
lo = mid
else:
hi = mid
return mid
F1_lo = KernelizedROLLoss._cdf(v1, lo, scores_pos)
F1_hi = KernelizedROLLoss._cdf(v1, hi, scores_pos)
if F1_hi - F1_lo < loss_tol:
break
return (lo + hi) / 2
@staticmethod
def _icdf_batch(v, scores, alphas, tol=1e-4, max_iters=60):
"""Vectorized bisection: find tau_k s.t. F_hat(tau_k) = alpha_k.
def _icdf_batch(v0, scores_neg, alphas, v1, scores_pos,
loss_tol=1e-4, max_iters=60):
"""Vectorized bisection: find tau_k s.t. F_hat_0(tau_k) = alpha_k
for all K alphas at once.
Each iteration evaluates F_hat at all K midpoints via a single
(K, n) broadcast, replacing K sequential _icdf calls.
Each bracket uses the exact per-alpha offset
h0*log(alpha_k/(1-alpha_k)) (same containment guarantee as
_icdf) rather than a shared fixed margin. Stopping criterion is
the same certified-loss bound as _icdf, evaluated for all K
brackets at once.
"""
if float(v) > 10000.:
if float(v0) > 10000.:
logging.warning(
f'Bandwidth v={float(v):.1f} exceeds maximum 10000, capping')
v = (v.clamp(max=10000.)
if isinstance(v, torch.Tensor) else 10000.)
f'Bandwidth v0={float(v0):.1f} exceeds maximum 10000, capping')
v0 = (v0.clamp(max=10000.)
if isinstance(v0, torch.Tensor) else 10000.)
K = len(alphas)
K = len(alphas)
alphas_t = torch.tensor(alphas,
dtype=scores.dtype, device=scores.device)
margin = 10.0 / float(v)
lo = torch.full((K,), float(scores.min()) - margin,
device=scores.device)
hi = torch.full((K,), float(scores.max()) + margin,
device=scores.device)
mid = (lo + hi) / 2
dtype=scores_neg.dtype, device=scores_neg.device)
h0 = 1.0 / float(v0)
# log(alpha/(1-alpha)) computed in float64 first: casting alpha
# to the working dtype before subtracting from 1 loses all
# precision once alpha is within ~1e-7 of 1, collapsing
# 1-alpha to exactly 0 and the offset to inf.
alphas_f64 = np.asarray(alphas, dtype=np.float64)
offsets_np = h0 * np.log(alphas_f64 / (1 - alphas_f64))
offsets = torch.tensor(offsets_np, dtype=scores_neg.dtype,
device=scores_neg.device)
lo = float(scores_neg.min()) + offsets
hi = float(scores_neg.max()) + offsets
for _ in range(max_iters):
diff = mid.unsqueeze(1) - scores.unsqueeze(0) # (K, n)
F_mid = torch.mean(
KernelizedROLLoss._sigma(v, diff), dim=1) # (K,)
if torch.all(torch.abs(F_mid - alphas_t) < tol):
break
too_low = F_mid < alphas_t
lo = torch.where(too_low, mid, lo)
hi = torch.where(too_low, hi, mid)
mid = (lo + hi) / 2
mid = (lo + hi) / 2
diff0 = mid.unsqueeze(1) - scores_neg.unsqueeze(0) # (K, n0)
F0_mid = torch.mean(
KernelizedROLLoss._sigma(v0, diff0), dim=1) # (K,)
too_low = F0_mid < alphas_t
lo = torch.where(too_low, mid, lo)
hi = torch.where(too_low, hi, mid)
return mid # (K,)
diff1_lo = lo.unsqueeze(1) - scores_pos.unsqueeze(0) # (K, n1)
diff1_hi = hi.unsqueeze(1) - scores_pos.unsqueeze(0) # (K, n1)
F1_lo = torch.mean(
KernelizedROLLoss._sigma(v1, diff1_lo), dim=1) # (K,)
F1_hi = torch.mean(
KernelizedROLLoss._sigma(v1, diff1_hi), dim=1) # (K,)
if torch.all(F1_hi - F1_lo < loss_tol):
break
return (lo + hi) / 2 # (K,)
# -- forward pass --
@@ -225,10 +272,10 @@ class KernelizedROLLoss(Function):
# threshold where the negative-class CDF equals alpha_k.
if len(alphas) == 1:
taus = KernelizedROLLoss._icdf(
v0, scores_neg, alphas[0]).unsqueeze(0)
v0, scores_neg, alphas[0], v1, scores_pos).unsqueeze(0)
else:
taus = KernelizedROLLoss._icdf_batch(
v0, scores_neg, alphas) # (K,)
v0, scores_neg, alphas, v1, scores_pos) # (K,)
# loss_k = F_1_hat(tau_k) = 1 - TPR@FPR_k
diff_pos = taus.unsqueeze(1) - scores_pos.unsqueeze(0) # (K, n1)
@@ -264,12 +311,14 @@ class KernelizedROLLoss(Function):
# dF_1_hat(tau_k)/d(tau_k) = (1/n1) sum_j sigma'(tau_k - s_j^(1))
dF1_dtau = sp_pos.mean(dim=1) # (K,)
# d(tau_k)/d(s_i^(0))
diff_neg = taus.unsqueeze(1) - scores_neg.unsqueeze(0) # (K, n0)
s_neg = KernelizedROLLoss._sigma(v0, diff_neg) # (K, n0)
kernel = s_neg * (1 - s_neg) # v0 cancels
grad_tau_neg = kernel / kernel.sum( # (K, n0)
dim=1, keepdim=True).clamp(min=1e-10)
# d(tau_k)/d(s_i^(0)), computed in log-space for numerical
# stability -- replaces sigma*(1-sigma), which saturates to
# exactly 0 in fp32 once v0*diff is large, even though the true
# derivative stays representable much further out.
diff_neg = taus.unsqueeze(1) - scores_neg.unsqueeze(0) # (K, n0)
log_kernel = KernelizedROLLoss._log_sigma_prime(
v0, diff_neg) # (K, n0)
grad_tau_neg = torch.softmax(log_kernel, dim=1) # (K, n0)
# d(sum_k loss_k)/d(s_i^(0))
grad_neg = (dF1_dtau.unsqueeze(1)
@@ -288,9 +337,9 @@ class KernelizedROLLoss(Function):
# -- public constructors --
def _kernelized_roll_fpr_internal(yh, y, fpr, gamma=1.0, normalize=True,
def _kernelized_roll_fpr_internal(yh, y, cdf_target, gamma=1.0, normalize=True,
bce_weight=0.0, bce_pos_weight=1.0):
roll = KernelizedROLLoss.apply([fpr], 1, yh, y, gamma, normalize)
roll = KernelizedROLLoss.apply([cdf_target], 1, yh, y, gamma, normalize)
if bce_weight == 0.0:
return roll
pos_weight = torch.tensor([bce_pos_weight],
@@ -302,14 +351,17 @@ def _kernelized_roll_fpr_internal(yh, y, fpr, gamma=1.0, normalize=True,
def kernelized_roll_fpr(fpr, normalize=True,
bce_weight=0.0, bce_pos_weight=1.0):
# Pass 1-fpr as CDF target: F_0_hat(tau) = 1-fpr iff P(neg > tau) = fpr
return partial(_kernelized_roll_fpr_internal, fpr=1 - fpr,
return partial(_kernelized_roll_fpr_internal, cdf_target=1 - fpr,
normalize=normalize, bce_weight=bce_weight,
bce_pos_weight=bce_pos_weight)
def _kernelized_roll_tpr_internal(yh, y, tpr, gamma=1.0,
bce_weight=0.0, bce_pos_weight=1.0):
# Reduce TPR@FPR to FPR@TPR by negating scores and swapping labels
# Reduce TPR@FPR to FPR@TPR by negating scores and swapping labels.
# `1 - yh` (not `-yh`) is used here; safe only because it differs
# from `-yh` by a constant shift, and KernelizedROLLoss is
# translation-invariant in its scores.
roll = KernelizedROLLoss.apply([tpr], 1, 1 - yh, 1 - y, gamma)
if bce_weight == 0.0:
return roll