Hopefully final draft
This commit is contained in:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user