Transformers have revolutionized machine learning by making attention a central mechanism for modeling interactions within a context. Despite the central role of attention, the theoretical capabilities of Transformers for representing contextual relations remain unclear. In this work, we address this question by developing a mathematical framework based on probability and optimal transport. We view a text as a distribution of its representations and attention as a probabilistic relation between them. This perspective reveals a connection between attention normalization and optimal transport: standard softmax normalization produces conditional relations, while Sinkhorn normalization produces joint relations with prescribed marginals. Thus, both mechanisms provide structured probabilistic relations from attention scores. Under mild conditions, we establish universal approximation results for both settings. We show that Transformer architectures with Sinkhorn normalization can approximate arbitrary contextual relations represented as joint probabilities, while standard softmax Transformers can approximate arbitrary contextual relations represented as conditional probabilities. These results provide a mathematical characterization of the expressive power of Transformers for contextual relations and show how the choice of normalization determines the probabilistic structure of the relations represented by attention.
Figures & tables
Figure 1: Contextual semantic graphs. (A) Single-text case: words in the same sequence are connected by weighted, possibly asymmetric relations. (B) Two-text case: words from different sequences are linked through cross-text relations. For visual clarity, edges with small weights are omitted.
Figure 2: Contextual Transformer Block. (A) Contextual Transformer Block architecture in the two-text setting. (B) Contextual Transformer Block in the original Transformer architecture Vaswani et al. (2017) two text setting. (C) Contextual Transformer Block in the original Transformer architecture Vaswani et al. (2017) one text setting.