Adaptive inference and function vectors in deep transformers
For researchers in deep learning theory, this work provides a theoretical framework explaining how transformers perform adaptive inference, though the results are demonstrated only on synthetic tasks.
The authors develop a theory of deep transformers as mean-field interacting systems that perform distributed inference, showing they can infer latent context variables at increasingly finer scales using 'function vectors'. They validate predictions on in-context regression tasks, demonstrating that depth and feedforward blocks enable richer in-context learning than previously known.
Transformers are widely used as a general-purpose substrate for learning complex correlations between a large collection of coupled variables, but their internal mechanisms have remained mysterious. We introduce a theory of a deep transformer as a mean-field interacting system that implements distributed inference, subject to constraints on communication, locality and depth. We show that such a system can exploit internal state representations ('function vectors') to infer a latent context variable at increasingly finer scales over its layers. In an in-context regression task, the theory predicts a non-trivial relationship between non-Gaussian, hierarchical structure in the latent context variable, and transformer depth. Predictions are tested using constrained linear attention transformers and demonstrate adaptive inference in deep architectures. Feedforward blocks and depth enable transformers to implement a much richer class of in-context learning algorithms than previously described.