Results overhaul
This commit is contained in:
@@ -0,0 +1,343 @@
|
||||
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.
|
||||
|
||||
\section{Gaussian ROLL}
|
||||
|
||||
The Gaussian variant fits a normal distribution to each class's scores and
|
||||
evaluates TPR at the requested FPR analytically via the Gaussian CDF/ICDF.
|
||||
|
||||
\begin{lstlisting}[language=Python]
|
||||
import torch
|
||||
from functools import partial
|
||||
import numpy as np
|
||||
|
||||
|
||||
# -- helpers --
|
||||
|
||||
def _torch_mean_std(yh):
|
||||
yhmean = torch.mean(yh)
|
||||
yhstd = torch.std(yh)
|
||||
return yhmean, yhstd
|
||||
|
||||
def _torch_normal_fit(yh):
|
||||
yhmean, yhstd = _torch_mean_std(yh)
|
||||
return torch.distributions.normal.Normal(yhmean, yhstd)
|
||||
|
||||
|
||||
# -- batch splitting --
|
||||
|
||||
def split_true_false_indeces(yh, y):
|
||||
true_indeces = torch.argwhere(y)[:, 0]
|
||||
false_indeces = torch.argwhere(torch.logical_not(y))[:, 0]
|
||||
return true_indeces, false_indeces
|
||||
|
||||
def split_true_false(yh, y):
|
||||
true_indeces, false_indeces = split_true_false_indeces(yh, y)
|
||||
return yh[true_indeces], yh[false_indeces]
|
||||
|
||||
|
||||
# -- Gaussian ROLL --
|
||||
|
||||
def _roll_loss_from_fpr_internal(yh, y, fpr):
|
||||
true_yh, false_yh = split_true_false(yh, y)
|
||||
true_normal = _torch_normal_fit(true_yh)
|
||||
false_normal = _torch_normal_fit(false_yh)
|
||||
return true_normal.cdf(false_normal.icdf(torch.tensor(1 - fpr)))
|
||||
|
||||
def roll_loss_from_fpr(fpr):
|
||||
return partial(_roll_loss_from_fpr_internal, fpr=fpr)
|
||||
\end{lstlisting}
|
||||
|
||||
\section{Kernelized (KDE) ROLL}
|
||||
|
||||
The KDE variant replaces the parametric Gaussian with a sigmoid-kernel KDE,
|
||||
enabling a custom autograd backward pass that directly implements the
|
||||
analytic gradients derived in \Cref{appendix:kde-grad}.
|
||||
|
||||
\begin{lstlisting}[language=Python]
|
||||
import torch
|
||||
import logging
|
||||
from functools import partial
|
||||
from torch.autograd import Function
|
||||
import numpy as np
|
||||
from KDEpy.bw_selection import silvermans_rule, improved_sheather_jones
|
||||
|
||||
|
||||
# -- bandwidth selection --
|
||||
|
||||
def _bandwidth(scores_1d, max_v=10000.0):
|
||||
"""Return 1/h (kernel precision v), clipped to max_v.
|
||||
|
||||
Tries ISJ; falls back to Silverman on convergence failure
|
||||
(common on small or near-constant batches early in training).
|
||||
"""
|
||||
data = scores_1d[:, np.newaxis]
|
||||
try:
|
||||
h = float(improved_sheather_jones(data))
|
||||
except Exception:
|
||||
h = float(silvermans_rule(data))
|
||||
return min(1.0 / h, max_v)
|
||||
|
||||
|
||||
# -- FPR grid for AOC objective --
|
||||
# 0.01, 0.02, ..., 0.99 (99 values); sum / 100 approximates
|
||||
# integral_0^1 (1 - TPR(alpha)) d alpha = AOC = 1 - AUC.
|
||||
|
||||
_AOC_FPRS = np.arange(0.01, 1.0, 0.01).tolist()
|
||||
_AOC_ALPHAS = [1.0 - fpr for fpr in _AOC_FPRS]
|
||||
_AOC_DIVISOR = 100
|
||||
|
||||
|
||||
class KernelizedROLLoss(Function):
|
||||
"""KDE-ROLL loss with custom backward pass.
|
||||
|
||||
Accepts a list of alpha values (CDF targets, alpha_k = 1 - FPR_k) and a
|
||||
divisor, so one class covers both objectives:
|
||||
point-FPR : alphas=[1-fpr], divisor=1
|
||||
AOC : alphas=_AOC_ALPHAS, divisor=_AOC_DIVISOR
|
||||
|
||||
Notation:
|
||||
v0, v1 -- KDE bandwidths for B_0 (negative) and B_1 (positive)
|
||||
taus -- per-alpha thresholds from inverting F_0_hat
|
||||
scores_pos -- positive-class model scores in the current batch
|
||||
scores_neg -- negative-class model scores in the current batch
|
||||
"""
|
||||
|
||||
# -- sigmoid kernel primitives --
|
||||
|
||||
@staticmethod
|
||||
def _sigma(v, x):
|
||||
"""sigma(x; v) = 1 / (1 + exp(-v*x)) [kernel CDF]."""
|
||||
return 1 / (1 + torch.exp(-v * x))
|
||||
|
||||
@staticmethod
|
||||
def _sigma_prime(v, x):
|
||||
"""sigma'(x; v) = v*exp(-v|x|) / (1+exp(-v|x|))^2 [kernel PDF].
|
||||
|
||||
Uses |x| to avoid overflow: exp(-v|x|)->0 for large |x|,
|
||||
whereas exp(-v*x)->inf for large negative x.
|
||||
sigma' is even so the two forms are equivalent.
|
||||
"""
|
||||
x = torch.abs(x)
|
||||
return (v * torch.exp(-v * x)) / ((1 + torch.exp(-v * x)) ** 2)
|
||||
|
||||
# -- KDE CDF and ICDF --
|
||||
|
||||
@staticmethod
|
||||
def _cdf(v, tau, scores):
|
||||
"""F_hat(tau; scores) = (1/n) sum_j sigma(tau - s_j; v)."""
|
||||
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.
|
||||
|
||||
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.
|
||||
"""
|
||||
if v > 10000.:
|
||||
logging.warning(
|
||||
f'Bandwidth v={v:.1f} exceeds maximum 10000, capping')
|
||||
v = 10000.
|
||||
|
||||
margin = 10.0 / v
|
||||
lo = scores.min() - margin
|
||||
hi = scores.max() + margin
|
||||
|
||||
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:
|
||||
lo = mid
|
||||
else:
|
||||
hi = mid
|
||||
|
||||
return mid
|
||||
|
||||
@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.
|
||||
|
||||
Each iteration evaluates F_hat at all K midpoints via a single
|
||||
(K, n) broadcast, replacing K sequential _icdf calls.
|
||||
"""
|
||||
if float(v) > 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.)
|
||||
|
||||
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
|
||||
|
||||
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
|
||||
|
||||
return mid # (K,)
|
||||
|
||||
# -- forward pass --
|
||||
|
||||
@staticmethod
|
||||
def forward(ctx, alphas, divisor, yh, y, gamma=1.0, normalize=True):
|
||||
# Normalise scores by std for numerical stability.
|
||||
# Chain-rule correction is applied in backward().
|
||||
# clamp avoids inf when all batch scores are identical (zero variance).
|
||||
if normalize:
|
||||
score_scale = 1 / torch.sqrt(
|
||||
torch.var(yh, correction=1).clamp(min=1e-8))
|
||||
else:
|
||||
score_scale = torch.tensor(1.0, device=yh.device)
|
||||
yh_norm = yh * score_scale
|
||||
|
||||
pos_idx, neg_idx = split_true_false_indeces(yh, y)
|
||||
scores_pos = yh_norm[pos_idx] # s^(1)
|
||||
scores_neg = yh_norm[neg_idx] # s^(0)
|
||||
|
||||
# Bandwidths via ISJ; gamma > 1 widens kernels early in training.
|
||||
v0 = (torch.tensor(
|
||||
_bandwidth(scores_neg.cpu().numpy().astype(np.float64)))
|
||||
/ gamma).to(yh.device)
|
||||
v1 = (torch.tensor(
|
||||
_bandwidth(scores_pos.cpu().numpy().astype(np.float64)))
|
||||
/ gamma).to(yh.device)
|
||||
|
||||
# tau_k = F_0_hat^{-1}(alpha_k):
|
||||
# threshold where the negative-class CDF equals alpha_k.
|
||||
if len(alphas) == 1:
|
||||
taus = KernelizedROLLoss._icdf(
|
||||
v0, scores_neg, alphas[0]).unsqueeze(0)
|
||||
else:
|
||||
taus = KernelizedROLLoss._icdf_batch(
|
||||
v0, scores_neg, alphas) # (K,)
|
||||
|
||||
# loss_k = F_1_hat(tau_k) = 1 - TPR@FPR_k
|
||||
diff_pos = taus.unsqueeze(1) - scores_pos.unsqueeze(0) # (K, n1)
|
||||
losses = torch.mean(
|
||||
KernelizedROLLoss._sigma(v1, diff_pos), dim=1) # (K,)
|
||||
|
||||
total_loss = losses.sum() / divisor
|
||||
|
||||
logging.debug(
|
||||
f'KDE-ROLL K={len(alphas)} gamma={gamma:.4f}'
|
||||
f' v0={v0:.4f} v1={v1:.4f} loss={total_loss:.4f}')
|
||||
ctx.save_for_backward(yh, scores_pos, scores_neg,
|
||||
pos_idx, neg_idx, taus, v1, v0,
|
||||
score_scale)
|
||||
ctx.divisor = divisor
|
||||
return total_loss
|
||||
|
||||
# -- backward pass --
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, grad_output):
|
||||
yh, scores_pos, scores_neg, pos_idx, neg_idx, \
|
||||
taus, v1, v0, score_scale = ctx.saved_tensors
|
||||
divisor = ctx.divisor
|
||||
n1 = scores_pos.shape[0]
|
||||
|
||||
diff_pos = taus.unsqueeze(1) - scores_pos.unsqueeze(0) # (K, n1)
|
||||
sp_pos = KernelizedROLLoss._sigma_prime(v1, diff_pos) # (K, n1)
|
||||
|
||||
# d(sum_k loss_k)/d(s_i^(1)) = -(1/n1) sum_k sigma'(tau_k - s_i)
|
||||
grad_pos = -(1.0 / n1) * sp_pos.sum(dim=0) / divisor # (n1,)
|
||||
|
||||
# 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(sum_k loss_k)/d(s_i^(0))
|
||||
grad_neg = (dF1_dtau.unsqueeze(1)
|
||||
* grad_tau_neg).sum(dim=0) / divisor # (n0,)
|
||||
|
||||
grad_input = torch.zeros(yh.shape[0], device=yh.device)
|
||||
grad_input[pos_idx] = grad_pos * grad_output
|
||||
grad_input[neg_idx] = grad_neg * grad_output
|
||||
|
||||
if torch.any(torch.isnan(grad_input)):
|
||||
logging.error('NaN gradient detected in KDE-ROLL backward')
|
||||
|
||||
# Chain-rule for score normalisation: yh_norm = yh * score_scale
|
||||
return None, None, grad_input * score_scale, None, None, None
|
||||
|
||||
|
||||
# -- public constructors --
|
||||
|
||||
def _kernelized_roll_fpr_internal(yh, y, fpr, gamma=1.0, normalize=True,
|
||||
bce_weight=0.0, bce_pos_weight=1.0):
|
||||
roll = KernelizedROLLoss.apply([fpr], 1, yh, y, gamma, normalize)
|
||||
if bce_weight == 0.0:
|
||||
return roll
|
||||
pos_weight = torch.tensor([bce_pos_weight],
|
||||
dtype=yh.dtype, device=yh.device)
|
||||
bce = torch.nn.functional.binary_cross_entropy_with_logits(
|
||||
yh, y.float(), pos_weight=pos_weight)
|
||||
return roll + bce_weight * bce
|
||||
|
||||
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,
|
||||
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
|
||||
roll = KernelizedROLLoss.apply([tpr], 1, 1 - yh, 1 - y, gamma)
|
||||
if bce_weight == 0.0:
|
||||
return roll
|
||||
pos_weight = torch.tensor([bce_pos_weight],
|
||||
dtype=yh.dtype, device=yh.device)
|
||||
bce = torch.nn.functional.binary_cross_entropy_with_logits(
|
||||
yh, y.float(), pos_weight=pos_weight)
|
||||
return roll + bce_weight * bce
|
||||
|
||||
def kernelized_roll_tpr(tpr, bce_weight=0.0, bce_pos_weight=1.0):
|
||||
return partial(_kernelized_roll_tpr_internal, tpr=tpr,
|
||||
bce_weight=bce_weight, bce_pos_weight=bce_pos_weight)
|
||||
|
||||
|
||||
def _kernelized_roll_aoc_internal(yh, y, gamma=1.0, normalize=True,
|
||||
bce_weight=0.0, bce_pos_weight=1.0):
|
||||
roll = KernelizedROLLoss.apply(
|
||||
_AOC_ALPHAS, _AOC_DIVISOR, yh, y, gamma, normalize)
|
||||
if bce_weight == 0.0:
|
||||
return roll
|
||||
pos_weight = torch.tensor([bce_pos_weight],
|
||||
dtype=yh.dtype, device=yh.device)
|
||||
bce = torch.nn.functional.binary_cross_entropy_with_logits(
|
||||
yh, y.float(), pos_weight=pos_weight)
|
||||
return roll + bce_weight * bce
|
||||
|
||||
def kernelized_roll_aoc(normalize=True,
|
||||
bce_weight=0.0, bce_pos_weight=1.0):
|
||||
return partial(_kernelized_roll_aoc_internal, normalize=normalize,
|
||||
bce_weight=bce_weight, bce_pos_weight=bce_pos_weight)
|
||||
\end{lstlisting}
|
||||
Reference in New Issue
Block a user