Bidirectional Attention as a Mixture of Continuous Word Experts

TL;DR

This paper provides a statistical interpretation of bidirectional attention as a mixture of experts model, revealing its connection to CBOW and MoE, and demonstrates its advantages in handling heterogeneous data and in out-of-distribution generalization.

Contribution

It introduces a novel statistical perspective of bidirectional attention as a mixture of experts, explaining its effectiveness and extending its application to categorical tabular data.

Findings

Bidirectional attention is equivalent to a mixture of experts model.

Extending this model to tabular data improves out-of-distribution generalization.

Stronger assumptions are needed for linear word analogies in attention-based models.

Abstract

Bidirectional attention $\unicode{x2013}$ composed of self-attention with positional encodings and the masked language model (MLM) objective $\unicode{x2013}$ has emerged as a key component of modern large language models (LLMs). Despite its empirical success, few studies have examined its statistical underpinnings: What statistical model is bidirectional attention implicitly fitting? What sets it apart from its non-attention predecessors? We explore these questions in this paper. The key observation is that fitting a single-layer single-head bidirectional attention, upon reparameterization, is equivalent to fitting a continuous bag of words (CBOW) model with mixture-of-experts (MoE) weights. Further, bidirectional attention with multiple heads and multiple layers is equivalent to stacked MoEs and a mixture of MoEs, respectively. This statistical viewpoint reveals the distinct use of MoE in bidirectional attention, which aligns with its practical effectiveness in handling heterogeneous data. It also suggests an immediate extension to categorical tabular data, if we view each word location in a sentence as a tabular feature. Across empirical studies, we find that this extension outperforms existing tabular extensions of transformers in out-of-distribution (OOD) generalization. Finally, this statistical perspective of bidirectional attention enables us to theoretically characterize when linear word analogies are present in its word embeddings. These analyses show that bidirectional attention can require much stronger assumptions to exhibit linear word analogies than its non-attention predecessors.