aner/thesis
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Thesis - wip

Browse Source
This commit is contained in:
Aner Zakobar
2026-07-15 20:58:17 +03:00
parent ce726d938f
commit a85bb2a551
13 changed files with 963 additions and 399 deletions
Download Patch File Download Diff File Expand all files Collapse all files
content/method/method.tex
+143 -394
View File
@@ -19,8 +19,8 @@ The positive and negative subsets of $\mathcal{D}$ are:
Applying a threshold $\tau$ produces predictions $\hat{y} = \mathbf{1}[f_\theta(\mathbf{x}) > \tau]$.
The \emph{true positive rate} (TPR) and \emph{false positive rate} (FPR) at $\tau$ are:
\begin{align}
\operatorname{TPR}(\tau) &= P\!\bigl(f_\theta(\mathbf{x}) > \tau \mid y = 1\bigr), \\
\operatorname{FPR}(\tau) &= P\!\bigl(f_\theta(\mathbf{x}) > \tau \mid y = 0\bigr).
\operatorname{TPR}(\alpha) &= P\!\bigl(f_\theta(\mathbf{x}) > \tau \mid y = 1\bigr), \\
\operatorname{FPR}(\alpha) &= P\!\bigl(f_\theta(\mathbf{x}) > \tau \mid y = 0\bigr).
\end{align}
In our scenario, we either wish to learn $f_{\theta}$ in order to maximize $\operatorname{TPR}$
subject to a fixed $\operatorname{FPR}$, or vice versa. Our fixed $\operatorname{TPR}$ or $\operatorname{FPR}$
@@ -112,7 +112,8 @@ This necessitates a different approach, to make these calculations differentiabl
The way to overcome this problem in differentiability is not to calculate the TPR and FPR
directly, but to estimate them by fitting probability distributions to the score output of the model at
every step. Contiguous probability distributions are contiguous w.r.t their inputs. Additionally, we can calculate the TPR, FPR given the probability distributions CDF and ICDF functions. If we can gaurantee these functions to be differentiable, which proved to be the main challenge of this work, we can gaurantee this process to be differentiable.
every step. Contiguous probability distributions are contiguous w.r.t their defining parameters, and a derivation of the parameters w.r.t the distributions themselves can be made.
Additionally, we can calculate the TPR, FPR given the probability distributions CDF and ICDF functions. If we can gaurantee these functions to be differentiable, which proved to be the main challenge of this work, we can gaurantee the entire process to be differentiable. Once we do so, a derivative of the scores themselves w.r.t. the loss can be acheived.
\subsection{General ROLL Formulation and Derivation}
@@ -157,11 +158,11 @@ Let $\hat{F}_1$ and $\hat{F}_0$ denote the CDFs of $P_{\hat{\phi}_1}$ and $P_{\h
respectively. Specific choices of $\mathcal{F}$ (Gaussian, Beta, KDE) are discussed in
\Cref{sec:roll-instantiations}.
Using this notation, at a threshold $\alpha$:
Using this notation, at a threshold $\tau$:
\begin{align}
\operatorname{TPR}(\alpha) &= 1 - \hat{F}_1(\alpha), \label{eq:tpr-cdf} \\
\operatorname{FPR}(\alpha) &= \hat{F}_0(\alpha). \label{eq:fpr-cdf}
\end{align}
\operatorname{TPR}(\tau) &= 1 - \hat{F}_1(\tau), \label{eq:tpr-cdf} \\
\operatorname{FPR}(\tau) &= 1 - \hat{F}_0(\tau). \label{eq:fpr-cdf}
\end{align}t
Let $\hat{F}_1^{-1}$ and $\hat{F}_0^{-1}$ denote the quantile functions of $P_{\hat{\phi}_1}$ and
$P_{\hat{\phi}_0}$, respectively. The threshold achieving a target TPR of $\alpha$ follows directly
@@ -172,34 +173,35 @@ from \eqref{eq:tpr-cdf}:
\]
and similarly for a target FPR of $\alpha$, from \eqref{eq:fpr-cdf}:
\[
\operatorname{FPR}(\tau) = \alpha \;\Longrightarrow\; \hat{F}_0(\tau) = \alpha
\;\Longrightarrow\; \tau = \hat{F}_0^{-1}(\alpha)
\operatorname{FPR}(\tau) = \alpha \;\Longrightarrow\; 1 - \hat{F}_0(\tau) = \alpha
\;\Longrightarrow\; \tau = \hat{F}_0^{-1}(1 - \alpha)
\]
This yields two versions of our objective. Substituting the ICDF threshold into the rate
expressions, our estimated TPR at a fixed FPR of $\alpha$ is $1 - \hat{F}_1(\hat{F}_0^{-1}(\alpha))$.
expressions, our estimated TPR at a fixed FPR of $\alpha$ is $1 - \hat{F}_1(\hat{F}_0^{-1}(1-\alpha))$.
Therefore, to maximise TPR we minimise:
\begin{equation}
\mathcal{L}_{\text{ROLL-TPR@FPR}}(f_{\theta}(\mathcal{B}) ; \alpha) = \hat{F}_1(\hat{F}_0^{-1}(\alpha))
\mathcal{L}_{\text{ROLL-TPR@FPR}}(f_{\theta}(\mathcal{B}) ; \alpha) = \hat{F}_1(\hat{F}_0^{-1}(1-\alpha))
\label{eq:roll-tpr-at-fpr}
\end{equation}
Our estimated FPR at a fixed TPR of $\alpha$ is $\hat{F}_0(\hat{F}_1^{-1}(1 - \alpha))$.
Our estimated FPR at a fixed TPR of $\alpha$ is $1 - \hat{F}_0(\hat{F}_1^{-1}(1 - \alpha))$.
Therefore, to minimise FPR:
\begin{equation}
\mathcal{L}_{\text{ROLL-FPR@TPR}}(f_{\theta}(\mathbf{X}) ; \alpha) = \hat{F}_0(\hat{F}_1^{-1}(1 - \alpha))
\mathcal{L}_{\text{ROLL-FPR@TPR}}(f_{\theta}(\mathbf{X}) ; \alpha) = 1 - \hat{F}_0(\hat{F}_1^{-1}(1 - \alpha))
\label{eq:roll-fpr-at-tpr}
\end{equation}
Observe that \eqref{eq:roll-tpr-at-fpr} and \eqref{eq:roll-fpr-at-tpr} share the same
functional structure: each evaluates one empirical CDF at the quantile of the other. This
symmetry means \eqref{eq:roll-fpr-at-tpr} can be reduced to \eqref{eq:roll-tpr-at-fpr} by
negating the model scores and exchanging the class labels. Concretely, replacing
inner structure: each evaluates one empirical CDF at the $(1-\alpha)$-quantile of the other.
The two objectives are interchangeable at the implementation level: replacing
$f_\theta(\mathbf{x})$ with $-f_\theta(\mathbf{x})$ and $y$ with $1 - y$ swaps the roles
of the two populations, turning a TPR@FPR objective into a FPR@TPR one. Because the
two objectives are therefore interchangeable at the implementation level, all derivations
of the two populations, and passing target rate $1-\alpha$ to the resulting
TPR@FPR objective recovers the FPR@TPR objective at rate $\alpha$. Because
the two objectives are therefore interchangeable at the implementation level, all derivations
and implementation details below are given for \eqref{eq:roll-tpr-at-fpr} only;
\eqref{eq:roll-fpr-at-tpr} follows by the same transformation applied to the inputs.
\eqref{eq:roll-fpr-at-tpr} follows by the same transformation applied to the inputs, with
the target rate complemented to $1-\alpha$.
\Cref{fig:roll-principle} provides a geometric view of \eqref{eq:roll-tpr-at-fpr}.
\begin{figure}[H]
@@ -221,20 +223,20 @@ and implementation details below are given for \eqref{eq:roll-tpr-at-fpr} only;
We now derive the gradient of \eqref{eq:roll-tpr-at-fpr} with respect to $f_\theta(\mathbf{x}_i)$.
If $y_i = 1$, then $(\mathbf{x}_i, y_i) \in \mathcal{B}_1$ and $f_\theta(\mathbf{x}_i)$ has no
effect on $\hat{F}_0^{-1}(\alpha)$. Therefore:
effect on $\hat{F}_0^{-1}(1-\alpha)$. Therefore:
\begin{equation}
\left.\frac{\partial \mathcal{L}_{\text{ROLL-TPR@FPR}}(f_\theta(\mathcal{B}) ; \alpha)}
{\partial f_\theta(\mathbf{x}_i)}\right|_{y_i = 1}
= \frac{\partial \hat{F}_1(\hat{F}_0^{-1}(\alpha))}{\partial f_\theta(\mathbf{x}_i)}
= \frac{\partial \hat{F}_1(\hat{F}_0^{-1}(1-\alpha))}{\partial f_\theta(\mathbf{x}_i)}
\end{equation}
If $y_i = 0$, then $(\mathbf{x}_i, y_i) \in \mathcal{B}_0$ and $f_\theta(\mathbf{x}_i)$ affects
$\hat{F}_0^{-1}(\alpha)$ instead. Applying the chain rule:
$\hat{F}_0^{-1}(1-\alpha)$ instead. Applying the chain rule:
\begin{equation}
\left.\frac{\partial \mathcal{L}_{\text{ROLL-TPR@FPR}}(f_\theta(\mathbf{X}) ; \alpha)}
{\partial f_\theta(\mathbf{x}_i)}\right|_{y_i = 0}
= \frac{\partial \hat{F}_1(\hat{F}_0^{-1}(\alpha))}{\partial \hat{F}_0^{-1}(\alpha)}
\cdot \frac{\partial \hat{F}_0^{-1}(\alpha)}{\partial f_\theta(\mathbf{x}_i)}
= \frac{\partial \hat{F}_1(\hat{F}_0^{-1}(1-\alpha))}{\partial \hat{F}_0^{-1}(1-\alpha)}
\cdot \frac{\partial \hat{F}_0^{-1}(1-\alpha)}{\partial f_\theta(\mathbf{x}_i)}
\end{equation}
Unifying both cases, the general gradient applicable to all ROLL instantiations is:
@@ -242,10 +244,10 @@ Unifying both cases, the general gradient applicable to all ROLL instantiations
\frac{\partial \mathcal{L}_{\text{ROLL-TPR@FPR}}(f_\theta(\mathbf{X}) ; \alpha)}
{\partial f_\theta(\mathbf{x}_i)} =
\begin{cases}
\dfrac{\partial \hat{F}_1(\hat{F}_0^{-1}(\alpha))}{\partial f_\theta(\mathbf{x}_i)}
\dfrac{\partial \hat{F}_1(\hat{F}_0^{-1}(1-\alpha))}{\partial f_\theta(\mathbf{x}_i)}
& \text{if } y_i = 1 \\[10pt]
\dfrac{\partial \hat{F}_1(\hat{F}_0^{-1}(\alpha))}{\partial \hat{F}_0^{-1}(\alpha)}
\cdot \dfrac{\partial \hat{F}_0^{-1}(\alpha)}{\partial f_\theta(\mathbf{x}_i)}
\dfrac{\partial \hat{F}_1(\hat{F}_0^{-1}(1-\alpha))}{\partial \hat{F}_0^{-1}(1-\alpha)}
\cdot \dfrac{\partial \hat{F}_0^{-1}(1-\alpha)}{\partial f_\theta(\mathbf{x}_i)}
& \text{if } y_i = 0
\end{cases}
\label{eq:roll-gradient}
@@ -253,6 +255,49 @@ Unifying both cases, the general gradient applicable to all ROLL instantiations
The concrete form of each partial derivative depends on the choice of distribution family
$\mathcal{F}$, and is derived for each instantiation in \Cref{sec:roll-instantiations}.
\subsection{ROLL variant - direct AUC optimisation}
In addition to optimising for a specific rate, ROLL provides us an opportunity to
optimize the AUC of the ROC directly.
The AUC is defined as the area under the curve of the ROC. The ROC is a function that
for every FPR gives us the performing TPR of the model we have fitted.
Formally, for the function $TPR(\alpha)$ as defined above, we may define AUC as:
\begin{equation}
\operatorname{AUC} = \int_{0}^{1}\operatorname{TPR}(\alpha) d \alpha
\end{equation}
The AUC is something we would like to maximize. In order to do so, we would like to
minimize the area over the curve, which is $1-AUC$.
We note that $1-\operatorname{TPR}(\alpha) = \mathcal{L}(TPR@FPR)$.
From here, we would like to minimize:
\begin{equation}
\mathcal{L}_{AOC} = \int_0^1\mathcal{L}_{\text{ROLL-TPR@FPR}}(\alpha) d \alpha
\end{equation}
A direct integration is not part of this work. However, integration can be approximated. Instead
of integrating directly, it is possible to evaluate the loss for every FPR between 0 and 1, sum them, and re-normalize.
Therefor, subdividing the interval $(0, 1)$ to 99 steps (every alpha starting from $0.01$ and ending with $0.99$, since a value of $0$ or $1$ would be illegal), our loss for minimizing AOC becomes:
\begin{equation}
\mathcal{L}_{ROLL-AOC} = \frac{1}{99}\sum_{i=1}^{99}\mathcal{L}_{ROLL-TPR@FPR}\left(\frac{i}{100}\right)
\label{eq:roll-aoc}
\end{equation}
The gradient per score then simply becomes the average gradient over every threshold.
\begin{equation}
\frac{\partial \mathcal{L}_{ROLL-AOC}(f_\theta(\mathbf{X}))}
{\partial f_\theta(\mathbf{x}_i)} = \frac{1}{99}\sum_{i=1}^{99}\left(\frac{\partial \mathcal{L}_{\text{ROLL-TPR@FPR}}\left(f_\theta(\mathbf{X}) ; \frac{i}{100}\right)}
{\partial f_\theta(\mathbf{x}_i)} \right)
\label{eq:roll-aoc-gradient}
\end{equation}
%------------------------------------------------
@@ -322,7 +367,7 @@ available in standard numerical libraries (e.g.\ \texttt{torch.erf} and
\texttt{torch.erfinv}). Plugging into the general ROLL framework, the threshold
$\tau$ achieving FPR $= \alpha$ and the resulting loss are:
\begin{align}
\tau &= \mu_0 + \sigma_0\sqrt{2}\,\text{ierf}(2\alpha - 1) \nonumber \\
\tau &= \mu_0 + \sigma_0\sqrt{2}\,\text{ierf}(1 - 2\alpha) \nonumber \\
\mathcal{L}_{\text{ROLL-TPR@FPR}}^{\text{GAUSSIAN}}(f_{\theta}(\mathcal{B}) ; \alpha)
&= \frac{1}{2}\left[1 + \text{erf}\!\left(\frac{\tau - \mu_1}{\sigma_1\sqrt{2}}\right)\right]
\label{eq:roll-tpr-at-fpr-gaussian}
@@ -331,87 +376,8 @@ $\tau$ achieving FPR $= \alpha$ and the resulting loss are:
\subsubsection{Gradient Derivation}
\label{sec:roll-gaussian-backward}
% Differentiate the Gaussian loss with respect to $f_\theta(\mathbf{x}_i)$.
% Gradients flow through $\mu_k$ and $\sigma_k^2$ (which are differentiable
% functions of the scores), and through $\Phi$ and $\Phi^{-1}$, which have
% simple closed-form derivatives.
In order to compute the gradient derivation for the gaussian estimates we must first compute the gradient to each of the parameters.
Firstly, for $y_i = 1$.
\[
\dfrac{\partial \hat{F}_1(\hat{F}_0^{-1}(\alpha))}{\partial \mu_1} =
\frac{\partial \frac{1}{2}\left[1 + \text{erf}\!\left(\frac{\hat{F}_0^{-1}(\alpha) - \mu_1}{\sigma_1\sqrt{2}}\right)\right]}{\partial \mu_1}
\]
\[
= -\frac{1}{\sigma_1\sqrt{2\pi}}\exp\!\left(-\frac{(\hat{F}_0^{-1}(\alpha) - \mu_1)^2}{2\sigma_1^2}\right)
\]
\[
\frac{\partial \mu_1}{\partial f_{\theta}(\mathbf{x}_i)} = \frac{1}{|\mathcal{B}_1|}
\]
\[
\dfrac{\partial \hat{F}_1(\hat{F}_0^{-1}(\alpha))}{\partial \sigma_1} =
\frac{\partial \frac{1}{2}\left[1 + \text{erf}\!\left(\frac{\hat{F}_0^{-1}(\alpha) - \mu_1}{\sigma_1\sqrt{2}}\right)\right]}{\partial \sigma_1}
\]
\[
= \frac{\mu_1 - \hat{F}_0^{-1}(\alpha)}{\sigma_1^2\sqrt{2\pi}} \exp\!\left( -\frac{(\hat{F}_0^{-1}(\alpha) - \mu_1)^2}{2\sigma_1^2} \right)
\]
\[
\frac{\partial \sigma_1}{\partial f_{\theta}(\mathbf{x}_i)} = \frac{f_\theta(\mathbf{x}_i) - \mu_1}{|\mathcal{B}_1| \sigma_1}
\]
Combining via the chain rule, the gradient for $y_i = 1$ is:
\begin{equation}
\dfrac{\partial \hat{F}_1(\hat{F}_0^{-1}(\alpha))}{\partial f_{\theta}(\mathbf{x}_i)}
= -\frac{1}{|\mathcal{B}_1|\,\sigma_1\sqrt{2\pi}}
\exp\!\left(-\frac{(\tau - \mu_1)^2}{2\sigma_1^2}\right)
\left(1 + \frac{(\tau - \mu_1)(f_\theta(\mathbf{x}_i) - \mu_1)}{\sigma_1^2}\right)
\label{eq:gauss-grad-y1}
\end{equation}
where $\tau = \hat{F}_0^{-1}(\alpha) = \mu_0 + \sigma_0\sqrt{2}\,\operatorname{ierf}(2\alpha - 1)$.
Given $y_i = 0$, the score $f_\theta(\mathbf{x}_i)$ affects the loss only through the threshold
$\tau = \hat{F}_0^{-1}(\alpha)$, which depends on $\mu_0$ and $\sigma_0$.
Applying the chain rule:
\[
\frac{\partial \hat{F}_1(\tau)}{\partial f_\theta(\mathbf{x}_i)}
= \underbrace{\frac{\partial \hat{F}_1(\tau)}{\partial \tau}}_{\text{PDF of class 1 at }\tau}
\cdot \frac{\partial \tau}{\partial f_\theta(\mathbf{x}_i)}
\]
The first factor is the Gaussian PDF evaluated at $\tau$:
\[
\frac{\partial \hat{F}_1(\tau)}{\partial \tau}
= \frac{1}{\sigma_1\sqrt{2\pi}}\exp\!\left(-\frac{(\tau - \mu_1)^2}{2\sigma_1^2}\right)
\]
For the second factor, $\tau = \mu_0 + \sigma_0\sqrt{2}\,\operatorname{ierf}(2\alpha - 1)$, so
$\frac{\partial \tau}{\partial \mu_0} = 1$ and
$\frac{\partial \tau}{\partial \sigma_0} = \sqrt{2}\,\operatorname{ierf}(2\alpha - 1) = \frac{\tau - \mu_0}{\sigma_0}$.
Combined with $\frac{\partial \mu_0}{\partial f_\theta(\mathbf{x}_i)} = \frac{1}{|\mathcal{B}_0|}$ and
$\frac{\partial \sigma_0}{\partial f_\theta(\mathbf{x}_i)} = \frac{f_\theta(\mathbf{x}_i) - \mu_0}{|\mathcal{B}_0|\sigma_0}$:
\[
\frac{\partial \tau}{\partial f_\theta(\mathbf{x}_i)}
= \frac{1}{|\mathcal{B}_0|}\left(1 + \frac{(\tau - \mu_0)(f_\theta(\mathbf{x}_i) - \mu_0)}{\sigma_0^2}\right)
\]
Therefore, the gradient for $y_i = 0$ is:
\begin{equation}
\dfrac{\partial \hat{F}_1(\hat{F}_0^{-1}(\alpha))}{\partial f_\theta(\mathbf{x}_i)}
= \frac{1}{|\mathcal{B}_0|\,\sigma_1\sqrt{2\pi}}
\exp\!\left(-\frac{(\tau - \mu_1)^2}{2\sigma_1^2}\right)
\left(1 + \frac{(\tau - \mu_0)(f_\theta(\mathbf{x}_i) - \mu_0)}{\sigma_0^2}\right)
\label{eq:gauss-grad-y0}
\end{equation}
\paragraph{Combined gradient.}
Substituting \Cref{eq:gauss-grad-y1} and \Cref{eq:gauss-grad-y0} into \Cref{eq:roll-gradient},
and letting $\tau = \hat{F}_0^{-1}(\alpha) = \mu_0 + \sigma_0\sqrt{2}\,\operatorname{ierf}(2\alpha-1)$:
The full derivation is given in Appendix~\ref{appendix:gauss-grad}.
Letting $\tau = \hat{F}_0^{-1}(1-\alpha) = \mu_0 + \sigma_0\sqrt{2}\,\operatorname{ierf}(1-2\alpha)$, the gradient is:
\begin{equation}
\frac{\partial \mathcal{L}_{\text{ROLL-TPR@FPR}}^{\text{GAUSSIAN}}}{\partial f_\theta(\mathbf{x}_i)} =
@@ -476,10 +442,10 @@ Substituting into the general ROLL framework \eqref{eq:roll-tpr-at-fpr}:
\mathcal{L}_{\text{ROLL-TPR@FPR}}^{\text{BETA}}(f_\theta(\mathcal{B});\,\alpha)
= I_{\hat\tau}(\hat{a}_1,\hat{b}_1),
\qquad
\hat\tau = I^{-1}_\alpha(\hat{a}_0,\hat{b}_0)
\hat\tau = I^{-1}_{1-\alpha}(\hat{a}_0,\hat{b}_0)
\label{eq:roll-tpr-at-fpr-beta}
\end{equation}
where $I^{-1}_\alpha(a,b)$ denotes the $\alpha$-quantile of $\operatorname{Beta}(a,b)$.
where $I^{-1}_{1-\alpha}(a,b)$ denotes the $(1-\alpha)$-quantile of $\operatorname{Beta}(a,b)$.
\subsubsection{Gradient Derivation}
\label{sec:roll-beta-backward}
@@ -492,12 +458,6 @@ Gradients of \eqref{eq:roll-tpr-at-fpr-beta} with respect to $f_\theta(\mathbf{x
are therefore propagated via automatic differentiation through the estimator
\eqref{eq:beta-params}, the CDF evaluation, and the quantile inversion.
Note that the gradient balance property (\Cref{prop:gradient-balance}) does not hold
for Beta ROLL: the sigmoid activation that maps raw scores into $(0,1)$ is nonlinear,
so a uniform additive shift of the raw scores produces a non-uniform shift of the
transformed inputs, violating the translation-invariance assumption required by the
proof.
\subsection{KDE ROLL}
\label{sec:roll-kde}
@@ -522,7 +482,7 @@ As before, to calculate the loss, we calculate:
\[
\mathcal{L}_{\text{ROLL-TPR@FPR}}^{\text{KDE}}(f_{\theta}(\mathcal{B}) ; \alpha) = \hat{F}_1(\hat{F}_0^{-1}(\alpha))
\mathcal{L}_{\text{ROLL-TPR@FPR}}^{\text{KDE}}(f_{\theta}(\mathcal{B}) ; \alpha) = \hat{F}_1(\hat{F}_0^{-1}(1-\alpha))
\]
To calculate this, we first define the kernel function as $\sigma '$, and the CDF of the kernel function as $\sigma$. The approach used here is to define $\sigma$ as a sigmoid function,.
@@ -542,24 +502,10 @@ The CDF of the KDE function then becomes:
\hat{F}_{KDE}(\tau ; \mathbf{X}) = \frac{1}{|X|}\sum_i \sigma(\tau - x_i)
\]
While the CDF has a nice formula, the inverse CDF has no known closed formula. It instead must be calculated numerically. Thus, the calculation of the decision threshold $\tau$ must be achevied
using a numerical calculation algorithm. In this work, the Newton--Raphson method~\cite{press2007numerical} is chosen.
Firstly we must calculate $\hat{F}_0^{-1}(\alpha)$ = \tau. We start with our initial guess, $\tau_0$, and for each step, calculate:
\[
\tau_{n+1} = \tau_n - \frac{\hat{F}_0(\tau_n ; \mathcal{X}_0) - \alpha}{\frac{\partial \hat{F}_0(\tau_n ; \mathcal{X}_0)}{\partial \tau}}
\]
This step is taken until $\hat{F}_0(\tau_n ; \mathcal{B}_0) \approx \alpha$ within some acceptable error (in our case $1e-4$).
Our derivative w.r.t. $\tau$ is
\[ %
\frac{\partial \hat{F}_0(\tau_n ; \mathcal{X}_0)}{\partial \tau} = \frac{1}{|\mathcal{X}_0|}\sum_i \sigma'(\tau - x_i)
\]
At this point, we have succesfully calculated $\tau$ for which $\tau = \hat{F}_0^{-1}(\alpha ; \mathbf{X}_0)$.
While the CDF has a closed form, its inverse has none and must be found numerically.
We use bisection, which requires no derivative evaluation and is guaranteed to converge
for any strictly monotone CDF; the bracket construction and iteration are detailed in
Appendix~\ref{appendix:kde-bisection}.
We can now relatively easily calculate $\hat{F}_1(\tau ; \mathbf{X}_1)$ using the formulat for $\hat{F}_{KDE}(\tau ; \mathbf{X})$ above.
@@ -574,144 +520,7 @@ Importantly, we must save the calculation of $\tau$ for the backward derivation,
% Gradients flow through the kernel evaluations back to the model scores.
% This is the section that connects to \Cref{sec:roll-backward} motivation.
We derive the concrete form of the two terms from \Cref{eq:roll-gradient} for the KDE instantiation,
where $\tau = \hat{F}_0^{-1}(\alpha ; \mathcal{B}_0)$.
\paragraph{Case $y_i = 1$ (positive class).}
When $\mathbf{x}_i \in \mathcal{B}_1$, the score $f_\theta(\mathbf{x}_i)$ appears directly in $\hat{F}_1$
and has no effect on $\tau$ (which depends only on $\mathcal{B}_0$). Therefore:
\[
\dfrac{\partial \hat{F}_1(\hat{F}_0^{-1}(\alpha))}{\partial f_{\theta}(\mathbf{x}_i)} =
\frac{\partial \frac{1}{|\mathcal{B}_1|}\sum_{\mathbf{x}_j \in \mathcal{B}_1} \sigma_1(\tau - f_{\theta}(\mathbf{x}_j))}{\partial f_{\theta}(\mathbf{x}_i)}
\]
\[
= \frac{\partial \left(\sum_{\mathbf{x}_j \in \mathcal{B}_1 ; j \neq i} \frac{1}{|\mathcal{B}_1|}\sigma_1(\tau - f_{\theta}(\mathbf{x}_j)) + \frac{1}{|\mathcal{B}_1|}\sigma_1(\tau - f_{\theta}(\mathbf{x}_i))\right)}{\partial f_{\theta}(\mathbf{x}_i)}
\]
Applying the chain rule to the only $f_\theta(\mathbf{x}_i)$-dependent term,
$\frac{\partial}{\partial f_\theta(\mathbf{x}_i)}(\tau - f_\theta(\mathbf{x}_i)) = -1$:
\begin{equation}
\frac{\partial \hat{F}_1(\hat{F}_0^{-1}(\alpha))}{\partial f_{\theta}(\mathbf{x}_i)} = -\frac{1}{|\mathcal{B}_1|}\sigma_1'(\tau - f_{\theta}(\mathbf{x}_i))
\label{eq:kde-grad-y1}
\end{equation}
\paragraph{Case $y_i = 0$ (negative class).}
When $\mathbf{x}_i \in \mathcal{B}_0$, the score $f_\theta(\mathbf{x}_i)$ does not appear directly in
$\hat{F}_1$, but influences it through the threshold $\tau$. Per \Cref{eq:roll-gradient}, we compute
each factor of the chain rule separately.
\medskip\noindent\textit{Derivative of $\hat{F}_1$ with respect to $\tau$.}
\[
\frac{\partial \hat{F}_1(\hat{F}_0^{-1}(\alpha))}{\partial \hat{F}_0^{-1}(\alpha)} = \frac{\partial \frac{1}{|\mathcal{B}_1|}\sum_{\mathbf{x}_j \in \mathcal{B}_1} \sigma_1(\tau - f_{\theta}(\mathbf{x}_j))}{\partial \tau}
\]
Applying the chain rule, $\frac{\partial}{\partial\tau}(\tau - f_\theta(\mathbf{x}_j)) = +1$:
\begin{equation}
= +\frac{1}{|\mathcal{B}_1|}\sum_{\mathbf{x}_j \in \mathcal{B}_1}\sigma_1'(\tau - f_{\theta}(\mathbf{x}_j))
\label{eq:kde-dF1-dtau}
\end{equation}
\medskip\noindent\textit{Derivative of $\tau$ with respect to $f_\theta(\mathbf{x}_i)$.}
This requires differentiating a quantity computed numerically --- $\tau = \hat{F}_0^{-1}(\alpha ; \mathcal{B}_0)$
is found via Newton's method rather than in closed form.
We apply the inverse derivative rule:
\[
f'(x) = \frac{1}{f^{-1}'(f(x))}
\]
This allows us to use the derivative of the inverse function instead. We need the inverse of
$\hat{F}_0^{-1}$ viewed as a function of $f_\theta(\mathbf{x}_i)$ --- that is, an expression for
$f_\theta(\mathbf{x}_i)$ in terms of $\tau$, $\alpha$, and $\mathcal{B}_0 \setminus \{\mathbf{x}_i\}$.
Solving $|\mathcal{B}_0|\cdot\alpha = \sum_{\mathbf{x}_j \in \mathcal{B}_0}\sigma_0(\tau - f_\theta(\mathbf{x}_j))$
for $f_\theta(\mathbf{x}_i)$ gives:
\[
f_\theta(\mathbf{x}_i) = \tau - \sigma_0^{-1}\left( |\mathcal{B}_0|\cdot\alpha - \sum_{\mathbf{x}_j \in \mathcal{B}_0 ; j \neq i}\sigma_0(\tau - f_{\theta}(\mathbf{x}_j))\right)
\]
Applying the inverse derivative rule to this inverse:
\[
\frac{\partial \hat{F}_0^{-1}(\alpha ; \mathcal{B}_0)}{\partial f_\theta(\mathbf{x}_i)} = \frac{1}{\frac{\partial \left(\tau - \sigma_0^{-1}\left( |\mathcal{B}_0|\cdot \alpha - \sum_{\mathbf{x}_j \in \mathcal{B}_0 ; j \neq i}\sigma_0(\tau - f_{\theta}(\mathbf{x}_j))\right)\right)}{\partial \tau}}
\]
Applying the chain rule:
\[
= \frac{1}{\sigma_0^{-1}{}'\left( |\mathcal{B}_0|\cdot \alpha - \sum_{\mathbf{x}_j \in \mathcal{B}_0 ; j \neq i}\sigma_0(\tau - f_{\theta}(\mathbf{x}_j))\right)\cdot\left( \sum_{\mathbf{x}_j \in \mathcal{B}_0 ; j \neq i}\sigma_0'(\tau - f_{\theta}(\mathbf{x}_j))\right) + 1}
\]
Since $|\mathcal{B}_0|\cdot\alpha = \sum_{\mathbf{x}_j \in \mathcal{B}_0}\sigma_0(\tau - f_\theta(\mathbf{x}_j))$,
the argument of $\sigma_0^{-1}{}'$ simplifies:
\[
|\mathcal{B}_0|\cdot \alpha - \sum_{\mathbf{x}_j \in \mathcal{B}_0 ; j \neq i}\sigma_0(\tau - f_{\theta}(\mathbf{x}_j)) = \sigma_0(\tau - f_{\theta}(\mathbf{x}_i))
\]
Therefore:
\[
\frac{\partial \hat{F}_0^{-1}(\alpha ; \mathcal{B}_0)}{\partial f_\theta(\mathbf{x}_i)} = \frac{1}{\sigma_0^{-1}{}'\!\left( \sigma_0(\tau - f_{\theta}(\mathbf{x}_i)) \right)\cdot\left( \sum_{\mathbf{x}_j \in \mathcal{B}_0 ; j \neq i}\sigma_0'(\tau - f_{\theta}(\mathbf{x}_j))\right) + 1 }
\]
Substituting $\sigma_0^{-1}{}'(\sigma_0(u)) = \frac{1}{\sigma_0'(u)}$ and multiplying numerator and denominator
by $\sigma_0'(\tau - f_{\theta}(\mathbf{x}_i))$:
\begin{equation}
\frac{\partial \hat{F}_0^{-1}(\alpha ; \mathcal{B}_0)}{\partial f_\theta(\mathbf{x}_i)} =
\frac{\sigma_0'(\tau - f_{\theta}(\mathbf{x}_i))}{\displaystyle\sum_{\mathbf{x}_j \in \mathcal{B}_0}\sigma_0'(\tau - f_{\theta}(\mathbf{x}_j))}
\label{eq:kde-grad-tau}
\end{equation}
\medskip\noindent\textit{Efficient computation via the sigmoid identity.}
Direct evaluation of \eqref{eq:kde-grad-tau} requires computing $\sigma_0'$ for every
point, then dividing by the per-point value — potentially unstable when a point is far
from $\tau$. For the sigmoid kernel, we can avoid this by exploiting the identity
\begin{equation}
\sigma(u;\,v)\,\bigl(1 - \sigma(u;\,v)\bigr) = \frac{\sigma'(u;\,v)}{v}
\label{eq:sigmoid-identity}
\end{equation}
which holds for all $u \in \mathbb{R}$ (proved by direct substitution: both sides equal
$\exp(-v|u|)/(1+\exp(-v|u|))^2$). Applying this with $u = \tau - f_\theta(\mathbf{x}_j)$:
\[
\sigma_0'(\tau - f_\theta(\mathbf{x}_j)) = v_0\,\sigma_0(\tau - f_\theta(\mathbf{x}_j))
\,\bigl(1-\sigma_0(\tau - f_\theta(\mathbf{x}_j))\bigr)
\]
Since $v_0$ cancels between numerator and denominator of \eqref{eq:kde-grad-tau},
the ratio can be written entirely in terms of the sigmoid values already computed
during the forward pass:
\begin{equation}
\frac{\partial \hat{F}_0^{-1}(\alpha ; \mathcal{B}_0)}{\partial f_\theta(\mathbf{x}_i)}
=
\frac{\sigma_0(\tau-f_\theta(\mathbf{x}_i))\,\bigl(1-\sigma_0(\tau-f_\theta(\mathbf{x}_i))\bigr)}
{\displaystyle\sum_{\mathbf{x}_j \in \mathcal{B}_0}
\sigma_0(\tau-f_\theta(\mathbf{x}_j))\,\bigl(1-\sigma_0(\tau-f_\theta(\mathbf{x}_j))\bigr)}
\label{eq:kde-grad-tau-efficient}
\end{equation}
This avoids recomputing $\sigma_0'$ (which requires an extra exponential), instead
reusing the sigmoid activations cached from the CDF computation.
\paragraph{Combined gradient.}
Substituting \Cref{eq:kde-grad-y1}, \Cref{eq:kde-dF1-dtau}, and \Cref{eq:kde-grad-tau}
into \Cref{eq:roll-gradient}:
The derivation is given in Appendix~\ref{appendix:kde-grad}; letting $\tau = \hat{F}_0^{-1}(1-\alpha ; \mathcal{B}_0)$, the gradient is:
\begin{equation}
\frac{\partial \mathcal{L}_{\text{ROLL-TPR@FPR}}}{\partial f_\theta(\mathbf{x}_i)} =
@@ -726,87 +535,10 @@ into \Cref{eq:roll-gradient}:
\label{eq:kde-grad-combined}
\end{equation}
where $\tau = \hat{F}_0^{-1}(\alpha ; \mathcal{B}_0)$.
where $\tau = \hat{F}_0^{-1}(1-\alpha ; \mathcal{B}_0)$.
\subsubsection{Sigmoid Kernel: Explicit Equations}
\label{sec:roll-kde-sigmoid}
The derivations above hold for any differentiable CDF kernel $\sigma$.
In our implementation we choose the logistic sigmoid:
%
\begin{equation}
\sigma(x ; v) = \frac{1}{1 + \exp(-vx)},
\qquad
\sigma'(x ; v) = \frac{v\,\exp(-v|x|)}{\bigl(1 + \exp(-v|x|)\bigr)^{2}}
\label{eq:sigmoid-kernel}
\end{equation}
%
where $v > 0$ is the bandwidth parameter. The derivative formula evaluates at
$|x|$ rather than $x$: since $\sigma'$ is an even function
($\sigma'(x) = \sigma'(-x)$ for any symmetric logistic), both forms are
mathematically identical, but $\exp(-v|x|) \to 0$ as $|x| \to \infty$ while
$\exp(-vx) \to \infty$ for $x \to -\infty$, so the $|x|$ form avoids
floating-point overflow when scores fall well below the threshold.
\paragraph{Forward pass.}
Substituting \eqref{eq:sigmoid-kernel} into the KDE CDF:
%
\begin{equation}
\hat{F}_{k}(\tau ; \mathbf{s}^{(k)}) =
\frac{1}{|\mathcal{B}_k|}
\sum_{\mathbf{x}_j \in \mathcal{B}_k}
\frac{1}{1 + \exp\!\bigl(-v_k\bigl(\tau - f_\theta(\mathbf{x}_j)\bigr)\bigr)}
\label{eq:kde-sigmoid-cdf}
\end{equation}
%
and its derivative with respect to $\tau$ — used both in the Newton--Raphson
inversion and in the gradient computation below — is:
%
\begin{equation}
\frac{\partial \hat{F}_{k}(\tau)}{\partial \tau}
=
\frac{1}{|\mathcal{B}_k|}
\sum_{\mathbf{x}_j \in \mathcal{B}_k}
\frac{v_k\,\exp\!\bigl(-v_k\bigl|\tau - f_\theta(\mathbf{x}_j)\bigr|\bigr)}
{\Bigl(1 + \exp\!\bigl(-v_k\bigl|\tau - f_\theta(\mathbf{x}_j)\bigr|\bigr)\Bigr)^{2}}
\label{eq:kde-sigmoid-pdf}
\end{equation}
\paragraph{Gradient.}
Define the sigmoid kernel PDF shorthand:
\begin{equation}
\varphi_k(u) \;=\;
\frac{v_k\,\exp(-v_k|u|)}{\bigl(1+\exp(-v_k|u|)\bigr)^{2}}
\label{eq:sigmoid-phi}
\end{equation}
This is simply $\sigma'(u ; v_k)$ from \eqref{eq:sigmoid-kernel}, a bell-shaped function
centred at zero with width controlled by bandwidth $v_k$.
Substituting \eqref{eq:sigmoid-kernel} into \eqref{eq:kde-grad-combined} and using this
notation, the explicit sigmoid-kernel gradient is:
\begin{equation}
\frac{\partial \mathcal{L}_{\text{ROLL-TPR@FPR}}}{\partial f_\theta(\mathbf{x}_i)} =
\begin{cases}
-\dfrac{\varphi_1\!\left(\tau - f_\theta(\mathbf{x}_i)\right)}{|\mathcal{B}_1|}
& \text{if } y_i = 1\\[14pt]
+\,\dfrac{\displaystyle\sum_{\mathbf{x}_j \in \mathcal{B}_1}
\varphi_1\!\left(\tau - f_\theta(\mathbf{x}_j)\right)}{|\mathcal{B}_1|}
\;\cdot\;
\dfrac{\varphi_0\!\left(\tau-f_\theta(\mathbf{x}_i)\right)}
{\displaystyle\sum_{\mathbf{x}_j \in \mathcal{B}_0}
\varphi_0\!\left(\tau-f_\theta(\mathbf{x}_j)\right)}
& \text{if } y_i = 0
\end{cases}
\label{eq:kde-sigmoid-grad}
\end{equation}
where $\tau = \hat{F}_0^{-1}(\alpha ; \mathcal{B}_0)$.
The signs reflect the direction of each class's effect on the loss: raising a
positive-class score pushes $\hat{F}_1(\tau)$ down (gradient negative), while
raising a negative-class score shifts $\tau$ right, also lowering $\hat{F}_1(\tau)$
(gradient positive). Positive-class samples are weighted by their individual kernel
The explicit sigmoid-kernel form of \eqref{eq:kde-grad-combined} — forward pass CDF, efficient gradient, and implementation notes — is given in Appendix~\ref{appendix:kde-sigmoid}.
Positive-class samples are weighted by their individual kernel
PDF at $\tau$; negative-class samples are weighted by the mean positive-class kernel
PDF scaled by each sample's share of the negative-class kernel mass.
@@ -859,7 +591,7 @@ From \Cref{eq:roll-gradient}, for $y_i = 0$:
We claim $\sum_{\mathbf{x}_i \in \mathcal{B}_0}\frac{\partial \tau}{\partial f_\theta(\mathbf{x}_i)} = 1$.
Shift all negative scores uniformly: $f_\theta(\mathbf{x}_i) \to f_\theta(\mathbf{x}_i) + \delta$ for
$\mathbf{x}_i \in \mathcal{B}_0$. By translation invariance, $\hat{F}_0$ simply translates, so
$\tau = \hat{F}_0^{-1}(\alpha)$ shifts by exactly $\delta$. Differentiating with respect to $\delta$ at
$\tau = \hat{F}_0^{-1}(1-\alpha)$ shifts by exactly $\delta$. Differentiating with respect to $\delta$ at
$\delta = 0$:
\[
1 = \frac{\mathrm{d}\tau}{\mathrm{d}\delta}\bigg|_{\delta=0}
@@ -887,25 +619,25 @@ Differentiating at $\delta = 0$:
Combining with \eqref{eq:grad-balance-b0} gives \eqref{eq:grad-balance}.
\end{proof}
The translation-invariance assumption holds for all distribution families considered in this
work: for Gaussian MLE, a uniform shift in scores moves $\mu_k$ by the same amount while
The translation-invariance assumption holds for both the Gaussian distribution,
as well as the KDE-based distribution, but not for the Beta distribution.
For Gaussian MLE, a uniform shift in scores moves $\mu_k$ by the same amount while
leaving $\sigma_k$ unchanged, which translates $\hat{F}_k$; for KDE, each kernel term is
evaluated at $\tau - f_\theta(\mathbf{x}_j)$, so a uniform shift in the $f_\theta(\mathbf{x}_j)$
is equivalent to an equal shift in $\tau$. The condition is necessary: without it, a uniform
score shift need not produce a rigid translation of $\hat{F}_k$, so the shift argument in the
proof breaks down and the gradient sums need not be equal. A natural example where this fails
is a family that first maps scores through a nonlinear link function before fitting, such as
passing scores through a sigmoid into $[0,1]$ before fitting a Beta distribution. A uniform
additive shift of the raw scores then produces a non-uniform shift of the transformed inputs,
changing the shape of the fitted distribution rather than merely translating it, and the
balance property no longer holds.
proof breaks down and the gradient sums need not be equal.
The Beta distribution only accepts values in the interval $\[0, 1\]$. Since it
is mapped only over an interval, the translation-invariance assumption does not hold
for this distribution.
\paragraph{Gradient locality.}
\label{para:gradient-locality}
A second structural consequence of the ROLL formulation is that gradient mass is
A second structural consequence of the ROLL-KDE formulation is that gradient mass is
\emph{concentrated near the operating threshold $\tau$}. Samples whose scores lie
close to $\tau$ receive disproportionately large gradient updates, while samples
close to $\tau$ receive larger gradient updates, while samples
already far on the correct side contribute very little.
This follows directly from \Cref{eq:kde-grad-combined}. For the positive class ($y_i = 1$),
@@ -922,15 +654,14 @@ again bell-shaped and centred on $\tau$.
The effect is that each gradient step focuses the model's attention on the region that
directly determines the operating-point performance: samples that could plausibly be
re-ranked relative to the threshold at the current step. Samples comfortably classified
on the correct side contribute negligible signal, not because they are deemed unimportant,
but because their contribution to the CDF estimate at $\tau$ is already saturated.
re-ranked relative to the threshold at the current step. Samples classified with strong
confidence relative to the threshold, with scores far enough away from said threshold,
receive small gradient updates, since a small change to the score of these samples will not
bring any change to the placement of the threshold itself, or the performance at said threshold.
This stands in contrast to standard losses such as cross-entropy, where the per-sample
gradient is determined by the predicted probability, with no direct dependence on an
operating threshold. Under heavy class imbalance this can cause the dominant class to
``push'' the threshold away from the intended operating point, a problem that gradient
locality avoids by construction.
operating threshold.
\begin{observation}[Gradient locality]
\label{obs:gradient-locality}
@@ -958,32 +689,34 @@ concentration of gradient mass around $\tau$, with smooth decay on either side.
\label{fig:threshold-locality}
\end{figure}
\subsection{Gradient Computation and the Custom Backward Pass}
\label{sec:roll-backward}
% Because the distribution parameters are themselves functions of model scores
% (not fixed), the backward pass must differentiate through the distribution
% fitting step. Describe the custom autograd implementation: how gradients flow
% from the ROLL loss back through the distribution parameters to the model
% weights. Highlight where standard autograd would fail and what the manual
% backward pass does differently.
%------------------------------------------------
\section{Implementation Considerations}
\label{sec:roll-implementation}
\subsection{Numerical Stability and the Scaling Trick}
\subsection{Combatting Numerical Instability With Score Scaling}
\label{sec:roll-numerical-stability}
% The KDE loss involves sums of kernel evaluations that can span many orders of
% magnitude, leading to floating-point underflow or overflow during training.
% Describe the scaling trick that normalizes these sums to a numerically stable
% range without changing the gradient direction, and show that it is equivalent
% to computing in log-space with a max-subtraction stabilizer (analogous to the
% log-sum-exp trick).
Before computing the KDE, all batch scores are divided by their overall standard
deviation:
\[
\tilde{s}_i = \frac{f_\theta(\mathbf{x}_i)}{\hat\sigma_{\mathcal{B}}},
\qquad
\hat\sigma_{\mathcal{B}} = \sqrt{\operatorname{Var}(f_\theta(\mathbf{x}_i) : (\mathbf{x}_i, y_i) \in \mathcal{B})}.
\]
This keeps the values given to the sigmoid kernel ($v(\tau - \tilde{s}_j)$)
near unit scale, preventing floating-point overflow and underflow in the exponential
evaluations.
The scaling has no effect on the loss value itself or derived gradients.
Silverman's rule derives the bandwidth proportionally to the data standard deviation,
$h \propto \hat\sigma \cdot n^{-1/5}$, so dividing all scores by $\hat\sigma_\mathcal{B}$
causes the bandwidth to scale by the same factor.
The product $v \cdot (\tau - s_j)$ — the only quantity entering the sigmoid — is
therefore invariant under the rescaling.
In the backward pass, gradients with respect to the normalised scores are multiplied
by $1/\hat\sigma_\mathcal{B}$ to apply the chain rule, providing gradients mathematically
identical to those that would have been obtained without normalisation.
\subsection{Bandwidth Estimation for KDE}
\label{sec:kde-bandwidth}
@@ -996,17 +729,33 @@ concentration of gradient mass around $\tau$, with smooth decay on either side.
% NOTE - not silverman's rule, we do use scheduling, but we also use different
% estimator - see impl.
When computing KDE, kernel selection and bandwidth estimation affect the final result greatly. Large bandwidths
don't accurately track the probability distribution, while small kernels result in a very jagged CDF, which for us means very low gradient values unless unless a point is very close to the decision threshold.
When computing KDE, bandwidth estimation strongly affects the result. Large bandwidths
produce an overly smooth CDF that fails to track the true score distribution; small
bandwidths produce a jagged CDF, yielding negligible gradients for all but the
nearest-threshold samples.
Use of the Improved Sheather-Jones (ISJ) method~\cite{botev2010kde} of bandwidth selection worked reliably well at creating very close-tracking, yet smooth CDFs.
We estimate the bandwidth using the Improved Sheather-Jones (ISJ)
method~\cite{botev2010kde}, which minimises the asymptotic mean integrated squared
error via an iterative plug-in procedure. ISJ tracks the true score distribution more
closely than closed-form rules, as illustrated in \Cref{fig:density-fit-comparison}.
Because ISJ is iterative, it can fail to converge on small or degenerate batches early
in training. In those cases we fall back to Silverman's rule of
thumb~\cite{silverman1986density},
\[
h = 1.06\,\hat{\sigma}\,n^{-1/5},
\]
which is closed-form and always produces a valid bandwidth.
The kernel parameter is then set to $v = 1/h$.
\subsubsection{Bandwidth Scaling for KDE}
\label{sec:kde-bandwidth-scaling}
While the ISJ method worked very well for selecting an acceptable bandwidth, this wasn't necessarily wanted during training. Near the start of the training session, scores tended to be all low. This naturally made the selected bandwidth to be too low. This caused very low training performance near the start of a training session, which could result in non-convergence of the model even after significant training epochs elapsed.
Although we may select a valid bandwidth at each step, it tends to be
too small near the start of training, when scores are not yet
spread out and $\hat{\sigma}$ is small. This causes the KDE to undersmooth the score
distribution, producing very low gradient magnitudes and stalling convergence.
The solution was to use a bandwidth scheduler. Starting with large bandwidths resulted in gradients that spread out more evenly across all scores, and less non-convergence runs occurred. The bandwidth scheduler artificially scales up the bandwidth at the beginning to speed up initial training, and then tapers back to restore the intended bandwidth in order to maximize resuslts.
The solution was to use a bandwidth scheduler. Starting with large bandwidths resulted in gradients that spread out more evenly across all scores, and fewer non-convergence runs occurred. The bandwidth scheduler artificially scales up the bandwidth at the beginning to speed up initial training, and then tapers back to restore the intended bandwidth in order to maximise results.
% TODO: algorithm box — full training loop (forward → ROLL loss → custom backward → weight update)