SeTformer is What You Need for Vision and Language

TL;DR

SeTformer introduces Self-optimal Transport as a drop-in replacement for dot-product self-attention, embedding inputs into a Reproducing Kernel Hilbert Space and aligning them with a reference set via entropically regularized optimal transport. This yields non-negative, nonlinear attention weights with efficient computation through Nyström approximations and linear positional encoding, built into a hierarchical architecture. Empirical results across ImageNet-1K, COCO, ADE20K, and GLUE show SeTformer achieving state-of-the-art or competitive performance with substantially fewer parameters and FLOPs than strong baselines. The work demonstrates that content-based interactions can be effectively modeled through kernel-OT attention, enabling scalable, versatile transformers for both vision and language tasks.

Abstract

The dot product self-attention (DPSA) is a fundamental component of transformers. However, scaling them to long sequences, like documents or high-resolution images, becomes prohibitively expensive due to quadratic time and memory complexities arising from the softmax operation. Kernel methods are employed to simplify computations by approximating softmax but often lead to performance drops compared to softmax attention. We propose SeTformer, a novel transformer, where DPSA is purely replaced by Self-optimal Transport (SeT) for achieving better performance and computational efficiency. SeT is based on two essential softmax properties: maintaining a non-negative attention matrix and using a nonlinear reweighting mechanism to emphasize important tokens in input sequences. By introducing a kernel cost function for optimal transport, SeTformer effectively satisfies these properties. In particular, with small and basesized models, SeTformer achieves impressive top-1 accuracies of 84.7% and 86.2% on ImageNet-1K. In object detection, SeTformer-base outperforms the FocalNet counterpart by +2.2 mAP, using 38% fewer parameters and 29% fewer FLOPs. In semantic segmentation, our base-size model surpasses NAT by +3.5 mIoU with 33% fewer parameters. SeTformer also achieves state-of-the-art results in language modeling on the GLUE benchmark. These findings highlight SeTformer's applicability in vision and language tasks.

Figures (7)

• Figure 1: Top-1 classification accuracy vs. FLOPs on the ImageNet-1k, where the bubble size represents the number of parameters. SeTformer improves upon the baseline accuracy yet require fewer parameters and FLOPs.
• Figure 2: Top-1 Classification accuracy vs. FLOPs on the ImageNet-1k. We compare SeTformer with other kernel-based transformers, while using Swin-B as the base.
• Figure 3: An input feature vector $x$ is transported onto reference $y$ via transport plan $T(x, y)$, that aggregates $x$ features w.r.t. $y$, yielding $A_y(x)$. In DPSA, each $x_i$ aggregates with all $x$ features, forming a large sparse matrix. Our model aggregates based on best-matched $x$ and $y$ features through OT.
• Figure 4: Visualization of our model on COCO dataset using $\mathcal{K}$, and $\mathcal{K}_y$. They are compared with the Swin liu2021swin and NAT hassani2023neighborhood. While $\mathcal{K}$ focuses solely on large objects, $\mathcal{K}_y$ accurately captures multi-scale objects without including sparse areas.
• Figure 5: SeTformer Architecture (Left). It starts with a down-sampling convolutional layer, followed by four sequence stages containing multiple SeT blocks. Consecutive stages are bridged by down-sampler layers that reduce the spatial size while doubling the depth. On the right, we illustrate our attention formulation: mapping $x$ and $y$ elements to the RKHS, then aggregating $x$ features if they align well with the corresponding reference through OT computation between $x$ and $y$.