Gradient Descent as Implicit EM in Distance-Based Neural Models

Alan Oursland

cs.LGDec 31, 2025

Gradient Descent as Implicit EM in Distance-Based Neural Models

Authors: Alan Oursland

Abstract

Neural networks trained with standard objectives exhibit behaviors characteristic of probabilistic inference: soft clustering, prototype specialization, and Bayesian uncertainty tracking. These phenomena appear across architectures -- in attention mechanisms, classification heads, and energy-based models -- yet existing explanations often rely on loose analogies to mixture models or post-hoc architectural interpretation. We provide a direct explanation. For any objective with log-sum-exp structure over distances or energies, the gradient with respect to each distance is exactly the negative posterior responsibility of the corresponding component: $\partial L / \partial d_j = -r_j$ . The identity is algebraic, requiring only differentiability; it is a specialization of Fisher's identity, and its significance here is its address: standard neural objectives instantiate it without modification. The consequence is that gradient descent on such objectives performs generalized expectation-maximization implicitly, with responsibilities arising as gradients to be applied rather than auxiliary variables to be computed. This result unifies three regimes of learning: unsupervised mixture modeling, where responsibilities are fully latent; attention, where responsibilities are conditioned on queries; and cross-entropy classification, where supervision clamps responsibilities to targets. Our claims live at training time: the responsibility-weighted gradient dynamics recently documented in transformers follow from the objective's geometry. The in-context Bayesian computation that trained transformers perform at inference time is the endpoint of these dynamics, not their per-step content.

Explore similar work

Jun 12, 2026cs.LG

EM-NeSy: Expectation Maximization for Neurosymbolic Learning

Neurosymbolic (NeSy) models integrate neural networks and symbolic reasoning for robust and interpretable AI. State-of-the-art NeSy models require that the symbolic component is expressed in a differentiable way, often complicating the use of approximate inference. We propose EM-NeSy which casts probabilistic NeSy learning as an instance of the Expectation-Maximization (EM) algorithm. In the expectation step, we compute the posterior over the neurally predicted symbols conditioned on the label via probabilistic inference. In the maximization step, we update the neural parameters based on this posterior using gradient descent only through the neural component. This formulation unlocks the full potential of the EM algorithm for NeSy learning. It allows NeSy to extend naturally to approximate reasoning without any additional modifications or differentiability requirements of the symbolic component. Furthermore, it recovers the standard end-to-end gradient-based NeSy setting under exact inference. Our experimental results demonstrate the scalability and computational efficiency of EM-NeSy.

Annegret Seibt, Luc De Raedt, Giuseppe Marra

Neuro-Symbolic Framework Expectation-Maximization Probabilistic Inference Differentiable Programming

Ask PDF

May 8, 2026cs.LG

Direct Bethe Free Energy Minimization for Bayesian Neural Networks

Bayesian neural networks are typically trained against the evidence lower bound (ELBO), whose Jensen gap closes only when the variational posterior is exact. We instead train by local consistency: gradient descent on the Bethe free energy, driving the belief at every factor toward agreement with its neighbours rather than placing a loss on the output. The resulting objective scores each observation by its own predictive density: a strictly proper rule whose optimum is the true conditional, for any likelihood with a tractable predictive convolution. Instantiated with a Gaussian last layer over a deterministic backbone, exact inference appears as one known corner: the neural-linear marginal likelihood. That corner is evidence-optimal; the shared-cavity, free-routed interior is predictive-optimal, improving NLL and calibration over it. This instance, SCROLL (Shared-Cavity fRee-rOuting Last-Layer), is a single-pass Bayesian neural network: batchable, any-likelihood, and implicitly empirical-Bayes-prior precision, observation noise, covariance, and backbone fit in one gradient pass. Prior work enters the interior only through the ELBO and its Jensen gap, even in this conjugate setting. At a single training run and forward pass per architecture-where the validation-tuned conventional references cross-validate

λ

and ensembles pay

5

50\times

at inference-a fixed SCROLL variant is best-or-tied on NLL and calibration on 7/8 UCI regression benchmarks, and best on 4/5 across three large tabular datasets (up to 515k examples) and two frozen text/vision embeddings.

Pavel Prochazka

Bayesian Neural Networks Bayesian Free Energy Principle Prior Knowledge Lower Bounds

Ask PDF

May 23, 2026stat.ML

How Neural Reward Models Learn Features for Policy Optimization: A Single-Index Analysis

Reward modeling is not only a prediction problem: in KL-regularized policy optimization, the learned reward is exponentiated to define the deployed policy, so downstream value depends on errors in reward-tilted regions. We study this feedback in a Gaussian single-index model with

r^*(x) = σ^*(\langle θ^*, x\rangle)

and

x \sim N(0, I_d)

. We analyze a two-stage neural reward model that first learns the hidden direction

θ^*

from reward-weighted samples and then fits the readout layer by weighted ridge regression. Exponential reward weighting changes the Hermite signal available to the first layer; for any feature-learning temperature

β_1

above a dimension-free

O(1)

threshold, a constant fraction of neurons recover the hidden direction, with weak-recovery complexity governed by the generative exponent. After feature recovery, we derive tilted-policy value-gap bounds for an idealized label-weighted fit with weights

e^{y/β_2}

and a more practical surrogate-weighted fit with weights

e^{r_{a_0}(x)/β_2}

. Keeping the

β_2

-dependence explicit yields an admissible set of deployment temperatures, balancing the gain from lowering

β_2

against the learning cost amplified by exponential weighting; in the surrogate-weighted case, proxy-dependent factors shrink this admissible set.

Rei Higuchi, Ryotaro Kawata, Akifumi Wachi +3

Topic pending

Ask PDF