Transformers perform probabilistic inference through expectation maximization
Neural expectation-maximization and meta-models
Machine Learning
Summary
This paper investigates how Transformer neural networks can be understood as performing a type of probabilistic inference on data with hidden structure. The authors explain that Transformers approximate a process called max-product Loopy Belief Propagation combined with Expectation-Maximization, typically used for reasoning about uncertain information in networks. Each layer of a Transformer is seen as an iteration that refines guesses about hidden factors and model parameters. They also introduce a new design idea called meta-models, where the model’s parameters change depending on each input, unifying several existing neural mechanisms under a common framework.
What this means in practice
- •For machine learning engineers: Create input-adaptive neural networks that dynamically adjust their internal structure and parameters for improved model expressiveness and efficiency.
- •For speech recognition developers: Develop speech models that better handle uncertainty and hidden dependencies by embedding probabilistic inference steps within Transformer architectures.
Authors
Zhuolin Fu
Abstract
We interpret Transformer networks as learned approximations to a hybrid probabilistic-inference procedure that combines max-product Loopy Belief Propagation (LBP) with Expectation--Maximization (EM). In this view, a Transformer performs inference in partially observed Markov Random Fields (MRFs) whose structure and parameters are not given a priori but are inferred from data. We show that self-attention implements a differentiable approximation to max-product LBP on an input-conditioned, dynamically structured MRF, with attention weights encoding learned structural affinities. We further establish an EM correspondence at the level of depth: each layer updates latent representations in an E-step-like forward pass, while the transition $θ^{(l)} \to θ^{(l+1)}$ across layers acts as an implicit M-step. Building on this interpretation, we introduce \emph{meta-models}, a design paradigm in which inputs generate instance-specific parameters, unifying mechanisms such as self-attention, mixture-of-experts, and gated units within a single framework. Empirical results support the proposed perspective.