Understanding and inverse design of implicit bias in stochastic learning: a geometric perspective

A key challenge in machine learning is explaining how learning dynamics select among many solutions that achieve identical loss values in overparameterized models—a phenomenon known as implicit bias. This paper develops a theoretical and constructive framework where implicit bias emerges as a geometric correction induced by the interplay between gradient noise and continuous symmetries of the loss, offering a unified mechanism for understanding and designing learned representations.

The Gist

Implicit bias is characterized not at the level of a single optimization trajectory, but at the level of distribution over solutions induced by stochastic training dynamics.

Theoretical Framework for Implicit Bias

The framework reformulates implicit bias by considering the distribution over equivalence classes of parameters corresponding to the same predictor, which is generated when parameter-space statistics are mapped onto predictor space. The core result is that under stochastic learning dynamics, mapping parameter-space statistics onto predictor space induces a geometric correction that reshapes the effective learning dynamics. This leads to an effective loss function:

Leff(θ) = L(θ) + σ2 / 2β log det G(θ) = L(θ) + LIB(θ)

Stochastic Learning Dynamics and Symmetry Breaking

The paper approximates Stochastic Gradient Descent (SGD) by overdamped Langevin dynamics with isotropic noise: dθt = −∇L(θt) dt + σ2 / 2β dWt. The formal stationary density in parameter space is proportional to exp − β σ2 L(θ) / dVolΘ(θ). When continuous symmetries exist in the parameter space, the orbit volumes can be infinite, necessitating a symmetry-breaking step. This is achieved by introducing a map χ: Θ → Rm such that Sχ:= χ−1(0) is a local slice intersecting each orbit locally at exactly one point. The induced stationary density on this slice is given by:

ρS(θ) ∝ exp − β σ2 L(θ) / (det Gχ(θ)) dσS

This leads directly to the effective loss:

Leff(θ) = L(θ) + σ2 / 2β log det Gχ(θ)

Inverse Design Principle

The framework establishes a constructive inverse-design principle: by engineering predictor-invariant parameterizations, one can induce targeted implicit biases directly at the level of predictors. This is demonstrated through two canonical constructions:

1. Hadamard factorization: By introducing a redundant Hadamard factorization z = u ⊙ v, the induced correction is LIB(u, v) = σ2 / 2β Σ log zi + const. Minimizing this over factorizations yielding the same invariant feature z leads to the balancedness condition ui = vi.

2. Matrix factorization: For a desired bias on features given by the logarithmic spectral penalty rank(X T)j=1 log σj (T), introducing a redundant factorization T = UV ⊤ induces a correction LIB(U, V) = σ2 / 2β rank(X T)j=1 log σj (T) + const. Minimization yields the balancedness condition λU,j = λV,j = q σ⋆ j.

Empirical Validation

The theory is validated through controlled experiments across various architectures:

1. Shallow ReLU networks: The theory predicts that stochastic dynamics select balanced representatives along neuron-wise positive rescaling orbits, leading to the condition∥W[j,:]∥2 → 1.

2. Single-head scaled dot-product attention (SDPA): The column-balancedness condition is predicted:∥Q[:,j]∥2 = ∥K[:,j]∥2 = q/r k C[:,j] squared.

3. Low-rank matrix completion: The theory predicts that stochastic learning selects balanced factorizations along each mode, leading to the condition∥U[:,j]∥2 = 1 and∥V[:,j]∥2 = 1 for the factorized model Tˆ=UV⊤.

The paper concludes that implicit bias emerges as a geometric effect of noise and symmetry breaking, providing a unified and constructive account that connects machine learning, stochastic dynamics, and geometric methods from statistical physics. This framework allows for the design of over-parameterized models with desired priors by controlling the geometry of the redundancy through which predictors are represented.

Data Availability

All data supporting the findings can be algorithmically generated from the code provided. safetensors files required to programmatically re-generate the pictures without re-running the experiments are also provided as part of the code.

Code Availability

Code to fully reproduce the experiments supporting the findings contained in this paper, and to re-generate the pictures shown above, can be acquired via github.com/emaballarin/understanding-design-ib.

Acknowledgments

We thank Liu Ziyin and Tomaso Poggio for useful discussions.

Here are the specific improvements that can be made to AI systems, based directly on the theoretical framework presented in this paper, and what those improved systems will be capable of doing:

The core capability derived from this paper is the ability to shift implicit bias from an uninterpretable empirical observation into a controllable, inverse-designed engineering principle.

Here are the specific improvements and their resulting capabilities:

Improved Model Architecture Design via Targeted Bias Engineering

By utilizing the constructive inverse-design principle (Equations 13–18), researchers can engineer redundant parameterizations (like Hadamard factorizations or matrix factorizations) that explicitly impose desired inductive biases onto the learning dynamics.

Specific Capability: Precision Sparsity Control

The system can be designed to enforce coordinate sparsity in learned representations, not just via traditional penalties, but by ensuring the optimization process selects balanced representatives along specific symmetry orbits. This allows for creating models that are inherently sparse in a mathematically rigorous way (as demonstrated by the Hadamard factorization example).

Specific Capability: Structured Low-Rank Discovery

The system can be engineered to favor low-rank solutions in complex data structures (like matrices or attention heads) by utilizing matrix factorization parameterizations. This moves beyond standard methods to ensure that the learned structure adheres to a specific spectral penalty, leading to more robust and generalizable low-rank approximations of high-dimensional tensors.

Specific Capability: Balanced Feature/Weight Optimization

The system can be designed to enforce balance conditions (e.g., ensuring the norm of incoming weights equals the magnitude of outgoing coefficients) during training in models like shallow ReLU networks or attention mechanisms. This leads to representations where feature contributions are uniformly weighted, improving stability and convergence properties.

Improved Interpretability through Geometric Correction Quantification

The framework provides a direct, computable expression for the geometric correction term:

Equation (1): Leff(θ) = L(θ) + σ2 / 2β log det G(θ) = L(θ) + LIB(θ).

This allows researchers to quantify exactly how much the implicit bias is deviating from the ideal loss landscape and how this deviation is driven by the geometry of the parameter space.

Specific Capability: Robustness to Over-parameterization

By understanding that implicit bias is a geometric consequence of noise propagating through parameter symmetries, systems can be designed to be inherently more robust across equivalent solutions. The framework predicts that stochastic learning selects balanced or canonical representatives on symmetry orbits, which are often the most stable points in the loss landscape.

Universal Bias Unification

The framework unifies previously disparate results (low-rank factorizations, sparsity penalties) under a single geometric mechanism: predictor-preserving continuous symmetries inducing a correction to the effective loss. This provides a common theoretical language for designing implicit priors across vastly different model classes (e.g., from linear networks to attention heads).

Can we design a model such that its stochastic training favours a desired class of solutions without enforcing an explicit penalty? Under suitable conditions, the interplay between symmetries of a model's weight parametrization and stochastic training favours particular solutions, inducing an implicit bias. Building on this mechanism, we develop a framework for inverse-designing such biases by constructing novel parametrizations and their associated symmetries. We show how holomorphic functions make this construction and calculation simple and explicit. Specifically, we introduce a new parametrization that biases learned weights toward the binary values-1,+1. Numerical experiments confirm the theoretical predictions. They also show that our parametrization reproduces the same preference induced by an explicitly regularized model without adding a penalty to the training loss.