Variational Deep Learning via Implicit Regularization

This paper proposes a method for regularizing variational neural networks by leveraging the implicit bias of stochastic gradient descent rather than explicit priors, theoretically characterizing this bias in overparametrized linear models and demonstrating empirically that it achieves strong in- and out-of-distribution performance with minimal computational overhead and no additional hyperparameter tuning.

The Gist

Imagine you are teaching a student (a computer program) to recognize pictures of cats and dogs.

The Problem: The Overconfident Student
In the world of modern AI, we have built students who are incredibly smart at memorizing their homework. If you show them 1,000 pictures of cats, they can identify them perfectly. However, these students have a flaw: they are overconfident.

If you show them a picture of a toaster, they might say, "I am 99.9% sure this is a cat!" because they haven't learned how to say, "I don't know." They are also brittle; if you slightly blur the image or change the lighting (like a real-world surprise), they often fail miserably.

The Old Solution: The Strict Teacher
For years, the standard way to fix this was to hire a "Strict Teacher" (called a Prior in math terms). This teacher would constantly interrupt the student's learning process, saying, "Wait, don't just memorize that! Remember, cats usually have fur, so if you see something smooth, be less sure."

While this works, it's expensive. The teacher takes up a lot of space in the classroom (computational memory) and requires a lot of time to explain their rules. Also, if the teacher's rules are slightly wrong, the student gets confused and performs worse.

The New Idea: The "Implicit" Coach
This paper proposes a brilliant new strategy: Stop hiring the Strict Teacher. Instead, rely on the Implicit Coach that is already there.

Here is the analogy:
Imagine the student is running a race to find the finish line (the perfect answer).

• Standard Training: The student runs as fast as possible. If there are many paths to the finish line, they just pick the first one they see.
• The Paper's Insight: The paper realizes that the way the student runs (the specific steps they take, the momentum they build, and how they start the race) naturally guides them to a specific, safer finish line. They don't need a teacher to tell them where to go; the physics of the run itself keeps them on a good path.

How It Works (The "Variational" Twist)
Usually, to make a computer "uncertain" (so it knows when it's guessing), we force it to learn a whole cloud of possible answers instead of just one.

• The Old Way: We tell the student, "Here is a cloud of answers. But you must stay close to the teacher's map." (This is the expensive part).
• The New Way (IBVI): The authors say, "Just run the race normally. Don't worry about the teacher's map. Just trust that the way you run (the optimization process) will naturally keep your 'cloud of answers' in a safe, sensible shape."

They proved mathematically that if you start the student in the right place and let them run with their natural momentum, they will naturally settle into a state where they are:

1. Confident when they know the answer (on training data).
2. Hesitant when they are unsure (on new, weird data).

The "Magic" Ingredient: The Starting Line
The paper also discovered that where you start the race matters more than you think.

• Standard Start: If you start a small student and a giant student at the same spot, they run differently. You have to re-teach the giant student how to run.
• The Paper's Fix (µP): They found a special "Starting Line" (called Maximal Update Parametrization) where a small student and a giant student run in the exact same way. This means you can train a tiny model, find the perfect speed (learning rate), and then just "copy-paste" that speed to a massive model without re-tuning anything. It's like finding the perfect gear ratio for a bicycle and knowing it works for a child's bike and a mountain bike alike.

The Results

• Cheaper: It uses almost the same computer power as a standard model (no expensive "Strict Teacher" needed).
• Smarter: The models are much better at saying "I don't know" when they see a blurry or weird image.
• Faster to Scale: Because of the new "Starting Line," it's much easier to train huge models without wasting time tuning them.

In a Nutshell
This paper says: "Stop micromanaging your AI with expensive rules. Trust the natural rhythm of the learning process. If you set the stage correctly, the AI will naturally learn to be humble, accurate, and robust, all while saving you money and time."

Technical Summary

1. Problem Statement

Modern deep learning models achieve remarkable in-distribution generalization despite being overparameterized and trained with little to no explicit regularization. This phenomenon is attributed to the implicit regularization (or inductive bias) imposed by the optimization algorithm (e.g., Stochastic Gradient Descent), architecture, and initialization.

However, standard deep neural networks (DNNs) often fail to generalize well to out-of-distribution (OOD) data and produce overconfident predictions. Bayesian Deep Learning (BDL) addresses this by learning a distribution over weights to provide uncertainty quantification and robustness. Traditional BDL methods (e.g., Variational Inference) rely on explicit regularization via a Kullback-Leibler (KL) divergence to a prior. This approach introduces significant challenges:

• Computational Overhead: Requires sampling and calculating complex regularizers.
• Prior Elicitation: Requires carefully chosen priors, which can override the beneficial implicit biases of optimization.
• Pathological Biases: Explicit priors combined with approximate inference can lead to poor uncertainty estimates (e.g., uncertainty collapse).

The core problem is how to achieve the robustness and uncertainty quantification of Bayesian methods without the computational cost and potential negative side effects of explicit regularization, by instead leveraging the implicit bias of optimization.

2. Methodology: Implicit Bias Variational Inference (IBVI)

The authors propose Implicit Bias Variational Inference (IBVI), a framework that trains variational neural networks by minimizing the expected loss over the variational distribution, omitting the explicit KL divergence term to the prior.

Core Objective

Instead of minimizing the Evidence Lower Bound (ELBO) which includes a KL term:
$$\mathcal{L}_{ELBO} = \mathbb{E}_{q_\theta}[\ell(y, f_w(X))] + \lambda \text{KL}(q_\theta \| p)$$

IBVI minimizes only the expected loss:
$$\theta^* \in \arg\min_\theta \mathbb{E}_{q_\theta(w)}[\ell(y, f_w(X))]$$

Theoretical Insight: Generalized Variational Inference

The paper rigorously characterizes the implicit bias of (Stochastic) Gradient Descent (SGD) when training this objective in overparameterized linear models.

• Key Finding: If the variational distribution is initialized at the prior $p(w)$, SGD converges to the distribution $q_{\theta^*}$ that minimizes the expected loss (zero training error) while being closest to the prior in 2-Wasserstein distance ($W_2$).
• Mathematical Formulation:
  $$q_{\theta^*} = \arg\min_{q_\theta} \{ W_2^2(q_\theta, p) \mid \theta \in \arg\min \bar{\ell}(\theta) \}$$
• Implication: The optimization process implicitly acts as a Generalized Variational Inference procedure where the regularizer is the 2-Wasserstein distance rather than the KL divergence. This preserves the inductive bias of the optimizer while naturally weighting models based on their proximity to the initialization (prior).

Practical Implementation Details

• Training with $M=1$: To minimize computational overhead, the authors demonstrate that training with a single parameter sample ($M=1$) per batch is sufficient, provided the learning rate is adjusted (reduced) to account for the increased noise in the gradient.
• Parametrization (µP): The authors extend Maximal Update Parametrization (µP) to the variational setting. This ensures that feature learning occurs even as network width increases and allows for hyperparameter transfer (tuning learning rates on small models and applying them to larger ones without re-tuning).
• Architecture: They use Gaussian variational distributions with low-rank covariance matrices for the input and output layers (and optionally hidden layers), keeping the parameter overhead minimal (approx. 10% increase over standard DNNs).

3. Key Contributions

1. Theoretical Characterization: The paper provides rigorous proofs (Theorems 1 and 2) showing that for overparameterized linear models, minimizing the expected loss via SGD is equivalent to Generalized Variational Inference with a 2-Wasserstein regularizer. This bridges the gap between standard deep learning optimization and Bayesian inference.
2. Implicit Bias VI (IBVI): A new training algorithm that removes the explicit KL term, relying solely on the implicit bias of SGD initialized at the prior.
3. Parametrization for Probabilistic Networks: An extension of µP to variational networks, enabling stable training across different widths and facilitating hyperparameter transfer.
4. Efficiency: The method achieves Bayesian-like performance with minimal computational overhead (comparable to standard DNNs) by using single-sample training and low-rank covariance structures.

4. Experimental Results

The authors benchmarked IBVI against standard DNNs, Temperature Scaling (TS), Laplace Approximation (LA), Weight-Space VI (WSVI), SWAG, and Deep Ensembles (DE) on image classification datasets (MNIST, CIFAR-10/100, TinyImageNet) and robustness benchmarks (corrupted versions of these datasets).

• In-Distribution Performance: IBVI achieves test error rates comparable to Deep Ensembles and standard DNNs, significantly outperforming standard Weight-Space VI and Laplace Approximation in many cases.
• Uncertainty Quantification: IBVI provides competitive Negative Log-Likelihood (NLL) and Expected Calibration Error (ECE), often outperforming Deep Ensembles in calibration while using far less memory.
• Out-of-Distribution (Robustness): On corrupted datasets (e.g., noise, blur, pixelation), IBVI demonstrates superior robustness compared to standard DNNs and most other Bayesian baselines, with performance approaching that of Deep Ensembles but with a fraction of the memory cost.
• Hyperparameter Transfer: Experiments confirm that using µP allows learning rates tuned on small networks to transfer effectively to larger networks, a property not observed with standard parametrization.

5. Significance and Impact

This work fundamentally shifts the perspective on Bayesian Deep Learning:

• Paradigm Shift: It suggests that explicit regularization (KL divergence) may be unnecessary and potentially detrimental. The "magic" of generalization in deep learning comes from the optimizer, and this bias can be harnessed for probabilistic modeling without explicit priors.
• Scalability: By eliminating the need for explicit prior calculations and complex sampling schemes, IBVI makes Bayesian deep learning scalable to large models and datasets, closing the gap between the theoretical promise of BDL and practical deep learning efficiency.
• Theoretical Foundation: The connection between SGD's implicit bias and 2-Wasserstein distance provides a new theoretical lens for understanding why overparameterized models generalize well and how to extend this to probabilistic settings.

In summary, the paper demonstrates that implicit regularization is sufficient to train robust, uncertainty-aware neural networks, offering a computationally efficient alternative to traditional Variational Inference that aligns closely with standard deep learning practices.