Skip to main content
Back to timeline
arXivSource publication:

Each Transformer layer read as one EM iteration: attention does belief propagation, the layer-to-layer jump updates parameters

Lead

Transformer attention is interpreted as max-product belief propagation on a Markov random field that varies from instance to instance, while the parameter change between layers acts as the parameter-update step of EM, a view that groups self-attention, mixture-of-experts, gated units and selective state-space models under one meta-model principle in which the input generates the parameters.

Source-provided article image: Neural expectation-maximization and meta-models
Figure 1 · arXiv

Story

The internal computation of a Transformer is recast as a probabilistic inference procedure rather than a set of separate architectural parts. Earlier explanations each cover part of the behavior, for example mapping in-context learning to gradient-descent-like or Bayesian updates, or relating attention to message passing on graphs, but they do not place self-attention, depth, latent-variable refinement, graph-structure selection and layerwise parameter change inside one probabilistic picture. The correspondence rests on a derivation: the Bayesian network taken to underlie an encoder Transformer is moralized and reduced by variable elimination to a fully connected Markov random field, and with the assumption that higher-arity factors can be approximated by pairwise factors, a mapping follows between attention and max-product belief propagation.

Self-attention is written as a one-step approximation to max-product belief propagation on a dynamically structured Markov random field. Belief propagation was built for partially observed models with known structure and parameters, and EM for models with known structure but unknown parameters, whereas real data often have both structure and parameters unknown, so either algorithm alone does not work effectively. Expanding each log message at its maximizer by a second-order Taylor expansion writes the belief maximizer as a precision-weighted average of the individual message maximizers; under an isotropic negative Hessian this matches the query, key and value form of self-attention, where attention weights correspond to learnable structural affinities and value vectors to message maximizers.

Transformer depth is read as an EM iteration: the forward pass updates latent variables and the change between layer parameters acts as the parameter update. Earlier work that read Transformer depth as iterative probabilistic fitting was either tied to a specific modeling choice or was an expressivity result for a particular model class, while this analysis frames the algorithm as probabilistic latent-variable inference. In a supervised setting the forward pass of layer l computes the latent variables, corresponding to the E step, but runs only one iteration of approximate max-product belief propagation to get an estimated maximizer; the next layer's parameter is the one that maximizes the objective, corresponding to the M step, and the Gibbs inequality relates improvement in that objective to improvement in the log likelihood.

The EM view motivates meta-models, a design paradigm in which the input first generates the parameters used by that forward pass. Hypernetworks, mixture-of-experts, gated convolutional and recurrent models, and selective state-space models are usually studied as separate design families, whereas here they are treated as instances or close relatives of one principle. Self-attention and mixture-of-experts can both be written as a meta-forward pass that determines parameters from the input followed by a forward pass; on two-dimensional synthetic data generated from Gaussian mixture models, meta-linear layers consistently outperform a standard linear layer across multi-modal datasets, with the SwiGLU variant reaching 0.6842 evaluation accuracy at the highest complexity while the vanilla linear layer degrades to 0.5171.

What to watch

The interpretation addresses encoder Transformers, whose underlying model is taken to be a partially observed Markov random field, and it relies on two assumptions: that any factor of arity greater than two can be reasonably approximated by a product of a finite number of pairwise factors, and that encoder Transformers implicitly make inferences on Markov random fields with pairwise factors that vary from instance to instance. In the experiments, output feature vectors of each token are taken from the last layer of several Huggingface Transformers on movie reviews and real-life texts from X, non-semantic or special tokens such as punctuation and newlines are removed, a cosine similarity matrix is computed, and single-linkage agglomerative hierarchical clustering is run and plotted as dendrograms. The text proposes future tests: comparing attention maps with known dependency graphs in synthetic graphical-model tasks, tracking whether layer-wise representations and parameters improve a surrogate objective across depth, and building architectures that generate parameters with controlled order and sparsity.

The text presents the interpretation as a perspective and states that empirical findings are strongly supportive of its claims; the two assumptions about pairwise-factor approximation and instance-varying Markov random fields are premises of the derivation. A meta-linear layer's parameter count is a multiple of a linear layer's, and the parameter count of a generalized meta-linear layer grows exponentially with its order. In the experiments the meta-linear layer reaches 0.6842 evaluation accuracy at the highest complexity while the vanilla linear layer reaches 0.5171, and a ground-truth MoE that sets the number of experts equal to the number of sub-models reaches 0.9400, 0.9022 and 0.8079, which the text treats as the upper bound on classification accuracy in these experiments.

Sources