ESPFormer: Doubly-Stochastic Attention with Expected Sliced Transport Plans

TL;DR

ESPFormer introduces a fully parallelizable doubly-stochastic attention mechanism using sliced optimal transport, improving efficiency and performance in various deep learning tasks without iterative normalization.

Contribution

The paper proposes a novel attention method based on sliced optimal transport that enforces doubly-stochasticity efficiently without iterative Sinkhorn normalization.

Findings

Improved attention regularization enhances model performance.

Consistent gains across image, text, and point cloud tasks.

Efficient, differentiable implementation suitable for deep learning.

Abstract

While self-attention has been instrumental in the success of Transformers, it can lead to over-concentration on a few tokens during training, resulting in suboptimal information flow. Enforcing doubly-stochastic constraints in attention matrices has been shown to improve structure and balance in attention distributions. However, existing methods rely on iterative Sinkhorn normalization, which is computationally costly. In this paper, we introduce a novel, fully parallelizable doubly-stochastic attention mechanism based on sliced optimal transport, leveraging Expected Sliced Transport Plans (ESP). Unlike prior approaches, our method enforces doubly stochasticity without iterative Sinkhorn normalization, significantly enhancing efficiency. To ensure differentiability, we incorporate a temperature-based soft sorting technique, enabling seamless integration into deep learning models. Experiments across multiple benchmark datasets, including image classification, point cloud classification, sentiment analysis, and neural machine translation, demonstrate that our enhanced attention regularization consistently improves performance across diverse applications. Our implementation code can be found at https://github.com/dariansal/ESPFormer.