Introduction: The Problem with Real-World Problems

Many of us are familiar with optimizing a single metric, like minimizing the error rate of a machine learning model. However, real-world problems are rarely that simple. They often involve juggling multiple, conflicting objectives simultaneously.

Real-world examples:

• Neural Network Design: High accuracy vs. low inference time (faster predictions).
• Mechanical Engineering: Maximum power vs. minimum fuel consumption.
• Cloud Computing: Low cost vs. high performance.

These objective functions are often:

• Expensive to evaluate (each evaluation takes hours or days).
• Black-box (no simple mathematical formula).
• Defined on complex search spaces (mixed real, integer, categorical, and conditional parameters like “if layer X exists, then set parameter Y”).

So, what algorithm can handle all of this efficiently?

The Limitations of Existing Methods

The standard tool for expensive black-box optimization is Bayesian Optimization (BO) using Gaussian Processes (GPs).

GP-based methods (like PESMO, ParEGO, SMS-EGO) are powerful, but they have major drawbacks:

• ❌ Not suitable for non-continuous or conditional (tree-structured) spaces.
• ❌ High computational complexity: O(n³) — they scale poorly with the number of observations.
• ❌ Approximation tricks help but cause performance degradation.

Single-Objective TPE (Tree-Structured Parzen Estimator) is a great alternative that:

• ✅ Naturally handles complex, conditional spaces.
• ✅ Scales to ~1000 observations and tens of variables.
• ✅ Outperforms GP-based methods on single-objective HPO.
• ❌ But: It is designed for single-objective only!

Enter MOTPE: Multiobjective Tree-Structured Parzen Estimator

MOTPE (Ozaki et al., 2020/2022) extends the powerful TPE algorithm to handle multiple objectives. It is designed to be:

How It Works: The Core Idea

1. The Split Rule

In single-objective TPE, observations are split into “good” (below a quantile threshold y*) and “bad” (above). For multiple objectives, MOTPE uses a dominance-based split:

• Good set (l(xᵢ)): Points that are either dominated by or incomparable to the current Pareto front approximation Y*.
• Bad set (g(xᵢ)): Points that weakly dominate Y* (i.e., are clearly worse).

This allows the algorithm to learn which parameter values tend to produce good (non-dominated or diverse) solutions.

2. Greedy Splitting with Hypervolume

To select the top γ% of observations for the good set, MOTPE:

1. Sorts observations by nondomination rank (Pareto sorting).
2. Greedily adds full fronts.
3. For the remaining slots, solves the Hypervolume Subset Selection Problem (HSSP) using a greedy algorithm with a (1 — 1/e)-optimality guarantee.

Points in the good set are then weighted by their hypervolume contribution — points that improve the Pareto front more get higher weight in the density model.

3. The Acquisition Function: Expected Hypervolume Improvement (EHVI)

MOTPE uses Expected Hypervolume Improvement (EHVI) as its acquisition function. Remarkably, after derivation, EHVI simplifies to:

This means: To maximize EHVI, simply maximize the ratio l(xᵢ) / g(xᵢ) — exactly the same as in single-objective TPE! No complex EHVI calculations are needed.

Benchmark Results: How Well Does It Work?

Experiment 1: WFG Benchmark (Low Dimension, Limited Budget)

MOTPE was compared against GP-based methods (PESMO, ParEGO, SMS-EGO) on the WFG test suite with a budget of only 250 evaluations.

Key findings:

• MOTPE achieves comparable or better results on most WFG problems.
• MOTPE is more robust to dimensionality than GP methods.
• Each MOTPE run took minutes; GP runs took hours to days.

Experiment 2: Real-World — CNN Design for CIFAR-10

Task: Design a CNN minimizing two objectives:

1. Classification error rate.
2. Prediction time (inference speed).

Search space: 13 parameters with complex conditionals (e.g., number of blocks determines which filter parameters are active).

Results:

• ✅ MOTPE outperformed all baselines (ParEGO, SMS-EGO, PESMO, HyperMapper 2.0).
• ✅ MOTPE found a better diversity of trade-offs between accuracy and speed.
• ✅ Baselines tended to find either fast-but-inaccurate or accurate-but-slow models; MOTPE found both.

Key Insights: The Quantile Parameter γ

The γ parameter controls the split between good and bad observations. The paper’s investigation of γ reveals:

Empirical recommendation: Start with γ = 0.10 (the winner in 17 out of 72 experiments).

Asynchronous Parallelization: Practical Speedups

MOTPE supports asynchronous parallelization — workers grab the latest observations, run the algorithm, and evaluate candidates without waiting for others.

Speedup results on WFG4:

• 1 worker → 250 minutes
• 10 workers → 28 minutes
• 30 workers → 13 minutes

On the CNN design problem: Parallelization (4 workers) achieved a ~4x speedup (555 min → 142 min to reach the same hypervolume).

When Should You Use MOTPE?

Use MOTPE when:

• ✅ Your search space has conditional parameters (e.g., neural architecture search).
• ✅ You have a limited evaluation budget (tens to hundreds).
• ✅ You need scalability to many parameters or observations.
• ✅ You want asynchronous parallelization without complicated batching.

Be cautious when:

• ⚠️ The objective space is extremely biased (WFG1 case).
• ⚠️ The problem is highly deceptive (WFG5 case) — but note: GP methods also struggled here.
• ⚠️ You have a very large budget (thousands of evaluations) — evolutionary algorithms like NSGA-II may eventually catch up.

Summary

Final Thoughts

MOTPE is a practical, scalable, and effective algorithm for expensive multi-objective optimization in complex, real-world search spaces. If you’re doing neural architecture search, hyperparameter tuning, or any engineering design with multiple objectives and a limited budget, MOTPE deserves a spot in your toolbox.

The algorithm is implemented in Optuna (as MOTPE) and available for use today.

Based on the paper: “Multiobjective Tree-Structured Parzen Estimator” by Ozaki, Tanigaki, Watanabe, Nomura, and Onishi (Journal of Artificial Intelligence Research, 2022).

This post explains how to turn a standard Gaussian-process-style regression problem into a robust, scalable state-space model. The core idea is simple: keep the analytical structure of Gaussian-process inference, but replace the fixed Gaussian uncertainty with a Student-t process so that the model becomes much more tolerant to outliers.

Introduction

Gaussian processes (GPs) are one of the most elegant tools for non-parametric regression. They provide closed-form posterior predictions, uncertainty estimates, and a principled way to encode prior structure through kernels. But in practice, standard GP regression has two major weaknesses:

1. Computational cost: the naive implementation requires matrix inversion and scales as $O(n^3)$.
2. Sensitivity to outliers: the posterior variance depends only on input locations, not on observed values, so a single bad observation does not automatically increase uncertainty.

For many real-world time series, these limitations are a problem. Sensor failures, missing values, and occasional large spikes are common, and a model should react to them instead of treating every observation as equally trustworthy.

The method discussed here replaces the Gaussian assumption with a Student-t process (TP) and then exploits a state-space representation to make inference efficient. The result is a model that is both robust and fast.

Gaussian processes in one page

A Gaussian process is a distribution over functions:

which means that for any finite set of inputs $x_1,\dots,x_n$, the vector of function values is jointly Gaussian:

For GP regression, the posterior mean and variance at a test point $x_*$ are

The problem is visible in the variance formula: the uncertainty depends only on the kernel geometry and the training inputs, not on the actual observed values $y$. If the data contains a strong outlier, the model does not automatically become more uncertain about that region.

Why Student-t processes help

The Student-t distribution has heavier tails than the Gaussian, which makes it much more forgiving to unexpected observations. For a vector $y \in \mathbb{R}^n$,

with density

As $\nu \to \infty$, this distribution approaches a Gaussian.

A useful way to think about the Student-t distribution is as a scale mixture of Gaussians:

So the Student-t is still Gaussian at its core, but with a random scale factor $\gamma$. That single random variable is what makes the model adaptive: unusual observations can be explained by a larger local scale, which effectively reduces their influence.

Student-t process regression

A Student-t process is the function-space analogue of the multivariate Student-t distribution:

if every finite collection of function values follows a multivariate Student-t law.

For regression, the crucial point is that the conditional distribution of one block of variables given another retains the same structure. If the joint vector is Student-t distributed, then the conditional posterior is also Student-t, with a mean similar to the GP case but with a variance scaling term that depends on the observations.

For a test point $x_*$, the predictive distribution can be written as

where

This is the most important difference from GP regression:

• in a GP, uncertainty is fixed by the kernel and the input geometry;
• in a TP, uncertainty grows when the data looks suspicious.

That makes the posterior variance much more informative in the presence of outliers.

The training objective

Hyperparameters are learned by maximizing the marginal likelihood, or equivalently minimizing the negative log marginal likelihood. For the TP model, the objective takes the form

where $\beta = (y-\mu)^T K^{-1}(y-\mu)$.

The learned parameters are:

• $\theta$: kernel hyperparameters,
• $\sigma_n^2$: noise level,
• $\nu$: degrees of freedom, which control tail heaviness.

Smaller $\nu$ means heavier tails and stronger robustness.

From Gaussian processes to state-space models

The key trick behind the efficient algorithm is that many one-dimensional temporal kernels admit a state-space representation. In that form, the GP is no longer computed by manipulating the full covariance matrix directly. Instead, it is represented as a linear dynamical system, and inference can be carried out with Kalman filtering.

A temporal GP

can be rewritten as a continuous-time stochastic differential equation (SDE):

with observations

In discrete time, this becomes

This transformation is the reason the method becomes scalable: rather than working with all pairwise correlations at once, we propagate a compact latent state forward in time.

The Student-t process as a state-space model

The same idea extends to the Student-t process by introducing a random scaling variable $\gamma$. The state-space model becomes a Gaussian SDE with a shared scale:

In discrete time, the same scaling appears in the process covariance:

This gives the model the heavy-tailed behavior of the Student-t process while preserving the linear-Gaussian structure needed for Kalman-style inference.

Filtering: the forward pass

Inference is carried out by a forward recursion that is closely analogous to the Kalman filter.
At each step, the algorithm maintains a current estimate of the latent state together with its uncertainty, plus two Student-t-specific variables that control how strongly the next observation should affect the posterior.

The recursion has four conceptual stages:

Prediction
The model first propagates the previous estimate forward in time using the state-space dynamics.

Innovation
It then compares the prediction with the new observation and measures how surprising that observation is.

Scale update
This is the key Student-t feature. If the innovation is unusually large, the model increases its scale variable, which inflates uncertainty and reduces the influence of that suspicious observation.

State update
Finally, the latent state estimate is corrected using the innovation, but now with the outlier-robust weighting produced by the scale update.

So the forward pass behaves like a Kalman filter, but with an adaptive mechanism that automatically softens the impact of outliers instead of treating every observation equally.

Smoothing: the backward pass

Filtering gives good online estimates, but if the full sequence is available, the posterior can be improved with a backward smoothing pass.
This second pass propagates information in reverse time and refines each latent state using both the past and the future.

The result is a globally consistent trajectory estimate with the same robustness benefits as the forward pass. In practice, smoothing usually produces a cleaner reconstruction than filtering alone, especially when the observations contain noise spikes or corrupted segments.

Marginal likelihood and learning

The model parameters are learned by iterating the filtering procedure and minimizing the negative log marginal likelihood. In the state-space form, the objective accumulates contributions from each step:

The payoff is computational: if the latent state dimension is $m \ll n$, the cost becomes approximately

instead of $O(n^3)$ for the naive covariance-matrix approach.

Experiments

The experiments in the presentation compare naive inference with the state-space implementation on both synthetic and real-world data.

Experiment 1: computational efficiency

The first experiment demonstrates the speed advantage of the state-space formulation. Instead of building and factorizing the full covariance matrix, the model updates a compact latent state sequentially. This is where the linear-time scaling becomes visible in practice.

Figure 1: computational efficiency of naive inference versus the state-space method.

Experiment 2: robustness on synthetic and real data

The second experiment compares several data regimes:

• Synth A: Gaussian noise,
• Synth B: Student-t noise,
• Synth C: a mixture with outliers,
• Electricity: hourly household electricity consumption,
• Stock (Apple): log-prices over a long time span.

These settings test both accuracy and robustness. On clean Gaussian data, the GP and TP behave similarly. On contaminated data, the Student-t model is more stable and produces more meaningful uncertainty estimates.

Figure 2: comparison of naive and state-space inference for GP and TP, with average MSE and log-likelihood.

Experiment 3: missing data interpolation

A particularly useful application is interpolation of missing or unreliable measurements. In such settings, the Student-t process is usually more conservative around suspicious regions and more confident elsewhere, which improves reconstruction quality.

Figure 3: interpolation of missing observations with wider uncertainty around corrupted regions.

Conclusion

Student-t process regression gives a simple but powerful upgrade over standard Gaussian-process regression:

• it keeps the Bayesian, analytical nature of GP inference;
• it reacts to outliers through adaptive uncertainty;
• and, with a state-space representation, it becomes scalable enough for long time series and online inference.

The main message is that robustness and efficiency do not have to be in conflict. By combining heavy-tailed probabilistic modeling with Kalman-style inference, we get a method that is practical for real data and still mathematically elegant.

Main takeaways

• GPs are elegant but brittle when the data contains outliers.
• Student-t processes add robustness through heavy tails and observation-dependent uncertainty.
• State-space inference makes the method scalable, turning $O(n^3)$ inference into an approximately linear-time algorithm for temporal kernels.
• Filtering and smoothing remain analytical, so the model stays interpretable and easy to optimize.

If you find this topic interesting, please check out the original paper!

Introduction

In the realm of statistics, there is a fascinating phenomenon known as Stein’s Paradox. It states that when you need to estimate the means of three or more Gaussian random variables, you actually get a better estimate if you compute them jointly using samples from all of them, rather than estimating each one separately—even if the variables are completely independent!

This mathematical quirk serves as an early motivation for Multi-Task Learning (MTL). In modern machine learning, MTL leverages the shared inductive bias across different tasks to improve overall performance. For instance, in autonomous driving, predicting depth and segmenting pedestrians are seemingly different tasks, yet they are governed by the same physical laws of optics and scene geometry. Why learn the rules of the visual world from scratch for every single task when you can learn them once and share the knowledge?

The Problem with the Standard Approach

Typically, MTL in deep neural networks is implemented via hard parameter sharing: the network has shared parameters $\theta^{sh}$ (the feature extractor/encoder) and task-specific parameters $\theta^{t}$ (the heads/decoders).

The most common way to optimize such a model is by taking a linear combination of empirical losses:

where $c_t$ are static or dynamically computed weights.

The flaw: This approach makes a massive assumption—that the tasks do not compete. In reality, tasks frequently conflict. Improving the loss for task A might degrade the loss for task B. When tasks compete, the linear combination forces an arbitrary trade-off, usually requiring an excruciatingly expensive grid search over the weights $c_t$ to find a “good enough” balance.

The Paradigm Shift: Multi-Objective Optimization

Instead of forcing tasks to cooperate through a weighted sum, Sener and Koltun propose casting MTL explicitly as Multi-Objective Optimization (MOO). In MOO, we accept that tasks conflict and that a single “global optimum” that minimizes all losses simultaneously simply might not exist.

Instead, the goal is to find an optimal trade-off, mathematically defined as a Pareto optimal solution.

Pareto Optimality Defined

To understand the objective, we must define two key concepts:

1. Dominance: A set of parameters $\theta$ dominates another set $\bar{\theta}$ if it is better or equal on all tasks, and strictly better on at least one task. Mathematically: $\forall t, \hat{\mathcal{L}}^t(\theta) \leq \hat{\mathcal{L}}^t(\bar{\theta})$ and $\exists i \text{ s.t. } \hat{\mathcal{L}}^i(\theta) < \hat{\mathcal{L}}^i(\bar{\theta})$.
2. Pareto Optimality: A solution $\theta^\ast$ is Pareto optimal if no other solution dominates it. The set of all such solutions forms the Pareto front.

Our new objective is to use gradient-based algorithms to smoothly navigate our model parameters until we land on this Pareto front.

The Math: MGDA and KKT Conditions

To achieve Pareto optimality, the authors turn to the Multiple Gradient Descent Algorithm (MGDA). MGDA relies on the Karush-Kuhn-Tucker (KKT) conditions for multi-objective optimization.

For a point to be Pareto stationary, two conditions must hold:

1. For task-specific parameters: $\nabla_{\theta^{t}} \hat{\mathcal{L}}^t(\theta^{sh}, \theta^t) = 0$
2. For shared parameters: there must exist weights $\alpha_1, \dots, \alpha_T \geq 0$ where $\sum_{t=1}^T \alpha_t = 1$, such that: \(\sum_{t=1}^T \alpha_t \nabla_{\theta^{sh}} \hat{\mathcal{L}}^t(\theta^{sh}, \theta^t) = 0\)

To satisfy the second condition and find a gradient direction that improves all tasks simultaneously, we must solve an optimization problem on a simplex at each training step:

Geometrically, this is equivalent to finding the minimum-norm point within the convex hull of the task gradients. If you think of each task as a player in a multi-directional tug-of-war over the shared weights, this algorithm calculates the exact center of force where everyone is satisfied.

The Computational Bottleneck

Here is where the elegant theory hits a brick wall of computational reality.

To solve the simplex optimization problem above, we need the gradient of each task’s loss with respect to the shared parameters: $\nabla_{\theta^{sh}} \hat{\mathcal{L}}^t$. In a deep neural network, computing this requires a separate backward pass for each task. If you have 40 tasks, you need 40 backward passes per training step. This linear scaling with the number of tasks makes standard MGDA entirely impractical for deep learning.

The Proposed Solution: MGDA-UB

To bypass this bottleneck, the authors exploit the architecture of neural networks. Let $Z = g(x; \theta^{sh})$ be the shared representations (the output of the shared encoder).

Using the chain rule, we can extract the Jacobian of the representation w.r.t the shared parameters: $\frac{\partial Z}{\partial \theta^{sh}}$. The authors prove the following upper bound:

Notice something beautiful here? The term $\left| \frac{\partial Z}{\partial \theta^{sh}} \right|_2^2$ does not depend on $\alpha$.

This means to find the optimal $\alpha$ weights, we can drop the expensive shared-parameter gradients and optimize only using the gradients with respect to the representations $Z$:

Why is this a game-changer? The gradients with respect to $Z$ ($\nabla_{Z} \hat{\mathcal{L}}^t$) can be computed for all tasks simultaneously in a single backward pass. The computational overhead of MOO effectively drops to zero.

The authors theoretically guarantee this approach with a theorem:

Theorem 1: Assuming $\frac{\partial Z}{\partial \theta^{sh}}$ is full rank, if $\alpha^{1,\dots,T}$ is the solution to the MGDA-UB problem, it will either yield a Pareto stationary point or provide a descent direction that strictly decreases all objective losses.

The Final Algorithm

The training loop elegantly comes together as follows:

1. Update task-specific parameters $\theta^t$ using their respective gradients.
2. Compute the gradients of task losses with respect to the shared representations $Z$.
3. Solve the Frank-Wolfe optimization algorithm on the simplex to find the optimal $\alpha$ weights.
4. Calculate the common descent direction $\Delta = \sum_{t=1}^T \alpha_t \nabla_{\theta^{sh}} \hat{\mathcal{L}}^t$ (effectively done via a standard backward pass weighted by $\alpha$).
5. Update shared parameters $\theta^{sh}$ using $\Delta$.

Experimental Results

To prove that this isn’t just theoretical wizardry, the authors evaluated MGDA-UB on three distinct benchmarks, increasing the complexity and the number of tasks.

1. MultiMNIST (Capacity Competition)

The first experiment uses MultiMNIST, where two digits are overlaid on a single image. The two tasks are classifying the top-left digit and the bottom-right digit.

Because the network has limited capacity, these two tasks heavily compete. As seen in the results, standard linear combinations (even after exhaustive grid search) fail to match the performance of training two entirely separate models (Single Task).

MGDA-UB, however, perfectly navigates the capacity competition, matching single-task performance without needing two separate networks.

Figure 1. Authors plot the obtained accuracy in detecting the left and right digits for all baselines. The grid-search results suggest that the tasks compete for model capacity. Proposed method is the only one that finds a solution that is as good as training a dedicated model for each task. Top-right is better.

Table 1. Comparison of proposed method vs baselines on MultiMNIST.

2. CelebA (Scaling to 40 Tasks)

To test scalability, the authors framed the CelebA facial attribute dataset as a 40-way multi-label classification problem. Grid search is mathematically impossible here.

Despite the massive number of tasks, MGDA-UB seamlessly scaled to 40 objectives, beating Uniform Scaling, Uncertainty Weighting (Kendall et al.), and GradNorm, achieving the lowest average error.

Figure 2. Radar charts of percentage error per attribute on CelebA. Lower is better. Authors divide attributes into two sets for legibility: easy on the left, hard on the right. Zoom in for details.

Table 2. Comparison of proposed method vs baselines on CelebA.

3. Cityscapes (Complex Scene Understanding)

Moving to a real-world autonomous driving analog, the model was tasked with jointly performing semantic segmentation, instance segmentation, and monocular depth estimation on the Cityscapes dataset using a ResNet-50 encoder.

Once again, MGDA-UB achieved state-of-the-art results across all three metrics. It allowed the tasks to actively cooperate, beating single-task baselines across the board.

Figure 3. Authors plot the performance of all baselines for the tasks of semantic segmentation, instance segmentation, and depth estimation. They use mIoU for semantic segmentation, error of per-pixel regression (normalized to image size) for instance segmentation, and disparity error for depth estimation. To convert errors to performance measures, they use 1 - instance error and 1/disparity error. They plot 2D projections of the performance profile for each pair of tasks. Although they plot pairwise projections for visualization, each point in the plots solves all tasks. Top-right is better.

4. The Role of the Approximation (Ablation Study)

A crucial question remains: Does approximating the true MGDA with the Upper Bound (MGDA-UB) hurt performance?

The authors compared exact MGDA (multiple backward passes) against MGDA-UB. The results were startling. On CelebA (40 tasks), the training time dropped from 42.9 hours to just 1.6 hours (a ~25x speedup).

Even more surprisingly, accuracy slightly improved with the approximation. The authors hypothesize that calculating the simplex optimization in the lower-dimensional space of $Z$ (thousands of dimensions) rather than $\theta^{sh}$ (millions of dimensions) significantly reduces gradient noise, leading to higher stability.

Conclusion

Sener and Koltun’s paper shifts the paradigm of Multi-Task Learning from empirical weight guessing to rigorous mathematical optimization.

Key Takeaways:

1. Mathematically Sound: Formulating MTL as finding a Pareto optimum removes the need for heuristic weight tuning.
2. Highly Scalable: The MGDA-UB upper bound reduces an $O(T)$ backward pass bottleneck to a single backward pass, making gradient-based MOO practical for massive deep neural networks.
3. Theoretically Proven: Optimizing the upper bound is mathematically guaranteed to yield a Pareto stationary point (under full-rank Jacobian assumptions).
4. State-of-the-Art: It efficiently utilizes shared model capacity, achieving top performance across digit classification, multi-label prediction (up to 40 tasks), and dense computer vision tasks.

If you find this topic interesting, please check out the original paper!

Introduction

Vector Approximate Message Passing (VAMP) is one of the cornerstones of modern high-dimensional statistical inference. Rooted in the statistical physics of spin glasses, VAMP and its predecessor AMP provide iterative algorithms for solving linear inverse problems of the form $\mathbf{y} = \mathbf{A}\mathbf{x}_0 + \mathbf{w}$. Their key feature is the Onsager correction — a seemingly mysterious term that removes harmful correlations between iterations and ensures that the algorithm’s dynamics can be exactly tracked by State Evolution (SE) in the large-system limit. When the priors and likelihoods are Gaussian, VAMP achieves Bayes-optimal performance.

However, classical VAMP faces a critical bottleneck when we try to apply it to real-world problems. The Onsager correction requires computing the divergence (trace of the Jacobian) of the denoiser $\eta_t(\cdot)$ at every iteration: \(\alpha_t = \frac{1}{N}\operatorname{div}(\eta_t) = \frac{1}{N}\sum_{i=1}^N \frac{\partial \eta_{t,i}}{\partial r_i}.\)

For simple denoisers (e.g., soft-thresholding) this is trivial. But for deep neural network denoisers — the kind that actually work well on natural images, medical imaging, or scientific data — computing this divergence via automatic differentiation is prohibitively expensive. It costs $\mathcal{O}(N^2)$ operations or requires memory-hungry Monte Carlo approximations, effectively negating the computational advantages of message passing.

In their 2026 paper, Wadayama and Takahashi propose Score-based VAMP (SC-VAMP), a elegant reformulation that eliminates the Jacobian entirely. The key insight is that both the optimal denoiser and its Onsager correction can be expressed purely in terms of the score function $\nabla \log p(\mathbf{r})$ and the Fisher information $\mathbb{E}[|\nabla \log p(\mathbf{r})|^2]$. Since score functions are exactly what modern diffusion models learn, SC-VAMP turns any pre-trained score network into a Bayes-optimal inverse problem solver — without ever computing a derivative of the network.

Background: From AMP to VAMP

AMP (Approximate Message Passing) is an iterative algorithm for linear inverse problems that alternates between a linear step and a nonlinear denoising step, with the Onsager correction subtracting the “self-interaction” of the iterate. Its asymptotic dynamics are exactly characterized by SE, a scalar recursion tracking the mean-square error.

VAMP generalizes AMP to handle matrices $\mathbf{A}$ of arbitrary singular value distributions (not just i.i.d. sub-Gaussian ones). It does so by using the SVD of $\mathbf{A}$ and maintaining two modules:

• Module A (LMMSE): processes the linear observations $\mathbf{y} = \mathbf{A}\mathbf{x} + \mathbf{w}$.
• Module B (Denoiser): applies a nonlinear estimator $\eta(\mathbf{r})$ to remove noise.

The Achilles’ heel is the divergence computation in Module B. For a neural denoiser with millions of parameters, $\operatorname{div}(\eta)$ is a nightmare.

Method: Score-Based VAMP

SC-VAMP resolves this by re-parameterizing the entire algorithm in terms of the score function. The paper makes three interconnected contributions.

Tweedie’s Formula and the Score Function

Suppose we observe a noisy vector $\mathbf{r} = \mathbf{x} + \gamma^{-1/2}\mathbf{z}$ where $\mathbf{z} \sim \mathcal{N}(\mathbf{0}, \mathbf{I})$. The optimal MMSE denoiser is given by Tweedie’s formula: \(\hat{\mathbf{x}} = \mathbb{E}[\mathbf{x}|\mathbf{r}] = \mathbf{r} + \frac{1}{\gamma}\nabla_{\mathbf{r}} \log p(\mathbf{r}; \gamma).\)

The gradient term is the $\textbf{score function}$. Modern diffusion models train neural networks $\mathbf{s}\theta(\mathbf{r}, \gamma) \approx \nabla{\mathbf{r}} \log p(\mathbf{r}; \gamma)$. Thus, we can implement the optimal denoiser using only a forward pass through a score network: \(\mathbf{x}_{\text{post}} = \mathbf{r} + \frac{1}{\gamma} \mathbf{s}_\theta(\mathbf{r}, \gamma),\) where $\gamma$ is the inverse input variance (so that $v_{\text{in}} = \gamma^{-1}$).

Jacobian-Free Onsager Correction via Fisher Information

Here is the theoretical centerpiece of the paper. The authors prove that the Onsager coefficient $\alpha(v_{\text{in}})$, which normally requires the divergence of the denoiser, can be computed directly from the conditional Fisher information: \(\alpha(v_{\text{in}}) = 1 - \frac{v_{\text{in}}}{N} J(\gamma),\) where \(J(\gamma) = \mathbb{E}_{\mathbf{r}}\left[\left\|\nabla_{\mathbf{r}} \log p(\mathbf{r}; \gamma)\right\|^2\right] = \mathbb{E}_{\mathbf{r}}\left[\left\|\mathbf{s}_\theta(\mathbf{r}, \gamma)\right\|^2\right].\)

In other words: the Onsager correction is determined by the expected squared norm of the score function. No Jacobians. No backpropagation through the denoiser. Just evaluate the score network on a mini-batch, average the squared $\ell_2$-norms, and plug into the formula.

The paper provides multiple derivations of this identity — via Stein’s identity, via the I-MMSE relationship (de Bruijn’s identity), and via the information-geometric interpretation of the Onsager term as a “curvature correction” governed by the local Fisher information.

The SC-VAMP Algorithm

Putting it together, one iteration of SC-VAMP looks like this:

1. LMMSE step (Module A): Update using linear measurements (standard VAMP).
2. Denoising step (Module B): \(\mathbf{x}_{1,t} = \mathbf{r}_{1,t} + v_{1,t} \mathbf{s}_\theta(\mathbf{r}_{1,t}, v_{1,t}).\)
3. Onsager correction: \(\alpha_{1,t} = 1 - \frac{v_{1,t}}{N} \hat{J}_\theta, \quad \hat{J}_\theta = \frac{1}{B}\sum_{i=1}^B \|\mathbf{s}_\theta(\mathbf{r}_{1,t}^{(i)})\|^2.\)
4. Extrinsic output: \(\mathbf{r}_{2,t} = \frac{\mathbf{x}_{1,t} - \alpha_{1,t}\mathbf{r}_{1,t}}{1 - \alpha_{1,t}}.\)

The algorithm is Jacobian-free: the only operation involving the neural network is a forward pass to get the score. This reduces the per-iteration cost by a factor of 10–50× compared to standard VAMP with AutoDiff-based divergence computation.

Using Pre-Trained Denoisers

A practical bonus: if you already have a state-of-the-art denoiser $\eta_{\text{opt}}(\cdot)$ (e.g., DnCNN, DRUNet) but no explicit score model, you can extract an implicit score via Tweedie’s formula: \(\hat{\mathbf{s}}(\mathbf{r}) = \frac{\eta_{\text{opt}}(\mathbf{r}) - \mathbf{r}}{v_{\text{in}}}.\)

Substituting this into the Fisher information estimator yields a fully plug-and-play Onsager correction, letting SC-VAMP leverage existing denoiser libraries without retraining.

Theoretical Guarantees

The paper demonstrates that SC-VAMP is not just a computational hack — it is theoretically sound.

Optimality in Scalar Gaussian Channels

Theorem 1 shows that in the classical linear Gaussian setting, SC-VAMP reduces exactly to standard Bayes-optimal VAMP. For a scalar channel $Y = X + Z$ with $X \sim \mathcal{N}(0, P)$ and $Z \sim \mathcal{N}(0, \sigma^2)$, the SE fixed point of SC-VAMP achieves the mutual information: \(I_{\text{VAMP}} = I(X;Y) = \frac{1}{2}\log\left(1 + \frac{P}{\sigma^2}\right).\)

Moreover, the point estimate converges to the Wiener filter $\hat{x} = \frac{P}{P+\sigma^2}y$. Thus, SC-VAMP preserves the optimality of VAMP exactly where VAMP is already optimal.

State Evolution and Decoupling

The authors verify empirically that the MSE trajectory of SC-VAMP follows the theoretical SE prediction precisely:

Figure 1: MSE convergence of SC-VAMP (blue) versus State Evolution theory (red dashed). The trajectories are indistinguishable, confirming that the Fisher-information-based Onsager correction is unbiased.

This confirms that the score-norm approximation does not break the decoupling principle. The algorithm still decomposes the high-dimensional problem into independent scalar Gaussian channels in the large-system limit.

Information-Theoretic Perspective

Perhaps the most conceptually rich part of the paper is the connection to the entropic Central Limit Theorem. The authors interpret the linear mixing step in VAMP as a “Gaussianizer”: each iteration reduces the non-Gaussianity (in KL-divergence) of the estimation error. This provides an information-theoretic justification for why the Gaussian approximation underlying SE remains valid even beyond idealized i.i.d. settings, including nonlinear regimes.

Experiments

The experimental section focuses on a linear observation system with a Bernoulli-Gaussian prior ($N=2000$, $M=1000$, $\rho=0.1$, SNR = 20 dB). The score function is learned via denoising score matching (DSM).

MSE and Convergence

As shown in Figure 1, SC-VAMP tracks the theoretical SE curve almost perfectly. The algorithm converges to the same fixed point as classical VAMP with exact divergence, but with drastically lower computational cost.

EXIT Chart Analysis

Figure 2: EXIT-style analysis showing Module A (observation) and Module B (denoiser) transfer characteristics. The SE trajectory (green) and actual SC-VAMP trajectory (gray dashed) follow the characteristic curves and converge to the same fixed point.

The EXIT chart confirms that the score-based SISO modules correctly implement the MMSE estimator and that the mini-batch Fisher information estimator provides an accurate Onsager term.

Extensions and Future Directions

The paper sketches several promising extensions:

• Random orthogonal/unitary mixing: To handle structured or correlated sensing matrices (where standard VAMP/AMP often fails), SC-VAMP can be combined with random rotations that “whiten” the problem.
• Nonlinear observations: The score-based formalism extends naturally to $\mathbf{y} = f(\mathbf{x}) + \mathbf{w}$ with deterministic nonlinearities $f$, such as sensor saturation or optical systems.
• Flow matching: The authors note that since score functions and velocity fields are algebraically equivalent under a given probability path, an SC-VAMP-like algorithm could be built using rectified flow or flow matching networks — potentially offering more stable training in low-noise regimes.

Conclusion

SC-VAMP represents a significant step in unifying classical statistical-physics-based inference with modern deep learning. By reformulating the denoiser and its Onsager correction entirely through the lens of score functions and Fisher information, it eliminates the Jacobian bottleneck that has long plagued neural AMP/VAMP methods.

Main strengths of SC-VAMP:

• Jacobian-free: The Onsager correction requires only the squared norm of the score network output, enabling 10–50× speedup per iteration.
• Plug-and-play: Works with any pre-trained score model or denoiser; no retraining or architectural constraints.
• Theoretically grounded: Recovers exact Bayes-optimal VAMP in Gaussian settings and tracks State Evolution perfectly.
• Universal: Extends to nonlinear observations, complex priors, and structured sensing matrices.

Limitations:

• Decoupling assumption: Like all AMP/VAMP variants, SC-VAMP relies on asymptotic statistical decoupling; for highly structured finite-dimensional problems, the Gaussian approximation may be less accurate.
• Score estimation quality: The method inherits any bias or variance from the learned score network. In low-noise regimes, score matching can be unstable (though flow-matching extensions may remedy this).
• Mini-batch variance: The Fisher information is estimated via Monte Carlo; very small batch sizes could introduce variance into the Onsager term.

SC-VAMP is a rare example of a method that is simultaneously cheaper, more general, and theoretically cleaner than its predecessor. It opens the door to applying message-passing algorithms to the complex, black-box inference problems that arise in computational imaging, scientific computing, and beyond.

Many machine-learning problems are really problems about hidden motion. We do not directly see the true state of a system: the position and velocity of a car, the internal state of a robot, the real phase of an epidemic, or the underlying dynamics of a noisy sensor. We only see imperfect measurements. A state-space model is a standard way to describe this situation.

The paper Ensemble Kalman Filtering Meets Gaussian Process SSM for Non-Mean-Field and Online Inference asks a practical question: can we learn an unknown nonlinear dynamical system from noisy observations while also estimating the hidden states? The authors combine two ideas that are powerful in different ways:

• Gaussian processes (GPs): flexible Bayesian models for unknown functions, with uncertainty.
• Ensemble Kalman filtering (EnKF): a fast filtering method that tracks hidden states using a cloud of particles.

Their method is called EnVI: EnKF-aided Variational Inference. The online version is called OEnVI. The main message is simple: instead of training a large neural inference network to guess hidden states, use a model-based filter to do that job, and let the GP focus on learning the dynamics.

One-sentence intuition: the GP learns “how the system tends to move”, while the EnKF continuously asks “given what we just observed, where is the system now?”

The modeling problem: hidden states and noisy observations

A state-space model has two parts. The first part says how the hidden state evolves. The second part says how observations are generated from the hidden state:

Here $x_t$ is the hidden state, $y_t$ is the observation, $f$ is the transition function, and $v_t, e_t$ are noise terms. The matrix $C$ maps the hidden state to the observation space.

If $f$ is known and the model is linear-Gaussian, the classical Kalman filter is almost ideal. But in many realistic problems, $f$ is not known. We may only have a sequence of noisy observations and need to learn both:

1. the hidden trajectory $x_0, x_1, \ldots, x_T$;
2. the transition rule $f$ that generated it.

This creates a chicken-and-egg problem. To learn $f$, we need good estimates of the hidden states. To infer the hidden states, we need a good $f$.

Gaussian processes: learning a function with uncertainty

A Gaussian process is a distribution over functions. Instead of saying “the transition function must be a neural network with these weights”, a GP says: before seeing data, plausible functions are those that look smooth according to a kernel $k$.

A compact way to write this is:

After seeing data, the GP gives two things at a new input:

• a mean prediction: the most likely function value;
• a variance: how uncertain the model is there.

This uncertainty is very useful in dynamical systems. If we ask the model to predict in a region it has not seen, it should not pretend to be confident. This is one reason GP state-space models, or GPSSMs, are attractive for small and medium datasets.

The difficulty is computational. A full GP becomes expensive for long time series, and in GPSSMs the inputs to the GP are hidden states, not observed data. The paper therefore uses a standard sparse-GP trick: inducing points. These are a small set of representative pseudo-inputs that summarize the transition function. Instead of carrying the whole GP over every time point, the algorithm learns a compact surrogate.

Kalman filtering and EnKF: tracking the hidden state

The Kalman filter alternates between two steps:

1. Predict: use the dynamics to move the previous state estimate forward.
2. Update: correct the prediction using the new observation.

In the linear-Gaussian case, this is exact. In nonlinear systems, the Ensemble Kalman Filter keeps a cloud of particles, called an ensemble, and moves each particle through the dynamics. Then it updates the whole cloud using a Kalman-style correction.

A simplified update for the mean looks like this:

where $\bar m_t$ is the predicted mean, $y_t - C\bar m_t$ is the surprise in the new observation, and $G_t$ is the Kalman gain. The gain decides how much to trust the observation versus the model prediction.

For this paper, the important point is not only that EnKF is fast. It is also differentiable when implemented carefully with reparameterized noise. That means gradients can flow through the filtering procedure, so the GP parameters and variational parameters can be optimized with tools such as automatic differentiation.

How EnVI unites GPSSMs and EnKF

The paper’s central idea is to put EnKF inside variational inference for GPSSMs. Variational inference normally introduces an approximate posterior distribution $q$ and optimizes an evidence lower bound, or ELBO. In many previous GPSSM methods, the distribution over hidden states is parameterized by many extra variables or by an inference network. That can be slow, unstable, and awkward for online learning.

EnVI changes the design:

Noisy observations y_t
        |
        v
EnKF estimates hidden states x_t
        |
        v
Sparse GP learns transition f(x_t)
        |
        v
ELBO balances data fit + regularization
        |
        v
Updated model for filtering and forecasting

The approximate objective derived in the paper can be read as:

The first term rewards predictions that explain the observations. The two KL terms act as regularizers: the initial state and the GP transition should not drift too far from their priors unless the data strongly supports it.

This is a nice objective because it is interpretable. The algorithm is not just fitting observations. It is also controlling model complexity and uncertainty.

Why “non-mean-field” matters

A mean-field approximation breaks dependencies between groups of variables. This often makes optimization easier, but in a dynamical model it can be too aggressive. The hidden states and transition function are deeply linked: changing the transition function changes the plausible hidden trajectory, and changing the hidden trajectory changes what transition function is learned.

EnVI keeps this relationship more naturally. The EnKF state estimates depend on the GP transition, and the GP is learned from those filtered states. This is why the paper calls the method non-mean-field: it does not pretend that the latent states and GP dynamics are independent.

Online learning: OEnVI

The online version, OEnVI, processes data one time step at a time. At each new observation it performs the same basic cycle:

1. sample or use the current GP surrogate;
2. predict the ensemble forward;
3. update the ensemble with the new observation;
4. update model and variational parameters using the local objective.

The online objective has the same spirit as the offline one:

This matters because many systems do not arrive as a fixed dataset. Sensors, robots, vehicles, and monitoring systems stream data continuously. A method that requires the entire sequence at training time is less convenient there. OEnVI is designed to update as data arrives.

What tasks does this combination solve?

The GP-EnKF combination is useful for several related tasks:

The key design choice is that EnKF handles the state-estimation part, while the GP handles the unknown-function part. Variational inference ties them together with a principled training objective.

Experimental results: what was most impressive?

The authors test the methods on synthetic and real datasets. The exact numbers are less important than the pattern: EnVI is usually more accurate, more robust, or faster to train than competing GPSSM and neural state-space methods.

1. Linear-Gaussian tracking: close to the Kalman filter

The authors first use a setting where the classical Kalman filter is available as a strong reference: a linear-Gaussian car-tracking model. EnVI and OEnVI do not receive the true physical transition model, only noisy observations.

Even so, their state estimates are close to the Kalman filter baseline. The reported latent-state RMSE values are:

• Kalman filter: 0.5252;
• EnVI: 0.6841;
• OEnVI: 0.7784;
• raw observations versus latent states: 0.9872.

This experiment is a sanity check. If a learned GPSSM cannot do well on a linear-Gaussian system, it is hard to trust it on nonlinear ones. EnVI passes this check convincingly. OEnVI is less accurate at the beginning because it learns sequentially, but the paper reports that after more online data it improves substantially.

2. Kink function: learning nonlinear dynamics under noise

The kink function is a classic GPSSM test. It is a one-dimensional nonlinear transition that is simple enough to visualize but hard enough to expose bad uncertainty estimates.

Across three observation-noise levels, EnVI obtains the best transition-function fit among the compared methods. For example, at the lowest noise level, EnVI reports MSE 0.0046, while AD-EnKF reports 0.0285, VCDT 0.2057, and vGPSSM 1.0410. At high noise, EnVI still performs best: MSE 0.5315, compared with 1.3489 for AD-EnKF, 1.4035 for VCDT, and 1.9584 for vGPSSM.

The visual result is also important: EnVI learns both the shape of the kink and a reasonable uncertainty band. The paper argues that AD-EnKF can become overconfident because it uses a deterministic neural transition model, while EnVI keeps uncertainty through the GP.

Another striking result is convergence. On the kink experiment, EnVI reaches good performance after roughly 300 iterations, whereas vGPSSM and VCDT require many more iterations and more runtime. This supports the authors’ claim that removing a large inference network makes optimization easier.

3. Real time-series forecasting: strong small-data performance

The paper also evaluates five public system-identification datasets: Actuator, Ball Beam, Drive, Dryer, and Gas Furnace. The model trains on the first half of each sequence and forecasts the second half. The reported metric is 50-step-ahead RMSE.

EnVI is best on four of the five datasets and competitive on the fifth. Its RMSEs are:

• Actuator: 0.657;
• Ball Beam: 0.055;
• Drive: 0.703;
• Dryer: 0.125;
• Gas Furnace: 1.388.

The Drive dataset is the main exception: PRSSM reports 0.647, better than EnVI’s 0.703. But overall, the results are strong, especially because these datasets are relatively small. That is exactly the regime where GP-based models can be attractive compared with large neural models.

4. Online NASCAR dynamics: OEnVI wins clearly

For online learning, the authors use a NASCAR-shaped latent trajectory and compare OEnVI with SVMC and VJF. The prediction RMSEs are:

• OEnVI: 1.8780;
• SVMC: 4.6682;
• VJF: 10.8499.

This is one of the clearest experimental wins in the paper. OEnVI tracks and predicts the latent trajectory much better than the alternatives. The authors attribute this to the EnKF-based approximation of the latent-state distribution: it is structured enough to be stable, but flexible enough to work online.

Takeaways and limitations

The paper is interesting because it does not simply replace everything with a neural network. It combines a probabilistic non-parametric model with a classical filtering algorithm, then trains the whole system with variational inference.

The main takeaways are:

• GPSSMs are useful when the transition dynamics are unknown and nonlinear.
• EnKF provides a practical way to infer hidden states without a heavy inference network.
• The resulting ELBO has a clear interpretation: fit the observations, but regularize the initial state and transition function.
• OEnVI makes the method suitable for streaming data.
• The strongest empirical results are on nonlinear dynamics learning and online tracking.

There are also natural limitations. EnKF relies on Gaussian-style updates, so it may struggle in strongly non-Gaussian settings where particle filters are more appropriate. The paper mainly uses a linear emission model, although EnKF can be extended to nonlinear emissions. Finally, the conclusion notes that time-varying dynamical systems remain an important direction for future work.

Useful links for readers

• Original paper on arXiv: Ensemble Kalman Filtering Meets Gaussian Process SSM
• Authors’ GPSSM code repository: zhidilin/gpssmProj
• A classic free book on Gaussian processes: Gaussian Processes for Machine Learning
• Intuitive Kalman filtering book with Python code: Kalman and Bayesian Filters in Python
• Practical GP library documentation: GPyTorch variational and approximate GPs

In 2018, Frankle and Carbin introduced the Lottery Ticket Hypothesis (LTH): a dense, randomly-initialized neural network contains a sparse subnetwork, a winning ticket, that, when trained in isolation from the original initialization, matches the accuracy of the full network. Empirically, such tickets can be 10 to 20 times smaller than the parent network.

LTH was a striking observation, but it left a theoretical gap. Why do such tickets exist? Are they an artifact of optimization, or a property of random networks themselves?

In Proving the Lottery Ticket Hypothesis: Pruning is All You Need (ICML 2020), Malach, Yehudai, Shalev-Shwartz, and Shamir prove a much stronger version of LTH. Their result, sometimes called the Strong Lottery Ticket Hypothesis, says:

A sufficiently overparameterized randomly-initialized network contains a subnetwork that approximates any target function, with no training at all. Pruning alone is enough.

This post walks through the statement, the proof idea, and what it does (and does not) say about deep learning in practice.

The setup

Fix a target ReLU network $f^{\star}$ of depth $\ell$ and width $n$, with bounded weights $\lVert w \rVert \le 1$. We want to approximate $f^{\star}$ within error $\varepsilon$ on the unit ball.

The construction is:

1. Take a randomly-initialized ReLU network $G$ of depth $2\ell$ and width polynomial in $n, d, 1/\varepsilon, \log(1/\delta)$.
2. Do not train it. Instead, select a binary mask $M$ over its weights.
3. Show that $M \odot G$ approximates $f^{\star}$ uniformly with high probability $1 - \delta$.

The notable points: depth only doubles, width is polynomial, weights are never modified; they are only kept or zeroed out.

The main theorem (informal)

Let $f^{\star}$ be any ReLU network of depth $\ell$, width $n$, and bounded weights. For any $\varepsilon, \delta > 0$, a random ReLU network of depth $2\ell$ and width polynomial in $n, \ell, d, 1/\varepsilon, \log(1/\delta)$ contains, with probability at least $1 - \delta$, a subnetwork that uniformly $\varepsilon$-approximates $f^{\star}$ on the unit ball. The number of active (non-pruned) weights in the subnetwork is $O(dn + n^2 \ell)$, the same order as the parameter count of $f^{\star}$.

In other words: expressivity of pruning a random net equals expressivity of training a net of the same target size, up to polynomial overhead in width.

The paper focuses on weight-pruning, where individual weights can be zeroed independently. A separate result (Theorem 3.2) shows that the strictly weaker model of neuron-pruning, in which only entire neurons can be removed, is equivalent to random features and therefore cannot achieve the same expressive power. This post discusses only the weight-pruning result.

The proof idea: approximate each weight by two random ReLUs

The core trick is local: replace each scalar weight $w \in [-1, 1]$ in the target network by a tiny gadget built from random neurons in the random network.

Step 1: a single weight from two ReLUs

The building block is the identity \(a = \sigma(a) - \sigma(-a)\) where $\sigma$ is the ReLU. Applied coordinate-wise, this lets us write any signed product $w \cdot x$ as \(w \cdot x = \sigma(w \cdot x) - \sigma(-w \cdot x)\) so two ReLU units with input weights $+w$ and $-w$ (and output signs $+1$ and $-1$) reproduce the linear function $x \mapsto w \cdot x$ exactly.

In the random network we do not have $\pm w$ available, but we have a pool of many random scalars. The pool is wide enough that, with high probability, it contains a pair of values close to $+w$ and $-w$ within accuracy $\varepsilon_0$. Pruning everything else leaves a two-neuron gadget that approximates $w \cdot x$.

Step 2: pick by pruning

The random network has many neurons in each layer. For each target weight, the construction prunes away all neurons in the pool except the two chosen ones approximating $+w$ and $-w$. Combined with the next layer’s $+1/-1$ outgoing weights (also obtained by pruning), the remaining two neurons implement an approximation of $w \cdot x$.

Step 3: error accumulation

Each weight is approximated to error $\varepsilon_0$. The approximation errors propagate through $\ell$ layers, and by Lipschitz arguments on bounded-weight networks, the total error is $\mathrm{poly}(\ell, n) \cdot \varepsilon_0$. Choosing $\varepsilon_0$ small enough gives uniform $\varepsilon$-approximation of $f^{\star}$.

Why depth doubles

Each target layer is implemented by two layers of the random network: one for the pool of random scalars feeding the ReLU pair, one for combining them with $\pm 1$ outgoing weights. Hence depth $2\ell$.

Why width is polynomial

For each target weight, the pool has to be large enough that two random scalars fall within $\varepsilon_0$ of any target value in $[-1, 1]$. A covering argument on $[-1, 1]$ shows this requires $\widetilde{O}(1/\varepsilon_0^{2})$ random samples per weight, and a union bound over all $n^2 \ell$ weights yields a polynomial width.

What this changes about how we think about pruning

Pruning is at least as expressive as training. Up to polynomial overhead in width, anything a trained network of size $n$ can represent can also be represented by pruning a random network of size $\mathrm{poly}(n)$. Empirical pruning algorithms are searching over a hypothesis class that is, in principle, rich enough.

The winning ticket is universal. The same random network contains tickets for every target; the choice of mask depends on $f^{\star}$, but the random weights do not.

It partially explains why pruning works so well after training. If subnetworks of the initial random network can already approximate good functions, it is not surprising that subnetworks of a trained network can match its accuracy.

Important caveats

The result is an existence theorem, and several gaps remain between the proof and practice:

1. No algorithm. The proof tells us a good mask exists but not how to find it. Finding the optimal subnetwork is, in the worst case, computationally hard. Practical methods (magnitude pruning, edge-popup, lottery-ticket rewinding) are heuristics.
2. Polynomial overparameterization. The width bound is polynomial in the target size, the input dimension, and $1/\varepsilon$. The exponents in the original proof are not tight, and for realistic $n$ the implied constants are large.
3. Bounded weights and bounded inputs. The construction requires weight and input norms to be bounded, which is the standard but non-trivial assumption.
4. No generalization claim. The theorem is about approximation, not learning. It does not say that a pruned random network will generalize from finite data the way a trained one does.

Follow-up work has tightened these bounds. Pensia et al. (2020) showed that the width overhead can be reduced to logarithmic in $1/\varepsilon$, bringing the strong LTH much closer to a tight statement.

Takeaways

Frankle and Carbin’s empirical LTH said: train, then prune to find a winning ticket. Malach et al.’s strong LTH says: you don’t even need to train. The winning ticket is already there in the random initialization, waiting to be uncovered by a mask.

The proof is a clean, constructive argument: approximate every target weight by a difference of two random ReLUs, prune away the rest, and propagate the error through the layers.

The result is theoretical, and finding the mask remains the hard part, but it cleanly separates expressivity (which pruning a random net already has) from the algorithmic question of how to find a good subnetwork.

Pruning, viewed this way, is not a compression heuristic applied after the fact. It is a legitimate alternative form of “training”, and the strong LTH is the formal statement that it is, in principle, sufficient.

References

1. Malach, E., Yehudai, G., Shalev-Shwartz, S., & Shamir, O. (2020). Proving the Lottery Ticket Hypothesis: Pruning is All You Need. ICML 2020. PDF
2. Frankle, J., & Carbin, M. (2018). The Lottery Ticket Hypothesis: Finding Sparse, Trainable Neural Networks. arXiv preprint, later published at ICLR 2019. arXiv:1803.03635
3. Pensia, A., Rajput, S., Nagle, A., Vishwakarma, H., & Papailiopoulos, D. (2020). Optimal Lottery Tickets via Subset Sum: Logarithmic Over-Parameterization is Sufficient. NeurIPS 2020. arXiv:2006.07990
4. Ramanujan, V., Wortsman, M., Kembhavi, A., Farhadi, A., & Rastegari, M. (2020). What’s Hidden in a Randomly Weighted Neural Network? CVPR 2020. arXiv:1911.13299

If you find this topic interesting, please check out the original paper!

Introduction

Hyperparameter optimization (HPO) is a rapidly developing direction in the field of machine learning and optimization. It considers the automatic optimization of hyperparameters (outer optimization), for example, optimizer parameters like learning rate or weight decay, on top of the optimization of model parameters (inner optimization). In this post we will be looking at one of the gradient-based HPO methods, i.e. methods that rely on the differentiability of certain hyperparameters for their optimization. Such methods are able to utilize gradient information rather than relying on trial-and-error and thus have earned a widespread popularity in the context of few-shot meta-learning. However, at the time these methods were broadly impractical for long-horizon tasks (tasks with many gradient steps in each training cycle). But why?

Problems of previous methods

Typically previous gradient-based HPO methods relied on backpropagation through time (BPTT). Unfortunately, this procedure is extremely expensive both in time and memory, and because of that most previously proposed methods were limited to toy models and datasets. Moreover, long optimization horizons cause hypergradient degradation (i.e. exploding or vanishing hypergradients).

One type of methods that allows to alleviate both these problems is greedy methods. This refers to finding the best hyperparameters locally rather than globally, typically by splitting the inner optimization problem into smaller chunks (often just one batch) and solving for hyperparameters over these smaller horizons instead. However, such methods had been found to introduce bias in the HPO process and thus solve for the wrong objective. The paper we’ll discuss focuses on extending gradient-based methods to the non-greedy setting.

Many previously existing methods were not gradient-based. The most popular ones at the time were black-box methods like Hyperband and its combination with Bayesian optimization called BOHB. These methods rely on trial-and-error, and, as we will see later on, reach optima way slower than gradient-based alternatives such as the one we’ll look at in this post.

Method: Forward-Mode Differentiation with Hyperparameter Sharing

In their paper, Micaelli and Storkey introduce Forward-mode Differentiation with hyperparameter Sharing (FDS), which proposes the following solutions to the aforementioned problems:

1) The use of gradients for optimization allows to reach global optima faster than by trial-and-error; 2) Forward-mode differentiation solves the memory efficiency problem, boasting a memory cost constant with optimization horizon size; 3) Hyperparameter sharing tackles gradient degradation by averaging hyperparameters over time.

Let’s see how this method works step by step.

Problem statement

Let’s denote:

• $\boldsymbol{\theta}$ - the weights of the given neural network model.
• $\mathcal{L}$ - the loss function to be optimized.
• $\mathcal{D}$ - a dataset with train split $\mathcal{D}_\text{train}$ and validation split $\mathcal{D}_\text{val}$.
• $\Phi$ - a gradient-based optimizer for $\boldsymbol{\theta}$.
• $\boldsymbol{\lambda}_{[t]}$ - the set of hyperparameters that $\Phi$ uses for the optimization step $t$ (to get $\boldsymbol{\theta}_{t}$ from $\boldsymbol{\theta}_{t-1}$). Note that this implies that $\boldsymbol{\theta}_t = \boldsymbol{\theta}_t(\boldsymbol{\lambda}_{[1:t]})$.
• $\boldsymbol{\lambda} = \boldsymbol{\lambda}_{[1:T]}$ - the full set of hyperparameters used by $\Phi$ for optimization.
• $T$ - the number of optimization steps $\Phi$ takes.

Our task is to find the optimal set of hyperparameters $\boldsymbol{\lambda}^*$ such that the result at time $T$ of the gradient process optimizing the train loss $\mathcal{L}_\text{train}$ also minimizes the generalization loss $\mathcal{L}_\text{val}$ on the validation set $\mathcal{D}_\text{val}$:

Here, the inner optimization loop, optimizing $\boldsymbol{\theta}$, expresses a constraint on the outer loop, optimizing $\boldsymbol{\lambda}$.

Let $H$ be the horizon, which corresponds to the number of optimization steps taken in the inner loop before a step is taken in the outer loop (optimizing the hyperparameters). If we solve this problem non-greedily, we have $T=H$. This means that non-greedy methods, like FDS, only update $\boldsymbol{\lambda}_{[t]}$ at time $T$. If we, on the other hand, consider a greedy approach, we get $H\ll T$. For example, Hypergradient Descent (HD), a standard gradient-based HPO method, uses $H=1$.

The memory cost of BPTT, the go-to method for solving the optimization problem above, is $\mathcal{O}(DH)$, where $D$ is the number of weights. This estimate scales unfavorably when the optimization horizon is long. Greedy approaches help mitigate the memory scaling problem by minimizing $H$, yet bring about problems with minimizing the real objective. Forward-mode differentiation aims to improve the memory cost even in the non-greedy case.

Forward-mode differentiation

Let’s consider the general case of using one hyperparameter $\boldsymbol{\lambda}_t$ per step. First, we use the chain rule, knowing that $\partial\mathcal{L}_\text{val}/\partial\boldsymbol{\lambda} = 0$ since the loss function doesn’t directly depend on the hyperparameters:

The first can be calculated as usual through backpropagation. The second term can be calculated recursively, again using the chain rule:

We can write this as

The expressions for $\mathbf{A}_t$ and $\mathbf{B}_t$ depend on the specific hyperparameters used. The authors give an example for SGD with momentum with learning rate $α_t$, momentum $β_t$, weight decay $ξ_t$ and velocity $\mathbf{\nu}_t=\beta_t\mathbf{\nu}_{t-1}+\partial\mathcal{L}_\text{train}/\partial\boldsymbol{\theta}_{t-1})+\xi_t\boldsymbol{\theta}_{t-1}$:

Here, a further recursive term $\mathbf{C}_t = (\partial\mathbf{v}_t/\partial\boldsymbol{\lambda})$ must be considered to get exact hypergradients.

Forward-mode differentiation scales in memory as $\mathcal{O}(DN)$, where $N$ is the number of learnable hyperparameters. The additional scaling by $N$ is a limitation in case we learn one hyperparameter per inner step ($N=T$). However, we can conveniently allow for smaller values of $N$ using hyperparameter sharing.

Hyperparameter sharing

As we noticed earlier, one problem of non-greedy HPO methods is gradient degradation. Specifically, small changes in initial parameters like weight initialization and minibatch ordering can drastically affect hypergradients, introducing large fluctuations. Ideally, hyperparameters should be agnostic to such factors, so we would like to average out their effect on hypergradients. However, the most obvious way of doing that, ensemble averaging, has very high computational and memory cost. FDS utilizes a different strategy - time averaging.

The idea of time averaging is to average out hypergradients across the inner training loop rather than the outer loop. Specifically, in FDS we average out hypergradients from $W$ neighboring time steps in the inner loop, which is, in fact, equivalent to sharing one hyperparameter over all these steps. This helps reduce gradient degradation, but introduces a bias, since in general ensemble averaging and time averaging aren’t equivalent. Nevertheless, the authors manage to prove that the hypergradient error $\text{MSE}_W$ with sharing (specifically, the mean variance of the hypergradient, given that it can be approximated with a Gaussian) has the following upper bound:

where $c$ is the maximum absolute correlation between hypergradients, $L$ is the Lipschitz constant for the network, and $\text{MSE}_1$ is the hypergradient error without hyperparameter sharing. In fact, this entails that for sufficiently small $c$ and $L$, we actually end up with $\text{MSE}_W < \text{MSE}_1$ for some positive $W$.

Experiments

The paper’s authors conducted several experiments to test FDS on long-horizon tasks in comparison with state-of-the-art methods at the time to test whether the proposed method shows its main benefits in practice.

The effect of hyperparameter sharing on hypergradient noise

Figure 1: Hypergradients on SVHN for 100 seeds in the non-greedy (left) and greedy (middle) setting. The mean squared error is also shown (right).

The first experiment tested how well hyperparameter sharing dealt with hypergradient noise and how it affected the hypergradient error, training the learning rate for LeNet on the SVHN dataset. The method did, in fact, outperform the greedy setting, significantly reducing noise for many values of $W$. Interestingly, the mean squared error of the gradient has also significantly reduced, showing the best result for $W=50$.

The effect of hyperparameter sharing on HPO

Figure 2: The learning rate schedule learned on MNIST and SVHN using LeNet.

The next two experiments analyze how hyperparameter sharing affects performance of HPO on various real datasets. In the Figure 2, we can see the results of training the learning rate for LeNet on MNIST and SVHN. Despite LeNet being a relatively small architecture, making non-greedy HPO a viable option, both greedy and non-greedy HPO fail to find reasonable learning rates for training on SVHN, probably due to hypergradient variance. On the other hand, FDS stabilizes non-greedy hypergradients and allows to find learning rates that even outperform reasonable off-the-shelf schedules (cosine annealing in this case).

Figure 3: FDS applied to SGD with momentum to learn the learning rate schedule $\alpha$, momentum $\beta$ and weight decay $\xi$.

Figure 3 shows the results of an experiment with a larger model, WideResNet-16-1, on the CIFAR-10 dataset. Due to the size of the model, non-greedy HPO without hyperparameter sharing becomes too computationally expensive, so the authors only compared FDS with a greedy method, Hypergradient Descent (HD). This experiment shows that in just 10 outer steps, FDS manages to converge to noticeably more reasonable values of the hyperparameters than HD, resulting in better test performance while still being a viable option in this setting.

FDS vs. other HPO methods

Figure 4: Performance of the most popular HPO methods on CIFAR-10 for a WideResNet-16.

The last experiment aims to demonstrate why FDS is such a valuable method compared to others used before it. We can see that non-greedy methods like random search (RS), Bayesian optimization (BO), Hyperband (HB) and the combination of the latter two (BOHB), while solving for global optima, rely on trial-and-error, which makes them very slow. On the other hand, greedy methods like Hypergradient Descent (HD) are faster but solve for local optima. FDS manages to take the best of both worlds, outperforming even the next best method while converging 20 times faster.

Conclusion

FDS has been proven to be a well-balanced alternative to the state-of-the-art methods at the time, and its notable performance still makes it a strong baseline to this day. While the field of HPO has since then developed quite significantly, FDS is still worth considering in various applications, given its strengths.

Main strengths of FDS:

• Hypergradient noise reduction: hyperparameter sharing allows to combat gradient degradation, reducing the hypergradient error for optimal configurations of $W$ (for many purposes $W=50$ proved to be quite optimal).
• Accuracy: offers better accuracy without trade-offs in comparison to greedy methods.
• Convergence speed: converges much faster than trial-and-error methods.

Limitations:

• Requires differentiable hyperparameters: to use FDS with discrete hyperparameters, you will need to perform relaxation.
• Memory requirements scale linearly with the amount of hyperparameters. For example, a 12 GB GPU can train up to ~$10^3$ hyperparameters.
• Recurrent formulas are parameter-specific: each type of hyperparameter will require the derivation of its own expressions for matrices $\mathbf{A}_t$ and $\mathbf{B}_t$.

Hyperparameter Optimization (HPO) remains one of the most resource-intensive bottlenecks in the machine learning pipeline. While the community has largely moved past Grid Search, the standard alternatives — Random Search and Bayesian Optimization — still suffer from a fundamental inefficiency: they treat the training process as a “black box” that must be run to completion.

In 2018, Lisha Li and researchers from Carnegie Mellon, Google, and the University of Washington published the paper “Hyperband: A Novel Bandit-Based Approach to Hyperparameter Optimization.” They proposed a paradigm shift: instead of trying to intelligently select configurations (as Bayesian methods do), we should focus on efficiently evaluating them using adaptive resource allocation.

In this post, we will deconstruct the Hyperband algorithm, the theoretical problem it solves, and why it often outperforms Bayesian methods.

The Core Problem: The “n vs. B/n” Tradeoff

The central challenge in random search-based HPO is resource allocation. Suppose you have a total finite budget $B$ (e.g., total GPU hours). You need to decide how many unique hyperparameter configurations $n$ to evaluate.

This creates a fundamental tradeoff:

1. Maximize $n$ (Width): You sample many configurations to cover the search space, but each gets a very small average budget ($B/n$).
   • Risk: You might stop a promising configuration too early (“false negative”).
2. Minimize $n$ (Depth): You sample few configurations, but train them to convergence.
   • Risk: You train a poor configuration for too long, wasting resources that could have been used to explore other areas of the search space.

For a fixed budget, it is impossible to know which strategy — width or depth — will yield the best model. Hyperband was designed specifically to solve this dilemma.

The Building Block: Successive Halving

To understand Hyperband, one must first understand its subroutine: Successive Halving (SH). Originally proposed for multi-armed bandit problems, SH operates like a tournament:

1. Initialize: Start with $n$ randomly sampled configurations.
2. Evaluate: Allocate a small budget $r$ (e.g., 1 epoch) to all configurations.
3. Select: Rank them by validation loss and discard the worst half.
4. Promote: The surviving configurations are promoted to the next round with a larger budget.
5. Repeat: Continue until one configuration remains.

While SH is efficient, it still requires the user to choose $n$. Given some finite budget $B$, if $n$ is too large, the initial budget $r$ might be too small to distinguish good models from bad ones. If $n$ is too small, SH behaves like standard Random Search.

The Hyperband Algorithm

Hyperband acts as a “wrapper” or an outer loop around Successive Halving. Instead of forcing the user to guess the optimal $n$, Hyperband iterates through different feasible values of $n$ for a fixed total budget.

It divides the total budget into several “brackets” (instances of Successive Halving):

• Most Aggressive Bracket ($s = s_{max}$): Starts with the maximum possible number of configurations ($n_{max}$) with the minimum resource per config. This is designed to identify “fast learners” quickly.
• Intermediate Brackets: Gradually decrease $n$ and increase the initial resource $r$.
• Most Conservative Bracket ($s = 0$): Starts with a small number of configurations but allocates the maximum resource immediately. This is essentially equivalent to standard Random Search (exploration) or simply training to convergence (exploitation).

Algorithm Inputs

Hyperband is notably easy to configure, requiring only two inputs:

1. $R$: The maximum amount of resource that can be allocated to a single configuration (e.g., 100 epochs, or the full dataset size).
2. $\eta$: The proportion of configurations discarded in each round of Successive Halving.

By iterating through these brackets, Hyperband performs a geometric search over the trade-off between “number of configurations” and “resource per configuration.”

Theoretical Framework: The Infinite-Armed Bandit

The authors frame HPO as a non-stochastic infinite-armed bandit problem.

• Infinite-armed: The hyperparameters are drawn from a continuous probability distribution.
• Non-stochastic: The algorithm does not make strong assumptions about the convergence curves of the loss functions.

This theoretical grounding is significant because it contrasts with Bayesian Optimization (BO). BO relies on fitting a probabilistic model to the function $f(x)$. In high-dimensional spaces, fitting this model becomes computationally expensive and often inaccurate. Hyperband avoids this complexity entirely by relying on principled random sampling and aggressive early stopping.

Empirical Results

The paper presents extensive evaluation comparing Hyperband against Random Search, SMAC, TPE, and Spearmint (popular Bayesian optimization frameworks) on several benchmarks.

1. Deep Learning (Iterations as Resource)

In this experiment, the authors tuned Convolutional Neural Networks (CNNs) on datasets like CIFAR-10 and SVHN. The resource budget was defined as the number of training iterations (epochs).

Average test error across 10 trials. Label “SMAC (early)” corresponds to SMAC with the early-stopping criterion and label “bracket s = 4” corresponds to repeating the most exploratory bracket of Hyperband. Result: Hyperband found high-quality configurations 5× to 30× faster than Bayesian methods.

2. Kernel Methods (Data Subsampling as Resource)

Here, the task was Kernel Least Squares classification. The resource was the size of the dataset subsample.

On left: Average test error of the best kernel regularized least square classification model found by each searcher on CIFAR-10. On right: Average test error of the best random features model. Result: Hyperband achieved a massive 70× speedup over Random Search.

3. Generalization (117 OpenML Datasets)

To test robustness, Hyperband was evaluated on a large-scale automated machine learning task involving 117 real-world datasets from OpenML.

Average rank across all data sets for each searcher. For each data set, the searchers are ranked according to the average validation/test error across 20 trials.

• Avoiding Overfitting: A key finding was that Bayesian optimization methods often “overfit” the validation set — they found configurations that looked good during search but performed worse on the test set. Hyperband, being closer to Random Search in its sampling strategy, showed better generalization.
• Cost vs. Benefit: On a subset of 21 datasets where subsampling yielded meaningful computational speedups, Hyperband was the clear winner. However, on very small datasets where training takes seconds, the overhead of Hyperband made it less effective than simple Random Search.

Conclusion

The Hyperband paper provides a compelling argument that in the era of expensive model training, adaptive resource allocation is more critical than adaptive configuration selection.

Key Practical Takeaways:

1. Efficiency: For problems where partial training (e.g., few epochs) correlates with final performance, Hyperband is superior to standard Random Search and often beats Bayesian Optimization.
2. Simplicity: It requires minimal tuning compared to the complex kernels and acquisition functions of Gaussian Processes.
3. Parallelism: The algorithm is easily parallelizable, making it ideal for modern distributed computing clusters.

Today, Hyperband has become an industry standard, available in major HPO frameworks such as Ray Tune, Optuna, and Scikit-learn (HalvingRandomSearchCV). For anyone dealing with computationally expensive model tuning, it is a valuable tool in the machine learning toolkit.

You may know that image classification models are generally invariant to visual details such as brightness, object pose, or background configuration. For example, if we are classifying animals, models don’t care whether it’s a white poodle or a black puppy - both will be recognized simply as a dog. In other words, classification models discard most information about the image other than what is necessary to predict its class.

In this paper, the authors try to reconstruct images from the logits of a classification model. Earlier work has already attempted to invert logits back into images, but in this article they propose a more effective way to generate images using conditional BigGAN (more about the generation process below).

How can we invert logits to image?

Methods developed to invert representations in classification networks generally fall into two categories: optimization based and learning based.

Optimization based

This method just uses backpropagation to find the image $x \in \mathbf{R}^{H \times W \times C}$ that minimizes the loss:

where $x_0$ is the original image, $\mathcal{R}$ is a prior function and $\lambda$ is a hyperparameter.

Optimization-based methods offer a training-free, model-agnostic way to probe and interpret learned representations.

But the drawbacks are clear: the generated image strongly depends on the random initialization, and the method is slow and computationally expensive.

Learning based

Learning-based methods use a training set of {logits, image} pairs to learn a decoder network that maps a logit vector to an image. In this work, the decoder is implemented as a generator, which produces an image from raw logits and some random noise. This is the approach used in the paper.

Experiment

Now let’s look at the scheme of the experiments conducted by the authors.

Here we have 3 model: a classifier $\Phi$, a generator $G$ and a discriminator $D$.

• Classifier

It’s just a pre-trained image classifier (in this paper, ResNet and Inception are used), but instead of a final label it outputs raw logits $z = \Phi(x)$.

• Generator

The generator takes raw logits $z$ and some random noise $\epsilon$, and produces an image $\hat{x} = G(z, \epsilon)$.

• Discriminator

Discriminator takes an image $\bar{x}$ (a real image $\bar{x} = x$ or an image $\bar{x} = \hat{x}$ generated by generator $G$) and outputs a real value $d = D(\bar{x}, z)$. If the image $\bar{x}$ is real then $d \gg 0$. If the image is synthetic then $d \ll 0$.

How It Works

Let’s understand how this setup works and why it answers the main question of the article.

The classifier $\Phi$ is fixed, so we only train the generator $G$ and discriminator $D$.

The discriminator trains to minimize:

Basically it means that we teach discrimantor to output positive value if it gets a real image and negative value if it gets a generated one.

The generator trains to minimize:

In other words, the generator learns to create such images that the discriminator will classify as real, essentially trying to fool it.

Thus, after training, we obtain a generator that can produce a realistic (or at least similar to real) image purely from raw logits.

Results

Example of reconstructions

Let’s look what the authors got from our trained generator.

Here are images generated from logits for different pre-trained classifiers. You can see that reconstructed images look surprisingly similar to the original ones.

Notice that the authors used different classifiers, Inception, and the robust and non-robust versions of ResNet.

ResNet reconstructions stay closer to the original image, while Inception outputs look more realistic.

Also the authors compare their architcture to the method of Dosovitskiy and Brox (2016), which fails in recovering object shape and details. For example, you can see that the rabbit here is just a fuzzy blur.

Visualizing Model Invariances

Besides reconstructing images, the authors also use the generator to visualize which properties of the image the classifier is invariant to.

The generator takes as input Gaussian noise in addition to the logit vector from classifier.

The noise affects non-semantic properties - shape, pose, size, position, showing which aspects are not encoded in the logits themselves.

Noise resampling for Robust ResNet. The top left images are the original ones.

Noise resampling for non-robust ResNet. The top left image are the original ones.

Notice how noise has a greater effect on the non-robust model than on the robust model.

Reconstruction of incorrectly classified images

What happens if we reconstruct an image from logits that corresponds to an incorrect prediction? Surprisingly, the reconstructed image still looks like the original:

Logit manipulations

What happens if we manipulate the logits themselves? The authors perform three types of modifications:

• logit shifting - adding a constant to each logit
• logit scaling - multiplying logits by a constant
• logit perturbation - adding Gaussian noise

In (a), for the robust model, shifting mainly affects contrast and sharpness, but also subtly changes shape. For example, in the hockey scene, three players gradually merge into one with larger shifts. In the non-robust model, the effect is much weaker.

In (b), scaling logits in robust models changes sharpness and contrast. In non-robust models there are fewer brightness changes, but the content itself starts to shift (e.g., the coral reef changes shape).

In (c), perturbing logits with Gaussian noise affects image content in both robust and non-robust models. For robust models, content changes are moderate; for non-robust models, noise changes the image much more dramatically, suggesting that their logits are more closely clustered in the output space.

Conclusion

We discussed that one can reconstruct remarkably accurate images from logits hat often look very close to the originals. Earlier we believed that classifiers discard irrelevant information. But in reality, the final logits contain more than just class-related features.

Even when the classifier gives incorrect prediction, we still manage to reconstruct an image similar to the original!

Also we discovered how logit manipulations affect reconstructed images. For robust ResNet-152, logit shifts and rescaling influence both contrast, sharpness and brightness, while for the non-robust model these manipulations have much stronger effects on the image content, highlighting how robustness affects the stability of reconstructed images.

Surprisingly, logits preserve more than just class information; they retain enough detail to reconstruct the original image, even though networks are expected to be invariant to differences among instances of a class.

If you find this topic interesting, please check out the original paper!

Motivation

In multi-class classification problems, the standard loss function used by the overwhelming majority of machine learning models is the log-softmax function. It is certainly a convenient function with interpretable values and works well — but is it necessarily the best choice of loss function, period? The answer may be: not always. In this blog post, we’ll dive into research conducted by Alexandre de Brébisson and Pascal Vincent on loss functions for multi-class classification tasks and explore some alternative methods from the family of spherical loss functions.

But let’s take this step by step.

The Multi-Class Classification Task

First, we’ll establish a mathematical model of a loss function. Given a neural classification model, let’s denote the output of its last hidden layer as $\mathbf{o}$, where $\mathbf{o}$ is a $d$-dimensional vector. Suppose we have $D$ classes in our classification task. Let us denote the target vector $y$ and its non-zero component’s index as $c$. Then, a loss function is a function of the last layer output and $c$:

Current Approach: Log-Softmax

The log-softmax loss function is defined in our notation as follows:

This loss function is standard in multi-class classification, but it may be a bit too computationally expensive in some cases: gradient updates require $\mathcal{O}(D\times d)$ calculations, which for large output dimensionality $D$ (i.e., in tasks with large numbers of classes, such as language modeling) may be suboptimal. Alternative loss functions, which we’ll discuss below, may provide a solution to this problem.

The Spherical Loss Family

Let’s introduce a new family of loss functions depending on symmetric statistics of the hidden layer output vector $\mathbf{o}$. A loss belongs to the spherical family if it depends only on:

• $s = \sum_i o_i$,
• $q = \sum_i o_i^2 = |\boldsymbol{o}|^2_2$,
• $o_c$,
• $y_c$ for the target class $c$.

Then such a function can be written in the form

This definition may sound restrictive, but this family of functions is in fact quite diverse. One simple example of a spherical loss function—though for regression rather than classification—is MSE (mean squared error).

A neat property of spherical loss functions, discovered and proven by the paper’s authors together with Guillaume Bouchard in their 2015 paper, is that they allow for efficient gradient computation in $\mathcal{O}(d^2)$ instead of $\mathcal{O}(D\times d)$, which for large numbers of classes $D$ is noticeably better.

Spherical Softmax

One simple example of a spherical loss function is the spherical softmax loss. It is defined as follows:

where $\varepsilon$ is a small additive term for numerical stability in case $q$ is very small. It is fairly trivial to prove that this function is spherical, but it has more useful properties than just that:

• $L_{\text{log sph soft}}$ is invariant to scaling of $\boldsymbol{o}$.
• It is an even function, i.e., it ignores the sign of $o_k$.

However, it should be kept in mind that this loss function also requires careful tuning of the hyperparameter $\varepsilon$ for numerical stability.

Taylor Softmax

Another loss function from the spherical loss family is based on the second-order Taylor decomposition of the exponent in the regular softmax loss: $\exp(x) \approx 1 + x + \frac{1}{2}x^2$. It is defined as follows:

It isn’t hard to see that this is also a spherical loss function that depends only on $s$, $q$, and $o_c$. Moreover, this function, in contrast to the previous one, doesn’t require any hyperparameters and is numerically stable. Another distinctive feature of this function is that it is slightly asymmetric around zero. Interestingly, the paper’s authors hypothesize that this is actually a positive feature, which we’ll see in action later on.

Spherical Upper Bound for Log-Softmax

We can’t go on without mentioning one more loss function discussed in the paper, which is a spherical upper bound of the log-softmax function. It is derived from an upper bound for the log-sum of exponentials proposed by Bouchard (2007) and has the following monstrosity of an expression:

You can guess that I didn’t type that one out by hand. Ironically, the most complex of the functions discussed in the paper performed the worst. So poorly, in fact, that the authors mentioned it briefly at the beginning of the experiments section and then withdrew it from consideration. In our discussion of experimental results, we will follow their example.

Experimental Results

To test the loss functions in action, the authors compared log-softmax and different spherical alternatives on several tasks: with low-dimensional outputs, like image classification on MNIST and CIFAR-10, and with higher-dimensional outputs, like classification on CIFAR-100 and a language modeling task on the PennTree dataset. The goal was not to reach state-of-the-art performance on each task but to compare the influence of each loss given the same classification model.

Low-Dimensional Outputs

Experiments on datasets with low numbers of classes showed that such tasks may indeed be a scenario where spherical loss functions are superior to the usual log-softmax. Particularly, Taylor softmax noticeably outperforms log-softmax on both tasks with low-dimensional outputs, and spherical softmax also achieves comparable results.

Higher-Dimensional Outputs

On the contrary, log-softmax performs better as output dimension increases. This is particularly prominent in language modeling tasks with huge vocabulary sizes: here the difference between log-softmax and spherical losses is really noticeable. The paper’s authors suggest that this may be caused by the exponential in softmax better handling high-dimensional competition. It should be noted that such tasks with higher-dimensional outputs are the ones where the gain in computational efficiency achievable by spherical losses is most visible.

Conclusion

So, what did we end up with? Spherical losses (especially Taylor softmax) can outperform log-softmax on small-output tasks (e.g., MNIST, CIFAR-10), but for high-dimensional outputs (CIFAR-100, language modeling), log-softmax remains superior. The paper we discussed thus highlights that alternatives to log-softmax are in fact worth considering, but the choice of loss function should be task-specific.

In tasks with high-dimensional outputs, spherical losses enable efficient training but may lack the discriminative power of log-softmax. They are worth considering for specific applications where efficiency is key. In tasks with low-dimensional outputs, on the other hand, the use of spherical loss functions, especially Taylor softmax, can benefit model accuracy — so in such tasks, we may have just discovered the softmax killer.