ESPFormer: Doubly-Stochastic Attention with Expected Sliced Transport Plans
Authors: Ashkan Shahbazi
Elaheh Akbari
Darian Salehi
Xinran Liu
Navid NaderiAlizadeh
Soheil Kolouri
Venue: ICML 2025 poster (2025)
Keywords: Geometric deep learning
Optimal Transport
Attention mechanism
doubly-stochastic matrices
Primary area: deep_learning->attention_mechanisms
TL;DR: A novel Doubly-Stochastic Attention mechanism using Expected Sliced Transport Plans offers faster, parallelizable computations with adaptive priors and competitive performance.

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 \url{https://github.com/dariansal/ESPFormer}.

ESPFormer:
Attention with Expected Sliced Transport Plans
Ashkan Shahbazi' Elaheh Akbari' Darian Salehi?
Xinran Liu
Navid NaderiAlizadeh Soheil Kolouri 14
Abstract
highly expressive and context-rich representations that accu- rately capture the essence of input data (Khan et al., 2022). Although first developed perform Natural Language Pro- cessing (NLP) tasks, the Transformer architecture has been adapted to wide range of domains such computer vi- sion (Dosovitskiy et al., 2021; Touvron et al., 202 1;;; 2022), graphs (Yun et al., 2019; Rampášek et al., 2022; Shirzad et al., 2023), point clouds (Zhao et al., 2021; Qin et al., 2022), and biological sequences (Jumper et al., 2021; Abramson al., 2024; Rives et al., 2021; Lin al., 2023; Hayes et al., 2025). Over the years, a considerable amount of research has been devoted to enhancing the classic Trans- former, focusing on better positional encoding (Dwivedi et al., 2021; Ying et al., 2021), efficiency of the attention mechanisms (Wang et al., 2020; Choromanski et al., 2020), and variants of self-attention (Hou et al., 2019; Ho et al., 2019), among others.
While self-attention has been instrumental in the success of Transformers, it can lead to over- concentration on a few tokens during training, re- sulting in suboptimal information flow. Enforcing doubly-stochastic constraints attention matri- ces has been shown to improve structure and bal- ance in attention distributions. However, existing methods rely on iterative Sinkhorn normalization, which is computationally costly. In this paper, we introduce a novel, fully parallclizable 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 en- hancing efficiency. To ensure differentiability, we incorporate a lemperature-based soft sorting tech- nique, enabling seamless integration into deep learning models. Experiments across multiple benchmark datasets, including image classifica- tion, 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 ht tps: / //githi hit. com/ dari waaaaa//sPForeeee....
Self-attention produces row-stochastic matrix, which can lead to few tokens dominating the attention distribution. To mitigate this, enforcing doubly-stochastic attention en- sures more balanced distribution across tokens. To that end, Sander et al. (2022) introduced Sinkformer, replac- ing the softmax normalization in the classic Transformer with Sinkhorn's algorithm (Sinkhorn, 1964), resulting in a doubly-stochastic attention matrix. Sander et al. (2022) es- tablish connection between self-attention matrices and the optimal transport problem by theoretically demonstrating that Sinkformers can be interpreted as a Wasserstein gradient flow for an energy minimization in the infinite depth limit. Indeed, doubly stochastic attention can be interpreted as a transport plan between queries and keys, further strengthen- ing the connection between optimal transport and attention mechanisms. Owing their inherent doubly stochastic structure, transport plans (i.e., couplings between keys and queries) naturally serve as strong candidates for attention matrices in this setting.
1. Introduction
The debut of TTrransfommers (Vaswani et al., 2017) marked a turning point in artificial intelligence and machine learning. Self-attention mechanisms excel at modeling the interac- tions among features, allowing Transformers to generate
'Department of Computer Science, Vanderbilt University, Nashville, TN, USA. 'Department of Computer Science, Duke University, Durham, NC, USA. 'Department of Biostatistics and Bioinformatics. Duke University. Durham, NC, USA. "Department of Electrical and Computer Engineering, Vanderbilt University, Nashville, TN, USA. Correspondence to: Ashkan Shahbazi vanderbilt.cdu>.
Nevertheless, computing optimal transport plans between keys and queries is computationally expensive, with a general complexity of O(N*) for N tokens. A more scalable approach is entropy-regularine 2013), which leverages the Sinkhorn algorithm to itera- tively approximate the transport plan, obtaining complex- ity OSSN2) where S denotes the number of iterations. The Sinkhorn algorithm (Sinkhorn, 1964), as used in Sinkform-
transport (Cuturi,
Proceedings of the 42nd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

ESPFormer: Doubly-Stochastic Attention with Expected Sliced Transport Plans
ers (Sander et al., 2022), repeatedly applies a series of row and column normalization steps until it reaches a desired level of convergence. Such an approach may introduce com- putational inefficiencies in scenarios where significant number of iterations are needed for the normalization to converge.
2. Background and Related Work
In this section, we first provide an overview of the exist- ing variants of softmax attention, including the doubly- stochastic attention using Sinkhorn's algorithm (Sander et al., 2022). Then, shift focus to the fundamen- tals of sliced optimal transport, with soft sorting reviewed for algorithmic concern.. Finally, we provide an overview of Expected Sliced Transport Plans (Liu et al., 2025) and soft sorting (Prillo & Eisenschlos, 2020), which serve as the cornerstones of our proposed doubly-stochastic attention mechanism, ESPFormer.
more computationally efficient alternative is the calcu- lation of optimal transport plans for one-dimensional dis- tributions, with a complexity of (N log N). This has motivated large body of work on sliced optimal trans- port methods that compare distributions by comparing their slices, i.e., -ieedinnsional marginals (Rabin et al., 2012; Kolouri et al., 2016; 2018; 2019; Deshpande et al., 2019; Nguyen et al., 2022; Nguyen & Ho, 2024). Although sliced optimal transport methods offer efficient metrics between distributions, they do not explicitly construct transport plans. Recent studies have addressed this limitation by developing methods to construct transport plans through slicing (Mahey et al., 2024; Liu et al., 2025). Liu et al. (2025) introduce the Expected Sliced Transport Plan (ESP), which leverages slicing to define computationally efficient metric while crucially providing an explicit transport plan by aggregat- ing lifted plans from all slices. This makes ESP strong candidate for doubly-stochastic attention mechanisms.
2.1. Variants of Softmax Attention
At the heart of the Transformer architecture the self- attention operation, crucial component that enables dy- namic pairwise interactions among tokens. In essence, it allows each position to "attend" to all others, with the degree of attention determined by how similar their representations are. Formally, let Wo: WK E R
Wv E Rdddd denote the query, key, and value matrices, respectively. Then, for a sequence (æ1, T", ,,,,, 32 E IRd, Vi, the output of the attention function for the i-th row, 2; can be written as
Σj= sim(Woz,, W W Σ5= sim(WQau WKT,)
(I)
In this work, we leverage the ESP framework to propose a novel doubly-stochastic attention mechanism. We refer to the resulting architecture as ESPFormer. Our specific contributions are as follows:
where sim(..) can be any similarity function, and the nor- malization for the similarities is applied row-wise. The clas- sic self-attention mechanism (Vaswani et al., 2017) lever- ages the softmax function to perform this row-wise nor- malization, i.c., the softmax of attention matrix C with Ci2 = (Wq.;) "Trg,,,TWWaaa Wr can be interpreted as row-wise nor- malization of exp(C). Alternative normalization operators have also been proposed in the literature. Some focus on normalizations that share the same properties as softmax but produce sparse outputs, such as SparseMax (Martins & Astudillo, 2016) and SparseK (Lou et al., 2024). Sparse- Max seeks the Euclidean projection of the input into the probabilistic simplex. Since this projection tends to hit the boundary of the simplex, SparseMax will output sparse probabilities. SparseK extends SparseMax by replacing the probabilistic simplex with a k-sum constraint. Others use different sim(:,-) functions in (1) to achieve linear com- plexity: Katharopoulos et al. (2020) and Han et al. (2023) use kernel smoothers to linearize the similarity calculations, and Li et al. (2020) propose to approximate the exponential function in softmax by the first-order Taylor expansion.
We propose ESPFormer, novel doubly stochastic at- tention mechanism built on the recently introduced Ex- pected Sliced Transport Plan (ESP) framework. ESP- Former ensures more balanced distribution of atten- tion across tokens while enabling control over the num- ber of tokens each token attends to via an inverse tem- perature parameter.
Through extensive experiments across diverse appli- cations, we demonstrate performance improvements over both classic Transformer and Sinkformer archi- tectures, along with enhanced computational efficiency compared to Sinkformer.
We show that replacing the classic attention mecha- nism in a pre-trained Transformer with ESPFormer and fine-tuning for a few epochs results in significant performance gains.
We show that by finetuning pretrained models with exponential temperature annealing schedule and then switching to hard sorting at inference, we obtain exact doubly stochastic matrices reducing complexity from O to (LN log N), and achieve consistent accuracy gains on the Cats vs. Dogs dataset.
Another line of research approaches the attention from the alignment/matching perspective and specifically utilizes concepts from the Optimal Transport (OT) theory. Zhang et al. (2021) propose to enforce the alignment between the distributions of keys and queries by adding a penalty term in the training objective, where the JensenShannon (JS)
We demonstrate the compatibility of our proposed at- tention mechanism with the recently introduced differ- ential attention architecture (Ye et al., 2025).
2

---

ESPFormer: Doubly-Stochastic Attention with Expected Sliced Transport Plans
divergence, the Wasserstein distance, and bi-directional con- ditional transport (Zheng & Zhou, 2021) are considered to define the matching cost. Xu & Chen (2023) introduce multimodal co-attention method that relies on OT to match patches of images and gene embeddings for the survival prediction task. Zhang et al. (2023) reinterpret slot attention within the OT framework and propose to enhance the per- formance of slot attention by minimizing the entropy of the Sinkhom divergence between two multisets, one containing inputs and the other containing slot features.
Kolouri et al., 2016; 2018; Deshpande et al., 2019; Kolouri et al., 2019; Nguyen et al., 2022; Nguyen & Ho, 2024). Sinkhorn's 's algorithm at the core of Sinkformer is used to solve the entropy retropy-regularized optimal transport problem with an improved complexity of O(22 (og N).
Sliced optimal transport operates by projecting high- dimensional distributions onto one-dimensional slices, lever- aging the key property that, in the one-dimensional case, a unique optimal transport plan exists, and the p- p-Wasserstein distance has a closed-form solution (Rabin et al., 2012),
2.2. -tobhy-Stochastic Attention
$
Wp(a", - (1
(w) du
(4)
In their pioneering work, Sander al. (2022) empirically observed that, during training, the row-stochastic attention matrices in classical Transformers tend to become approxi- mately stobhssstochaastc, with most column sums approach- ing 1. This suggests that Transformers inherently learn to distribute attention more evenly across tokens. In light of this finding, the authors propose the Sinkformer architecture, which replaces the softmax operation by the Sinkhorn's 's al- gorithm (Sinkhorn, 1964; Cuturi, 2013; Peyré et al., 2019) to enforce doubly-stochastic attention as an informative prior. They show that although both classic Transformers and Sinkformers can be understood as models that operate on discrete probability distributions, the Sinkformers have a special property that, in the infinite depth limit, they behave as a Wasserstein gradient flow for an energy minimization.
where FØi and F are the quantile functions of p² and 2² respectively, and the optimal transport map is given by T(x) = Fiz (F (z)) when /² is absolutely continuous. For empirical measures with N samples, the quantile func- tions correspond to the sorted samples that can be calculated in O log N). Then, by integrating the 1-dimensional Wasserstein distance over a set of L slices, the sliced Wasser- stein distance reduces the computational cost significantly to OLL log N).
2.4. Expected Sliced Transport Plans
Although the sliced Wasserstein distance offers a rapid and well-defined metric, it has one limitation: it does not gen- erate transport plan between the probability measures. It thus fails to explicitly provide how one distribution could be transported into another.
2.3. Sliced Optimal Transport
Consider the space of Borel probability measures on C Rd with finite p-th moment (p 2 1). denoted by (!)).. For μ 2² (2)), the classic optimal transport (in the Kantorovich formulation) solves the optimization problem
Liu et al. (2025) recently proposed the Expected Sliced Transport Plan (ESP), which defines an efficient transport plan as an aggregation of lifted optimal transport plans on 1-dimensional slices. Let μul EzER p(z)6, be a dis- crete probability measure in P(R)), that is, p(x) 2 0 for all x E Fd and EvER. p(z) that p(z) # 0 for at most countable many points x E Rd Similarly, let y2 = EyERs q(y),, be another probability measure in P( with at most countable support. For a given 0 E gd-1 the projected measures 0g and 0z are I-dimensional probability measures in P(R). and there exists unique optimal transport plan Ag" "N² between them. Equivalently, 0 and 0zpu² can be interpreted as proba- bility measures over quotient space IRd ~g where ~ig is defined as follows:
inf [ c(x, ((r,8))((zz))) TED(M)
(2)
We further assume
where F(u', ,22) denotes the set of all couplings between jet and y22, and : n2 R2 is lower semi-continuous function representing the transportation cost. Specifically, when c(., :) is the p-th power of metric, the Wassersteen distance is defined as
TET(U wa) (1s Cic-uuPdna vid--iid(a..)) (III--//ara
Wyln', ,p²)== ==
min
In practical settings, the Wasserstein distance between dis- crete probability measures can be obtained by solving linear program (Peyré et al., 2019) with an expensive com-
x ~ x' and only if 0. 0-r=0-2'
as each point on the slice R corresponds to an equiva- lent class of points in Hd that gets mapped to it by 0. With a slight abuse of notation, we denote the equiva- lent class of x E Rd by z0, referring to either a point in the quotient space Rd/ ~o or the set of points {a' e
putational complexity of (N log N) for sample size N.
Therefore, faster alternatives have been extensively stud- ied, including mntropy-regularinn optimal transport (Cu-
turi, 2013) and sliced optimal transport (Rabin et al., 2012;
3

---

ESPFormer: Doubly-Stochastic Attention with Expected Sliced Transport Plans
Doubly Stochastic Attention for Slice
Queries
Keys
Sorted Q
Sorted
Sort Along
Tokens
Tokens
Axis
Stices
Slices
Correspondence Matrix
Doubly Stochastic Attentions for All Slices
Doubly Stochastic ESP Attention
T 
Aggregate
T == 100
Figure 1. overview of the proposed ESP attention mechanism. By integrating the slicing operator into the key and query matrices, each dimension is treated as learnable slice. For cach slice, tokens are (soft) sorted, and doubly-stochare correspondence matrix computed between the keys and queries. Finally, these correspondence matrices across dimensions are aggregated to form a single attention matrix.
Rd 8 8 = interchangeably. Then, we can write 04p EXER\/O P( where P(z") EVER' p(x'), and 22 = EgeEr~o (0r") where Eyg (1)). This quotient space interpretation of the I-dimensional distributions 0gp' and u allows us to construct a lifted transport plan in the original space Rd using the optimal transport plan Add uh'.
2.5. Soft Sorting
Calculating the sliced Wasserstein distance involves evalu- ating the quantile functions of the distributions, which, in the discrete case, be boiled down the sorting oper- ation. Sorting is one of the most common operations in computer science. Yet, the piecewise- linear sorted value function and the integer-valued rank/argsort operators pose a significant obstacle for gradient-based optimization tech- niques, which are essential deep learning, as neither of them is differentiable. To incorporate sorting operations into the backpropagation framework, differentiable approxi- mations, known soft sorting, have been explored. Exam- ples include smoothed rank operators by adding Gaussian noise (Taylor et al., 2008) and by using sigmoid surrogate functions (Qin et al., 2010), parameterizing permutations in terms of a differentiable relaxation (Mena al., 2018), and relaxing the permutation matrices to be only row-wise stochastic (Grover et al., 2019). Of note, Cuturi al. (2019) propose differentiable proxy by viewing sorting optimal assignment problem and relaxing it to entropic optimal transport problem from the input values to an auxil- iary probability measure supported on an increasing family of target values.
ΣΣ FERd
Σ yERR
ug u u² " (z,)))(z,)),
16
where the transported mass un n²² is defined as
p(z)q(U) No (((F".09))
ut (I. P(=")Q((")
Then for a given distribution of slicing directions (f's): a E P(gd-1), the ESP qud .²² E l(p', .1²) is defined as an expectation o/2 n² over o:
zu'ta Egno h ue2
iec 70t (((x,))))
gg- 1tth (((2)))))do00))
Prillo & Eisenschlos (2020) introduce simple yet highly effective continuous relaxation of the argsort operator using the softmax function. Given vector R
Liu et al. (2025) have shown that the associated cost,
[d[ssor(
(2) (v)I" 1
SoftSort (v) = softmax
(7)
,x²)
Σ
Σ x
Pll 2²² (((x)))})
yERd
approximates the sorting permutation matrix, where the softmax operator is applied row-wise, d(.,-) can be any differentiable remi-metric, and is temperature parameter that controls the degree of the approximation.
a well-defined distance and equivalent to the Wasserstein distance.

---

ESPFormer: Doubly-Stochastic Attention with Expected Sliced Transport Plans
3. Method
with Dz(ut μ²: 0) representing the p-th power of the in- duced transport cost by Uo, defined
3.1. ESP for Uniform Diserete Distributions
N N
Given our application of interest, we specifically focus on uniformly distributed discrete distributions with an equal number of support points, while the method can be readily extended
DEGU ,2:00)
ΣΣ 1-1
1., - Us (U (10)
where (f(((T; (((((x;))))))) is the (i.j) component of the Ug matrix. When T = 0, = = for all I E [1,2, L], then G is simply the mean of all the lifted plans.
unequal number of support points. For given N E N, denote the space of uniform discrete distribu- tions with N support as
When handling uniformly distributed discrete distribu- tions with different numbers of support points (e.g., in cross-attention), the transport plan involving mass split- ting can be approximated using linear interpolation. For 1
N Σ
} fI, .N .
(R)) =
6x RR
Eラ 5ys E (R*) and X = [21, D2 _N E RNx = (21, 1/, . UN E Nxd be the corresponding matrix forms. Denote the symmetry group of order N in matrix representation as Sy. that is. Sny contains all permutation matrices of set of N elements.
Let
N NV Ei=1 drg E P(N) (R=), 12 === N
4² N E dng u²² M ンj= dy, with N + M, let A' E R RNNN and B' E IR Mx respectively denote their sorting permutation matrices. We define an interpolation matrix I E R where
it1
M
Consider the I-dimensional slice of 4² and n²² in the 0 direc- tion: 4 1/4 Σi= 80.2, and 2xp² N Σj=1 60-y,- with XO- = 00 0 22 .0 IN 1"" E EN and
[b.
if
if N
S Mi
(1I)
0,
elsewhere.
The transport plan is then given by Ug 2/4 IB'
YO- 00 y
0 1/2. ....... 0 yNm" E RR
There exists
permutation matrices A, B SN such that AX0 e R and BYO R.V are in sorted order, i.e.,
3.2. ESP Attention
To bridge the abstract formulation of the Expected Sliced Plan (ESP) with its practical use in attention mechanisms, we interpret the query and key matrices as empirical prob- ability measures due to their permutation-equivarian na- ture. Following our notation in Section 2.1, let We, WK e Rmdd Wv E Rdxd denote the query, key, and value ma- trices, respectively, and let X = [x1, T2, ................................ .N]] v, e Rd Vi denote the input tokens. Also, let Q = WqX RmxN and K Wx X RmxN denote the queries and keys, respectively. Then, we think of doubly tention as a transport plan between the empirical measures n& and y.x defined as
(AX0); S (AX0)2 S ... S (AX0)N; (BY S (BYO), S ... (BYO)N.
Then the optimal matching between 0² and 0z can be described by
(AXO), + (BYO),, Vn E 11,2..... .N,
or equivalently, given that AT A11, A B represents the transport map from X0 to YO, i.e.,
at-
(X0), + (A" BYO),, Vn E 11,2,...,AA. NN.
N
N
1
By lifting this transport map from the 0 slice to the origi- nal space Rd, we have Up = 1 A which represents a transport plan from X to Y.
4
Q = N
Σ
8q
Σ 3=1
6ky
N
N
To create such transport plan using the ESP framework presented in Section 3.1, we use a set of L slicing directions θ = 00, 0z]" E Lxm .f E gm Then, the rows of OK and OQ contain the projected keys and queries for each slice. The transportation plan Uo, is, then, calculated for each slice by soft-sorting the projected keys and queries as described in Section 3.1. Finally, the attention matrix is calculated from (8).
Finally, for a given histogram of o Et_ with Ei_ `. = 1, the ESP for uniformly-distee distrributt discrete distributions with the same number of supports is defined as
G== Σ
of Uo.
where 7 denotes the "inverse temperature" hyprruarretpyprrrrmtmteeeee and of is defined as
A key consideration is the choice of slices, O. While classic sliced OT uniformly samples slices from gmm -1 prior work suggests learning them (Deshpande et al., 2019; Naderial- izadeh al., 2021; NaderiAlizadeh et al., 2025). Though
Dz(n' (10
of
L Dv=1 e- -D w2:00)
5

---

ESPFormer: Doubly-Stochastic Attention with Expected Sliced Transport Plans
g
S
5
g
5
Sinkhorn
==0
7 = 0.1
==
T 100
ESP
Figure The attention weights between an example pair of keys (red) and queries (green) obtained by Sinkhom's 's algorithm (top row) with different numbers of iterations and by Expected Sliced Transport Plans (bottom row) with different inverse temperature values. Note that for Sinkrhorn and zero iterations, i.e. S the computed attention reduces to classic self-attention. The weights are represented by the width of the lines connecting each pair of points.
this approach is common, learning θ introduces additional parameters, increasing the total count in ESP attention and complicating fair comparisons with other mechanisms. No- tably, since the input distributions (keys and queries) are themselves learned, optimizing θ may be unnecessary. In- stead, the distributions can adapt to a fixed slicing scheme, as done in prior work, such as FSPool (Zhang al., 2020). To avoid extra parameters, in this work, we propose ing axis-aligned slices by setting O Ixx, the identity matrix. This leads to
Algorithm 1 ESPFormer's Doubly-Stochastic Attention
input Query matrix Q E R
N Key matrix K E RR*N,
Value Matrix V E Id SoftSort hyperparameter t, and "inverse temperature" hyperparameter 7.
output Attention weighted output matrix Calculate the pairwise distance matrix:
[C], l|Q K.yl1²
for 1 - 1 to m do SoftSort the projected samples using (7):
ESP-Attention (1 K, V) VG,
(12)
A = SoftSort, (Qt:), B₁ = SoftSort
where = Wyx E Rd is the value matrix. Figure I illustrates the construction of the ESP matrix G, which serves the proposed transport-based attention map, direct- ing how information flows from V to the output. Moreover, Figure 2 compares classic attention (Sinkhorn 5º), Sinkhorn attention from (Sander et al., 2022) across different iteration counts, and our ESP attention under varying "inverse tem- perature" hyperparameters. As can be seen, while ESP pro- vides a balanced distribution of attention due to its double stochasticity, the "inverse temperature" parameter controls the number of other tokens each token should pay atten- tion to. The overall pipeline of ESPFormer can be found in Algorithm 1. For completeness, we also report results from explicitly learning the slices and observe no significant benefits compared to axis-aligned slices.
4: Calculate the transport plan Ui = # B:
Calculate Dj = E,,[Cl,
end for
Calculate the o~ softmax D;
8: Aggregate the plans from all slices G = >m oF
9: Return: VG
tions, where d corresponds to the output dimension. The subsequent SoftSort operation, which involves computing pairwise distances, incurs a complexity of O(N2) per slice. When applied across all m slices, this yields com- plexity of (m.N2). The computation of transport plans through matrix multiplications with soft permutation trices contributes an additional ((mN2) operations. The final aggregation of transport plans requires summing m plans of size N x N, also contributing OmmN2) operations. Therefore, the overall runtime complexity of ESPFormer is O(mN (N + d)). In comparison, Sinkhorn's 's algorithm
3.3. Runtime Complexity
ESPFormer begins with query and key projections (Q = WoX X and K = Xx.),, each requiring O(mNd) opera-
6

---

ESPFormer: Doubly-Stochastic Attention with Expected Sliced Transport Plans
formers (Guo al., 2021). Table 2 presents the classification results across four runs, comparing different attention mech- anisms integrated into Set Transformers and Point Cloud Transformers, Notably, ESPFormer outperforms competing methods, demonstrating superior performance.
Rundine (seconds)
4.2. Sentiment Analysis
We next evaluate ESPFormer on the IMDB dataset (Maas et al., 2011) for sentiment analysis. Following Sander et al. (2022), our architecture comprises an attentionbased encoder followed by maxpooling layer trained to clas- sify movie reviews as either positive negative. Table 3 presents the accuracy improvements achieved by using ESP- Former as the core attention module, compared to baseline models, showcasing its robustness in text classification.
16-1
Sequence Length
ESPFomer (Softsarting)
Sinkformer [5=3)
Sinkformer (5-15)
ESPFormer (Hardsortingg
Sinkformer 1S=10)
Sink (Sm20)
Figure 3. Runtime analysis of ESPFormer versus Sinkformer
with varying iteration counts (S) for sequence lengths N E (50, 100, 500, 1000} with d = 1024, averaged over 10 runs. ESPFormrr achieves superior computational efficiency across all sequence lengths with hard sorting, and outperforms Sinkformer for S 3 under soft sorting, while maintaining model expressivity.
We additionally evaluate ESPFormer on the TweetEval senti- ment dataset (Barbieri al., 2020), where tweets are labeled positive or negative. We use the same encoder + maxpooling architecture and training setup for IMDb. Table 4 reports the corresponding accuracy metrics.
exhibits a runtime complexity of O((S + m2) where S denotes the number of iterations. A key distinction lies in the parallelization capabilities: while Sinkhorn necessitates sequential processing over S iterations. ESPFormer enables parallel computation across the m slices. This parallelizabil- ity allows ESPFormer to scale efficiently with increasing m. As demonstrated in Figure 3. ESPFormer offers superior efficiency at all sequence lengths with hard sorting, and surpasses Sinkformer for S > 3 when using soft sorting. The wall clock runtime analysis of ESPFormer compared to all the baselines can be found in Appendix A.
4.3. Neural Machine Translation
We evaluate ESPFormer and Sinkformer using two reference models: the Transformer and its Transraasfrr counter- part (Ye et al., 2025). Both models are trained using the fairseq sequence modeling toolkit (Ott et al., 2019) on the IWSLT' 14 German toe-tt-----e to-English dataset (Cettolo et al., 2014). The architecture of both models consists of an encoder and decoder, each with a depth of 6 layers. Initially, we trained both models for 25 epochs using the standard training pro- cedure. After this phase, we performed a Plug-and-Play evaluation, where the attention heads were plugged into the pre-trained models and evaluated on their performance, as shown in Table 5. In this phase, we tested ESPFormer and Sinkformer attention heads, denoted by the respective mod- els in the table, comparing their performance the base Transformer and DiffTransformer models. Following the Blug-and-Play evaluation, we performed a Fine-Ture Tune Boost phase, where the models were further fine-tuned for an addi- tional 10 epochs. The teeituning led to further performance gains, as observed the table, with ESPFormer showing the highest improvement in both the Transformer and DiffTrans- former settings, achieving the best BLEU score of 34.64 and 34.83, respectively. The fine-tuned results show the effectiveness of incorporating ESPFormer into the model architecture.
4. Experiments
In this section, we evaluate the performance of ESPFormer across a diverse set of domains, including image classifica- tion using ViTs (Dosovitskiy al., 2021), point cloud clas- sification using set transformers (Lee et al., 2019) and point cloud transformers (Guo et al., 2021), as well as sentiment analysis and neural machine translation via Transformers, highlighting the applicability of our approach in comparison to Differential Transformer (Ye et al., 2025), Sinkformer (Sander et al., 2022), and the Vanilla Transformer (Vaswani al., 2017). Throughout our experiments, we employed an axis-aligned slicer. our experimental setup and implementa- tion can be found in Appendices B and C.
4.1. ModelNet40 Classification
The ModelNet40 dataset (Wu et al., 2015) comprises 40 widely recognized 3D object categories and serves as a stan- dard benchmark for point cloud classification. Transformers designed for point clouds and sets have been extensively evaluated on MedelNet40, with notable examples including Set Transformers (Lee et al., 2019) and Point Cloud Trans-
4.4. Vision Transformers
To evaluate the generalizability of the models under lim- ited data scenarios, we conducted experiments on the Cats and Dogs dataset (Kaggle, 2013) using varying fractions of the training data: 1%, 10%, 25%, and 100%. We train a
7

---

ESPFormer: Doubly-Stochastic Attention with Expected Sliced Transport Plans
Table Average and standard deviation (over 3 runs) of ESPFormer's classification accuracy (%) vs. baselines the Cats and Dogs dataset under varying data availability. ESPFormer's performance is reported in three modes: initial soft sort, sharp soft sort, and hard sort.
Baselines
ESPFormer
Data Fraction
Sinkformer
DifTTransformer
Transformer
Initial Soft Sort
Sharp Soft Sort
Hard Sort
$5.07 3.34
53.78 0.28
49.71 0.31
$5.66 3.95
$7.86 3.77
58.52 3.73
10%
69.56 0.32
67.34 + 0.11
57.25 0.22
71.49 0.43
72.22 0.37
72.71 0.36
25% %
74.56 + 0.58
74.86 + 0.17
72.25 + 0.16
75.40 + 0.38
75.92 0.31
75.92 0.28
100%
79.12 + 0.17
78.85 + 0.11 78.49 + 0.09
79.47 + 0.12
80.61 + 0.11
81.23 0.11
Table Test accuracy (%) on the ModelNet40 dataset over runs.
Table S. Plug-and-Play and Fine-Tune Boost performance of Trans- former and DranTransformrr Baselines on the IWSLT14 German- to-English dataset, reported as the median over 4 runs. Results with a sign denote the Plug-and- angaan performance of a diffcrent attention modulc than the base attention module.
Accuracies marked with are reported from (Sander et al., 2022).
Model
Best
Median
Mean
Worst
Set Translormer'
87.8
86.3
85.8
84.7
Set Transforrre
89.0
88.7
88.7
88.6
Set Sinkformer"
89.1
88.4
88.3
88.1
Set ESPFormer
89.6
89.5
89.4
89.1
Model
Plug-and-Play Fine-Tune Boost
Point Cloud Transformer'
93.2
92.5
92.5
92.3
Transformer
33.40
34.61
Point Cloud Dift Transfrmerr
93.1
92.8
92.7
92.6
Sinkformer
33.36*
34.61
Point Cloud Sinkformer"
92.8
92.7
92.5
ESPFormer
33.38*
34.64
Point Cloud ESPFormer
93.2
92.9
92.7
92.6
DiffTraasformer
33.85*
34.78
Table 3. Test accuracy (%) for Sentiment Analysis on IMDb.
Sinkformer
33.67*
34.81
Model
Best
Median
Mean
Worst
ESPFormer
33.72*
34.83
Transformer
85.30
85.25
85.25
85.20
DiffTraasformer
85.50
85.45
85,45
85.40
TRANSITION TO HARD SORTING
Sinkformer
85.40
85.39
85.37
85.30
During training, the use of soft sorting- depends on temperature parameter-results in approximate doubly stochastic matrices (DSMs). However, this approximation can be addressed by gradually annealing the temperature during training, effectively transitioning toward hard sorting. Consequently, inference time, soft sorting can be replaced with hard sorting to yield exact DSMs. This not only ensures exact computations but also reduces the inference time from O to O(LN logNN. where L is the number of slices.
ESPFormer
85.50
85.50
85.47
85.40
Table 4. Test accuracy (%) for Sentiment Analysis on TweetEval.
Model
Best
Median
Mean
Worst
Transformer 71.50
71.35
71.31
71.10
DiffTransforerr
72.60
72.35
72.31
72.00
Sinkformer
72.40
72.30
72.23
71.90
ESPFormer 72.60
72.40
72.36
72.10
To further elarify this point, we conducted additional ex- periments the Cats vs. Dogs dataset using the same frac- tions of the training data to evaluate the effect of hard sorting inference time. We finetuned pretrained mod- els for 40 epochs using an exponential temperature anneal- ing schedule to progressively sharpen the sorting behav- ior. Specifically, the temperature at epoch e is defined as: temp(e) = y temp(e 1), with 7 = 0.8. This decay schedule reduces the temperature to approximately 10 making the soft sorting operation nearly equivalent a hard permutation. After finetuning. we replaced soft sorting with hard sorting for final evaluation. Table 1 demonstrates that fine-tuning with temperature annealing and switching to hard sorting yields consistent improvements in test accuracy.
ViT and modify the attention mechanism accordingly. The models compared include ESPFormer, Sinkformer, Diff- Transformer, and the standard Transformer. The goal was to assess how well each model performs when faced with pro- gressively larger amounts of data, particularly in resource- constrained settings. Table summarizes the classification accuracy for each model under different data fractions. Our results demonstrate that ESPFormer consistently outper- forms the other models across all data fractions, achieving the highest classification accuracy at each level of data avail- ability. With only 1% of the data, ESPFormer achieves 55.66% accuracy, significant improvement over the Trans- former at 49.71% %. As the data availability increases, ESP- Former continues to show superior performance. This sug- gests that ESPFormer is particularly robust in data-scarce environments, showcasing its ability to generalize well even under limited training data.
ON THE CHOICE OF SLICER
We employ axis-aligned slicing because our keys and queries are learned parameters; any optimal slice orientation
8

---

ESPFormer: Doubly-Stochastic Attention with Expected Sliced Transport Plans
Table Accuracy (%) over three runs (mean + standard deviation) across slicer types, slice counts I, and inverse temperature 7.
L
L
L 128
0.1
Learnable
74.30 + 0.48
78.70 + 0.32
79.10 + 0.22
78.40 + 0.26
76.20 0.42 79.60 + 0.28
Frozen
66.50 + 0.52
72.80 + 0.38
78.30 + 0.18
79.20 + 0.30
Axis-Aligned
79.47 0.12
1.0
Learnable
74.30 = 0.48
79.07 + 0.30
78.20 + 0.35
77.80 - 0.21
74.10 + 0,46
Frozen
66.50 0.52
73.11 + 0.43
77.95 + 0.25
78.80 0.29
78.40 0.27
Axis-Aligned
78.85 1 0.31
T
Learnable
74.30 + 0.48
79.15 + 0.29
78.06 + 0.27
77.10 + 0.23
74.15 0.44
Frozen
66.50 + 0.52
73.45 + 0.41
76.85 + 0.24
77.85 + 0.30
78.10 0.26
Axis-Aligned
77.75 + 0.32
is thus implicitly encoded in the projection matrices WQ and WK- This avoids introducing extra learnable parame- ters and ensures a fair comparison with baselines such as vanilla attention and Sinkformer.
based attention mechanisms as principled alternative to existing methods, paving the way for future research in efficient and structured attention models.
Limitations
However, for ablation study, we ran additional experi- ments the Cats vs. Dogs dataset, varying the number of slices L. the inverse temperature т, and whether the slicer was learnable or frozen. Tables 6 summarizes the results. Across all settings, axis-aligned slicing provides a simple, in- terpretable baseline without parameter overhead. Learnable slicers can boost performance when L is small, by focus- ing on the most informative directions, but their advantage diminishes grows. Frozen slicers likewise require suffi- ciently many slices approximate the distribution structure effectively, since performance degrades when L is low in high-dimensional spaces.
Our method has two main limitations. First, incom- patible with strictly causal Transformer or autoregressive architectures, since enforcing both lower-triangularity and doubly stochasticity reduces the mapping to the identity, Second, during training its memory usage scales linearly with the number of slices, we must retain each slices permutation matrix; after fine-tuning and switching to hard sorting, however, these matrices can be discarded and all inference computations can be run off-GPU.
Impact Statement
ESPFormer significantly accelerates the computation of bal- anced attention by introducing an efficient mechanism for enforcing doubly stochastic constraints through the math- ematies of sliced optimal transport. The proposed atten- tion improves model robustness and interoperability, and is shown to be effective even in a plug-and-play setting, This enables practical use of balanced attention as a strong inductive bias, particularly effective in low-data regimes, Furthermore, by establishing a connection between sliced optimal transport and attention mechanisms, ESPFormer opens a new research direction for structured and efficient attention through transport-based formulations. Although ESPFormer has the potential to impact broad range of applications, we do not foresee any specific societal risks requiring special attention.
5. Conclusion
We introduced ESPFormer, a novel Transformer architecture that integrates fast, stobllcstochastic attention mechanism based on Expected Sliced Transport Plans (ESP). By lever- aging differentiable soft sorting and fully parallelizable approach, ESPFormer provides a computationally efficient alternative to iterative Sinkhorn normalization while ensur- ing a balanced distribution of attention across tokens. Our experiments across diverse domains, including image clas- sification, point cloud processing, sentiment analysis, and neural machine translation, demonstrate that ESPFormer consistently outperforms both classical self celf-attention and Sinkhorn-based alternatives in terms of accuracy and effi- ciency. Furthermore, we showed that our proposed attention can be seamlessly integrated into pre-trained Transform- ers, improving performance even with minimal fine-tuning. The flexibility of ESPFormer also makes it compatible with emerging differential attention mechanisms, e.g., (Ye al., 2025), expanding its applicability to broad range of archi- tectures. These findings highlight the potential of transport-
A Acknowledgment
This work was partially supported by the NSF CAREER Award #2339898 and the Wellcome Leap 'Surgery: Assess/- Validate/Expand (SAVE)' program.

---

ESPFormer: Doubly-Stochastic Attention with Expected Sliced Transport Plans
References
Han, D., Pan, X., Han, Y., Song, S., and Huang, G. Flatten transformer: Vision transformer using focused linear at- tention. In Proceedings of the IEEE/CVF international conference computer vision, pp. 5961-5971, 1 2023.
Abramson, J., Adler, J., Dunger, J., Evans, R., Green, T.,
Pritzel, A., Ronneberger, O., Willmore, L., Ballard, A. J., Bambrick, J. et al. Accurate structure prediction of biomolecular interactions with alphafold 3. Nature, pp-
Hayes, T., Rao, R., Akin, H., Sofroniew, N. J., Oktay, D., Lin, Z., Verkuil, R., Tran, V. Q., Deaton, J., Wiggert, M., et al. Simulating 500 million years of evolution with a language model. Science, pp. cads0018, 2025.
1-3, 2024.
Barbieri, F., Camacho-Collados, J., Espinosa-Anke, L., and Neves, L. EveetEval:Unified Benchmark and Compara- tive Evaluation for Tweet Classification. In Proceedings of Findings of EMNLP, 2020.
Ho, J., Kalchbrenner, N., Weissenborn, D., and Salimans, T. Axial attention multidimensional transformers. arXiv preprint 1exiv:/9122 2800, 2019.
Cettolo, C., Niehues, M., and Federico, M. The iwslt 2014 evaluation campaign. In Proceedings of the In- ternational Workshop on Spoken Language Translation (IWSLT), 2014.
Hou, R., Chang, H., Ma, B., Shan, S., and Chen, X. Cross attention network for few-shot classification. Advances in neural information processing systems, 32, 2019.
Choromanski, K., Likhosherstov, V., Dohan, D., Song. X., Gane, A., Sarlos, T., Hawkins, P., Davis, J., Mohiuddin, A., Kaiser, L., et al. Rethinking attention with performers. arXiv preprint arXiv:2009. 14794, 2020.
Jumper, J., Evans, R., Pritzel, A., Green, T., Figurnov, M., Ronneberger, O., Tunyasuvunakool, K., Bates, R., Zídek, A., Potapenko, A., et al. Highly accurate protein structure prediction with alphafold. nature, 596(7873):583-589) 2021.
Cuturi, M. Sinkhorn distances: Lightspeed computation of optimal transport. Advances in neural information processing systems, 26, 2013.
Kaggle. Dogs cats. https: wwwp///wwww//.oo http://wwwwwwwgg...eeoo www.ke.cg.l..oo..m... c ///dogs-ooe -oggs-vee vs cats, 2013. Accessed: 2025-01-29.
Cuturi, M., Teboul, O., and Vert, J. Differentiable rank- ing and sorting using optimal transport. In Wallach, H., Larochelle, H., Beygelzimer, A., d'Alché-Buc, F., Fox, E., and Garnett, R. (eds.), Advances in Neural Information Processing Systems, volume 32, Curran Associates, Inc., 2019.
Katharopoulos, A., Vyas, A., Pappas, N., and Fleuret, F. Transformers are mns:: Fast autoregressive transformers with linear attention. In International conference ma- chine learning, pp. 5156-5165. PMLR, 2020.
Khan, S., Naseer, M., Hayat, M., Zamir, S. W., Khan, F. S., and Shah, M. Transformers in vision: A survey. ACM computing surveys (CSUR), 54 54(10s): 1-11, 2022.
Deshpande, I., Hu, Y.-T., Sun, R., Pyrros, A., Siddiqui, N., Koyejo. S., Zhao, Z., Forsyth, D., and Schwing, A. G. Max-sliced Wasserstein distance and its use for GANs. In Proceedings of the IEEE/CVF conference computer vision and pattern recognition, pp. 10648-10656, 2019.
Kingma, D. P. and Ba, J. Adam: A method for stochastic op- timization. 3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Track Proceedings, 2015. URL http:/ /ett:///rxiv..r org abs 111211122 6980.
Dosovitskiy, A., Beyer, L.. Kolesnikov, A., Weissenborn, D., Zhai, X., Unterthiner, T., Dehghani, M., Minderer, M., Heigold, G., Gelly, S., Uszkoreit, J., and Houlsby, N. An image is worth 16x words: Transformers for image recognition at scale. ICLR, 2021.
Kolouri, S., Zou, Y., and Rohde, kernels for probability distributions. In Proceedings of the IEEE Conference Computer Vision and Pattern Recognition, pp. 5258-5267, 2016.
Sliced wasserstein
Dwivedi, V. P., Luu, A. T., Laurent, T.. Bengio, Y., and Bresson, X. Graph neural networks with learnable structural and positional representations. arXiv preprint arXiv:2/ 10.0787,. 2021.
Kolouri, S., Rohde, G. K., and Hoffmann, H. Sliced Wasser- stein distance for learning Gaussian mixture models. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pp. 3427-3436, 2018.
Grover, A., Wang, F., Zweig, A., and Ermon, Stochastic optimization of sorting networks via continuous relax- ations. arXiv preprint arXiv:
2019.
Kolouri, S., Nadjahi, K., Simsekli, U., Badeau, R., and Rohde, G. Generalized sliced Wasserstein distances. Advances in neural information processing systems, 32, 2019.
Guo, M.-H., Cai, J.-X., Liu, Z.-N., Mu, T.-.,, Martin, R. R., and Hu, S.-M. Pet: Point cloud transformer. Computers & Graphics, 102:1-13. 2021.
10

---

ESPFormer: Doubly-Stochastic Attention with Expected Sliced Transport Plans
LeCun, Y., Bottou, L.,, Bengio, Y., and Haffner, P. Gradient- based learning applied to document recognition. Proceed- ings of the IEEE, 86( 36(11):2278-2324, 1998.
NaderiAlizadeh, N., Salehi, D., Liu, X., and Kolouri, S. Con- strained sliced Wasserstein embedding. arXiv preprint arXiv:2506.02200 02203, 2025.
Lee, J., Lee, Y., Kim, J., Kosiorek, A. R., Choi, S., and Teh, Y. W. Set transformer: A framework for attention- based permutation-invarian neural networks, In Interna- tional Conference on Machine Learning, pp. 3744-3753. PMLR, 2019.
Nguyen, K. and Ho, N. Energy-based sliced wasserstein distance. Advances in Neural Information Processing Systems, 36, 2024.
Nguyen, K., Ren, T., Nguyen, H., Rout, L., Nguyen, T., and Ho, N. Hierarchical sliced Wasserstein distance. arXiv preprint arXiv:2209.13572 2022.
Li, R., Su, J., Duan, C., and Zheng. S. Linear attention mech- anism: An efficient attention for semantic segmentation. arXiv preprint arXiv:2007. 14902, 2020.
Ott, M., Edunov, S., Grangier. D., and Le, Q. V. fairseq: A fast, extensible toolkit for sequence modeling. In Pro- ceedings of NAACL-HLT 2019: Demonstrations, 2019. 2019..
Lin, Z., Akin, H., Rao, R., Hie, B., Zhu, Z., Lu, W., Smetanin, N., Verkuil, R., Kabeli, O., Shmueli, Y., al. Evolutionary-scale prediction of atomic-level protein structure with a language model. Science, 379(6637); 1123-1130, 2023.
Peyré, G., Cuturi, M., et al. Computational optimal trans- port: With applications data science. Foundations and Trends@ in Machine Learning, 11(5-6):355-6077 2019.
Prillo, S. and Eisenschlos, J. Softsort: A continuous relax- ation for the argsort operator. In International Conference on Machine Learning. pp. 7793-7802. PMLR, 2020.
Liu, X., Martín, R. D., Bai, Y., Shahbazi, A., Thorpe, M.. Aldroubi, A.. and Kolouri, S. Expected sliced transport plans. In The Thirteenth International Confer- ence on Learning Representations, 2025. URL https: //openreview. net/ forum? wttprummmddppppoedeee
Qin, T., Liu, T.-Y., and Li, general approximation framework for direct optimization of information retrieval measures. Information retrieval, 13:375-397 77775--97727 2010.
Lou, C., Jia, Z., Zheng, Z., and Tu, K. Sparser is faster and less more: Efficient sparse attention for long-range transformers. arXiv preprint arXiv:2406.1677
Qin, Z., Yu, H., Wang, C., Guo, Y., Peng, Y., and Xu, K. Geometric transformer for fast and robust point cloud registration. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 11143- 11152, 2022.
2024.
Maas, A. L., Hunley, R., Chen, D., Y, N. B. N.. and Ng, A. Y. Learning word vectors for sentiment analysis. In Proceedings of the 49th Annual Meeting of the Associa- tion for Compulational Linguistics (ACL), pp. 142-150, 2011.
Rabin, J., Peyré, G., Delon, J., and Bernot, M. Wasser- stein barycenter and its application to texture mixing. In Scale Space and Variational Methods in Computer Vision: Third International Conference. SSVM 2011, Ein-Gedi, Israel, May 29-June 2, 201 1, Revised Selected Papers 3, pp. 435-446. Springer, 2012.
Mahey, G., Chapel, L.,, Gasso, G., Bonet, C., and Courty, N. Fast optimal transport through sliced generalized wasser- stein geodesics. In Seirth-seventh Conference on Neural Information Processing Systems, 2024. URL https: //openreview net /oret//foummm forum? 3xmnnnxXdnnn.....
Rampášek, L.,, Galkin, M., Dwivedi, V. P., Luu, A. T., Wolf, G., and Beaini, D. Recipe for general, powerful, scal- able graph transformer. Advances in Neural Information Processing Systems, 35:14501-14515, 2022.
Martins, A. F. T. and Astudillo, R. From soft- max to sparsemax: A sparse model of attention and multi-label classification. ArXiv, abs/1602.02068. 2016. URL https:/ /tpp//app..a//// semanticscholas DB:/CoooppSS::1143320
Rives, A., Meier, J., Sercu, T., Goyal, S., Lin, Z., Liu, J., Guo, D., Ott, M., Zitmick, C. L., Ma, J., et al. Biological structure and function emerge from scaling unsupervised learning 250 million protein sequences. Proceedings of the National Academy of Sciences, 118(15):c20162391 188, 2021.
Mena, G., Belanger, D., Linderman, S., and Snock, J. Leam- ing latent permutations with gumbel- sinnbel- networks. arXiv preprint arXiv: / 1802.08665, 2018.
Naderializadeh, N., Comer, J. F., Andrews, R., Hoffmann, H., and Kolouri, S. Pooling by sliced- Wasserstein em- bedding. Advances in Neural Information Processing Systems, 34:3389-3400, 2021.
Sander, M. E., Ablin, P.,, Blondel, M., and Peyré, G. Sink- formers: Transformers with doubly stochastic attention. In International Conference on Artificial Intelligence and Statistics, pp. 3515-3530 PMLR, 2022.

---

ESPFormer: Doubly-Stochastic Attention with Expected Sliced Transport Plans
Shirzad, H., Velingker, A., Venkatachalam, B., Sutherland, D. J., and Sinop, A. K. Exphormer: Sparse transform- ers for graphs. In International Conference Machine Learning, pp. 31613-31632, PMLR, 2023.
Ying, C., Cai, T., Luo, S., Zheng, S., Ke, G., He, D., Shen, Y., and Liu, T.-Y. Do transformers really perform badly for graph representation? Advances in neural information processing systems, 34:28877 288888 2021.
Sinkhorn, R. A relationship between arbitrary positive ma- trices and doubly stochastic matrices. The annals of mathematical statistics, 35(2)877-877, 1964.
Yun, S., Jeong, M., Kim, R., Kang, J., and Kim, H. Graph transformer networks. Advances in neural information processing systems, 32, 2019.
Taylor, M., Guiver, J., Robertson, S., and Minka, T. Soft- rank: optimizing non non-smooth rank metrics. In Proceed- ings of the 2008 International Conference on Web Search and Data Mining, pp. 77-86, 2008.
Zhang, S., Fan, X., Zheng, H., Tanwisuth, K., and Zhou, M. Alignment attention by matching key and query dis- tributions. Advances Neural Information Processing Systems, 34:13444-13433 2021.
Touvron, H., Cord, M., Douze, M., Massa, F., Sablayrolles, A., and Jegou, H. Training data-efficient image transform- & distillation through attention. In International Con- ference on Machine Learning, volume 139, pp. 10347- 10357, July 2021a.
Zhang, Y., Hare, J., and Prgel-Bennett, A. Espool: Learn- ing set representations with featurewise sort pooling. In International Conference on Learning Representations, 2020. URL https: http://///ppeeeeeeeeeeee net/ ffrrmm?r i id=HJgBA2VYWH:
Zhang, Y., Zhang. D. W., Lacoste-Julien, S., Burghouts, G. J., and Snoek, C. G. Unlocking slot attention by changing optimal transport costs. In International Con- ference on Machine Leurning, pp. 2023.
Touvron, H., Cord, M., Sablayrolles, A., Synnaeve, G., and Jégou, H. Going deeper with image transformers. In Proceedings of the IEEF/CVF International Conference Computer Vision (ICCV) pp. 32-42, October 2021b.
PMLR,
Touvron, H., Cord, M., and Jegou, H. Deit ii:: Revenge of the vit. arXiv preprint 2rXiv:2204.07118, 2022.
Zhao, H., Jiang. L.. Jia, J., TomI P. H., and Koltun, V. Point transformer. In Proceedings of the IEEE/CVF interna- tional conference computer vision, pp. 16259-16268, 2021.
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L. u., and Polosukhin, I. At- tention is all you need. In Guyon, I., Luxburg, U. V., Bengio, S., Wallach, H., Fergus, R., Vishwanathan, S., and Garnett, R. (eds.). Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017.
Zheng, H. and Zhou, M. Comparing probability distribu- tions with conditional transport. stat, 1050:9, 2021.
Wang, S. Li, B. Z.. Khabsa, M., Fang, H., and Ma, H. Linformer: Self-attention with linear complexity. arXiv preprint 20XX::0000604768. 2020.
Wu, Z., Song, S., Khosla, A., Yu, F., Zhang, L., Tang, X., and Xiao, shapenets: deep representation for volumetric shapes. In Proceedings of IEEE conference on computer vision and pattern recognition, pp. 1912- 1920, 2015.
Xu, Y. and Chen, H. Multimodal optimal transport-based co-attention transformer with global structure consistency for survival prediction. Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 21241- 21251, 2023.
Ye, T., Dong, L., Xia, Y., Sun, Y., Zhu, Y., Huang, G., and Wei, F. Differential transformer. In The Thirteenth International Conference on Learning Representations, forum?
2025. URL https:/ stttp:://pperreieee forum? 1=G
12

---

ESPFormer: Doubly-Stochastic Attention with Expected Sliced Transport Plans
A. Full runtime wall-clock analysis, and Patch size impact
A.1. Full runtime wall-clock
The runtime comparison Figure 4 highlights the computational efficiency of different attention mechanisms across varying sequence lengths. Transformer and DiffTransformer exhibit the lowest runtime due to their standard self-attention mechanism, making them computationally lightweight. In contrast, ESPFormer achieves balance between efficiency and expressivity, maintaining lower runtimes compared to Sinkformer across all sequence lengths. These results emphasize the trade-off between computational efficiency and model expressivity when selecting attention model for large-scale tasks,
10
Runtime (seconds)
10-4
100
500
10000
Sequence Length (N)
Transformer
ESPFormer (Hardsorting) Sinkformer (S=3)
Sinkformer (S=15)
DiffTransformrr
Sinkformer (S=20)
ESPFormer (Softsorting)
Sinkformer (S-10)
Figure Runtime analysis Transformer, DiffTranstormer, ESPFormer, and Sinkformer with varying iteration counts (S) sequence lengths N E {50, 100, 500, 1000} with d 1024, averaged over 10 runs. ESPFormer consistently demonstrates superior computational efficiency compared to Sinkformer across all sequence lengths, while Transformer and Trinsfrrso remain the most lightweight.
A.2. Impact of Patch Size
Similar to Sander et al. (2022) and to analyze the effect of patch size model performance, trained Transformer, Diffrranffrr,r, Sinkformer, and ESPFormer on the MNIST dataset (LeCun et al., 1998). To isolate the impact of the attention mechanism, we used single-layer self-attention module without feed-forward layers. Figure illustrates the test accuracy as a function patch size for each model. ESPFormer consistently outperforms competing models, particularly for smaller patch sizes, highlighting its superior ability to capture fine-grained details. As the patch size increases, the performance gap narrows, with all models converging similar accuracies likely due to the reduced information content within each individual patch.
Acounny
Transformer
Striformee
ESPFomer
Figure 5. Comparison of MNIST test accuracy across different patch sizes for Transformer, Dif.Transformer Sinkformer, and ESPFormer architectures. Results averaged over 3 runs.
13

---

ESPFormer: Doubly-Stochastic Attention with Expected Sliced Transport Plans
B. Implementation Details
Sinkhorn's Algorithm
We implement Sinkhorns algorithm in the log domain enhance numerical stability, Given a matrix S -R2x= defined as S c for some cost matrix C E IR*"" the algorithm iteratively updates the scaling vectors (f.9) RR" R such that the limiting matrix is given by
g = diag(e" (isagge)))"0 diag(e" eM
Starting with g" = On. the updates are performed in the log domain as follows:
g+ log(1n/n)) log(Se")), if is even,
+1 log(lnn////gooo (n log(s" et), I is odd.
This log-domain formulation allows for efficient and numerically stable computations by leveraging the Log-sum-exe trick to evaluate expressions like log(Se") and Tog(STef) ef").
Implementation of Sinkhorn Algorithm:
def sinkhorn_log_dee dorgg_ddmainnnmmmm (C, epsilen-le-3, nun_iters=50); C. shape [0]
np.exp epsilon)
Compute 0-S
f, = np. (n), zerozero (n) # Initialize log scaling vectors
for in range interss
if
f up. (1 / n) Leg__uu (np. log (S) + g)
else:
np. 1og (1 / n) _og_sur_exp (np log (8.T) f)
return
f, g
For further details and implementation, please refer to the SinkFormer GitHub Repository.
Differential Transformer
The Differential Transformer extends traditional self-attention by introducing differential attention mechanisms, which modulate the contribution of multiple attention maps using learnable parameter This allows the model to capture finer differences the relational structure of the input data.
Given an input sequence X, the attention mechanism computes different sets of queries and keys, denoted (Q1, K1) and (Q2, K2). The attention output is computed using weighted difference of the attention maps:
A = softmax(Q1K)) À x s)
The resulting attention map is then applied to the values V, followed by normalization and projection:
0 = GroupNorm
Implementation of Differential Attention:
def Di fiffAttn(e (X, W_q, W_k, W_,, \lambda): Q1, Q2 split(X (X G (_g)
14

---

ESPFormer: Doubly-Stochastic Attention with Expected Sliced Transport Plans
K1, K2 (x W_k)
V X W_y
8 = 1 / sqrt (d)
Al Q1 @ Kl.transpose -2)
A2 Q2 t2.transpose (-1, -2)
return (softmax (A1) \lambda softmax (A2)) V
Implementation of Multi-Head Differential Attention:
def Mult Head (X, R_q, W_k, W_,, c )lambda) GroupNorm (Difffttn (X, w_qi, W_k1, W_vi, (lambda) for i in range (h) 1) /lambda ((mbda(init)) (init))
return Concat (0) 0 W_
For further details and implementation, please refer to the Differential Transformer Github Repository.
C. Experiment Details
C.1. Table of Notations
Table 7 includes the notations used in this section.
Table 7. Experimental Notations
Notation
Description
Temperature parameter in the Soft Sorting algorithm
The "inverse temperature" hyperparameter in ESP
The entropy regularization parameter in the Sinkhorn algorithm
S
Number iterations Sinkformer Initial value in DiffTraasformer Initial learning rate for all methods The interpolation matrix in ESPformer
Ainitial
Insittial
C.2. ModelNet 40
In our experiments on ModelNet40 utilizing Set Transformers, we begin by preprocessing the dataset and uniformly sampling 5000 points from each object.
Set Transformers
The model architecture consists of two Induced Set Attention Blocks (ISAB) in the encoder, followed by decoder incorporating Set Attention Block (SAB) and a Pooling by Multihead Attention (PMA) module. The training procedure employs a batch size of 64 and utilizes the Adam optimizer (Kingma & Ba, 2015). The network trained for 300 epochs, with an initial learning rate 10-3 which reduced by a factor of 10 after 200 epochs.
S
Anitial
I Irmittall
10 3
10
10
21
0.8
10-3
Identity XM
Table 8. Hyperparameters used in training for Set Transformers on ModelNet40 dataset.
Point Cloud Transformers
The training is conducted with a batch size of 32 using Stochastic Gradient Descent (SGD) (Ruder, 2016). The model undergoes 300 training epochs, starting with an initial learning rate of 10 which is reduced by a factor after 250 epochs.
15

---

ESPFormer: Doubly-Stochastic Attention with Expected Sliced Transport Plans
Amitial
Iriitiall
10-3
10-1
10-1
21
0.8
10-4
IdenlilyNxN
Table 9. Hyperparameters used in training for Point Cloud Transformers ModelNet40 dataset.
C.3. Sentiment Analysis
For our sentiment analysis experiments, we utilize the publicly available implementation from the nip-tutorial' repository, where pretrained Transformer model is fine-tuned the IMDb dataset. In our experimental setup, reset the parameters of the pretrained Transformer and train it from scratch the IMDb dataset. The model architecture consists of depth of 6 layers and employs 8 attention heads. Training is conducted using the Adam optimizer with batch size of 32 over 15 epochs. The initial learning rate is set to 10 and is reduced by a factor of 10 after 12 epochs.
7
Ainitial
Iniitiall
10
10+1
10
15
0.33
10-4
IdentityxxN
Table 10. Hyperparameters used in training for Transformers on IMDb dataset.
We also utilize this implementation for TveetEval, and similar hyperparameters for this sentiment analysis experiment can be found in table 11.
S
Ainitial
Irinitial
10-3
10+1
10-1
15
0.33
10-4
IdentilyNxx
Table 11. Hyperparameters used in training for Transformers on TweetFval dataset.
C.4. Neural Machine Translation
For our neural machine translation experiments, we adopt the Transformer model from fai rser along with its Draf'raan- former counterpart, training both from scratch for 25 epochs. We then fine-tune them alongside other baselines for an additional 10 epochs on the IWSLT' 14 dataset?. When fine-tuning ESPFormer and Sinkformer, we modify the original training schedule by reducing the learning rate by a factor of 10.
S
Amitial
Ir Irinitial
10-3
10+1
10-1
15
0.0
10-4
IdentityNxh
Table 12. Hyperparameters used in training for neural machine translation IWSLT14 dataset.
C.5. Vision Transformers
C.5.1. CATS AND DOGS CLASSIFICATION
For this experiment, we use the VIT model (Dosovitskiy et al., 2021) with different attention mechanisms. The images are resized to 224 224. We use a ViT architecture with an embedding and MLP dimension of 128, 6 layers, attention heads, and patch size of 16. For methods and percentages of data, we train the model for 300 epochs. use an initial learning rate of 10 5 After 250 epochs, the learning rate reduced by a factor of 10. Training done using the Adam optimizer (Kingma & Ba, 2015) and batch size of 64. Our experimental setup, including the normalizations and data augmentations, are consistent with the Cats and Dogs experiment in Sinkformer (Sander et al., 2022). For each percentage of the data, three random seeds are used to initialize and generate new subsets of data. The subsets are consistent across all methods to ensure the same training set is used.
'https: h/tp::///ttt thub. commm11eee com/1yeoni Inip-tutorial /ttpe///asterrccttttllaaaase ota trrnsfffrree Pttpps: //g www.bb.c..pppp.cc www.///www//////oooooooooooooeet fairseq/tb ob/ /aaommannn /wwn///eaalllee/rrassslssse RRADDRR.
16

---

ESPFormer: Doubly-Stochastic Attention with Expected Sliced Transport Plans
S
Aniitil
Irinitial
10-3 3
10-1
3
0.5
3
10-5
IdenlilyNxx
Table 13. Hyperparameters used in training for ViT Cats and Dogs dataset.
C.5.2. IMPACT OF PATCH SIZE
To analyze the effect patch size on final accuracy, we conduct experiments using a batch size and the Adam optimizer. The model architecture consists of single-layer Transformer (depth with one attention head, no non-linearity, and varying patch sizes. Training performed over 45 epochs, with an initial learning rate of 10-3 for the Transformer and DiffTransfrrmrr and 2 10 for the ESPFormer and Sinkformer. The learning rate is decayed by factor of 10 after 35 epochs, and again by another factor of 10 after 41 epochs.
S
Anitial
Iriittial
10 3
1
0.5
x 10
2
10
Identity N
Table 14, Hyperparameters used in training for shallow- VIT on MNIST dataset.
17


Reviews:
[{"review_id": "YbvUVCcKcD", "reviewer": "Reviewer_eE5t", "summary": "The paper introduces ESPFormer, a attention mechanism that enforces doubly-stochastic constraints in attention matrices without requiring iterative Sinkhorn normalization. Instead, it leverages Expected Sliced Transport Plans (ESP) to achieve a fully parallelizable and computationally efficient solution. The authors integrate a temperature-based soft sorting technique to ensure differentiability, making ESPFormer compatible with deep learning models. Experimental evaluations across diverse applications, including image classification, point cloud classification, sentiment analysis, and neural machine translation, demonstrate the advantages of ESPFormer over traditional self-attention and Sinkhorn-based attention methods.\n\n## update after rebuttal\n\nI increased the score", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their constructive feedback. Below are our responses.\n\n**1. Reliance on Soft sorting:**\n\nSoft sorting introduces a temperature hyperparameter that requires tuning. To address this, we use *temperature annealing*, an exponential decay schedule, to smoothly transition from soft to hard sorting. This allows the model to adapt as sorting sharpens, enabling hard sorting at inference without performance loss.\nThis improves runtime, reducing complexity from \\( O(N^2) \\) to \\( O(N \\log N) \\), and reduces our dependency on soft-sorting. See [Figure 3](https://github.com/Anon4142/ESPFormer/blob/main/figures/esp_sink_runtime.png) for updated runtimes. The table below summarizes results on Cats vs. Dogs across training fractions.\n\n| Data Fraction | Initial Test Acc (Soft Sort) | After Temp Decay (Sharp Soft Sort) | After Switching to Hard Sort |\n|---------------|------------------------------|-------------------------------------|-------------------------------|\n| 1%            | 59.87%                        | 60.50%                              | 61.02%                        |\n| 10%           | 71.71%                        | 72.16%                              | 72.66%                        |\n\nExperiments for the 25% and 100% training splits are currently in progress and will be included in the **camera-ready version** of the paper.\n\n\n**2. Additional experiments:**\n\nWe conducted experiments on TweetEval. The table below summarizes the test performance of several models, including ESPFormer, on this benchmark. ESPFormer demonstrates competitive results compared to the baselines, outperforming both Sinkformer and the standard attention mechanism (denoted by \"Vanilla\"). For a controlled comparison, we replaced the attention module in a standard 6-layer Transformer (as proposed by Vaswani et al., 2017) with alternative components—Sinkformer, DiffAttention, and ESPFormer—while keeping the rest of the architecture unchanged. As illustrated, all variants outperform the vanilla baseline, with ESPFormer and DiffAttention achieving the highest accuracies on this benchmark.\n\n\n| Model    | Vanilla | Sink. | Diff.Att. | ESP. |\n| -------- | ------- | ----- | --------- | ---- |\n| Accuracy | 71.5    | 72.0  | 72.6      | 72.6 |\n\nNote that this evaluation is limited to a single architecture and run due to the rebuttal timeframe. We plan to extend it in the camera-ready version with plug-and-play experiments across diverse attention modules and pretrained language models.\n\n\n**3. Axis-aligned slices:** \n\nWe chose axis-aligned slices because the keys and queries are learned parameters. Thus, any potential optimization of slice orientations is implicitly captured by the query and key matrices, $W_Q$ and $W_K$. This choice enabled us to avoid introducing extra learnable parameters to ensure a fair comparison with baseline methods, such as vanilla attention and Sinkformer.\n\nTo investigate the effect of the number and type (learnable vs. frozen) of slices, we conducted additional experiments on the Cats vs. Dogs benchmark, varying the number of slices $L$, the inverse temperature parameter $\\tau$, and whether the slices were frozen or learnable. The results are presented below.\n\n| **$\\tau=0.1$**   | $L = 1$ | $L = 8$ | $L = 32$ | $L = 64$ | $L = 128$ |\n| ---------------- | ------- | ------- | -------- | -------- | --------- |\n| **Learnable**    | 74.45%  | 78.56%  | 79.33%   | 78.25%   | 76.15%    |\n| **Frozen**       | 66.61%  | 72.75%  | 78.41%   | 79.39%   | 79.66%    |\n| **Axis-Aligned** | —       | —       | —        | 79.47%   | —         |\n\n| **$\\tau=1.0$**   | $L = 1$ | $L = 8$ | $L = 32$ | $L = 64$ | $L = 128$ |\n| ---------------- | ------- | ------- | -------- | -------- | --------- |\n| **Learnable**    | 74.45%  | 79.07%  | 78.10%   | 77.64%   | 74.34%    |\n| **Frozen**       | 66.61%  | 73.09%  | 77.89%   | 78.88%   | 78.50%    |\n| **Axis-Aligned** | —       | —       | —        | 78.86%   | —         |\n\n| **$\\tau=10$**    | $L = 1$ | $L = 8$ | $L = 32$ | $L = 64$ | $L = 128$ |\n| ---------------- | ------- | ------- | -------- | -------- | --------- |\n| **Learnable**    | 74.45%  | 79.24%  | 78.06%   | 77.07%   | 74.20%    |\n| **Frozen**       | 66.61%  | 73.50%  | 76.79%   | 77.91%   | 78.13%    |\n| **Axis-Aligned** | —       | —       | —        | 77.78%   | —         |\n\nLearnable slicers may help with fewer slices by focusing on informative directions, but their advantage diminishes as the slice count increases. In contrast, frozen slicers avoid added complexity but require many slices to capture distributional structure effectively, especially in high dimensions.\n\nOnce again, we thank the reviewer for their time and consideration,.\n\n **P.S.** Due to the limited rebuttal period, the results presented are based on a single run. We recognize the importance of reporting performance across multiple seeds and will include averaged results in the camera-ready version.", "claims_and_evidence": "1. ESPFormer provides a more computationally efficient alternative to Sinkhorn-based doubly-stochastic attention mechanisms.\n\n2. The proposed attention mechanism leads to better-balanced attention distributions, improving model performance.\n\n3. ESPFormer enhances downstream tasks such as classification and translation compared to classical self-attention and Sinkhorn-based approaches.\n\n4. Fine-tuning pre-trained models with ESPFormer leads to significant performance improvements.\n\nThese claims are supported by experimental evidence on multiple datasets. The runtime complexity analysis demonstrates computational advantages, and performance metrics from classification and translation tasks confirm improved accuracy.", "methods_and_evaluation_criteria": "The methods proposed are well-motivated for addressing the limitations of standard self-attention and Sinkhorn normalization. The evaluation criteria include:\n\n1. Accuracy on benchmark datasets for classification tasks.\n\n2. BLEU scores for machine translation tasks.\n\n3. Runtime complexity comparisons between ESPFormer and Sinkformer.\n\n4. Ablation studies on hyperparameters such as inverse temperature.\n\nThese metrics appropriately reflect the benefits of the proposed method in terms of both efficiency and performance.", "theoretical_claims": "N/A", "experimental_designs_or_analyses": "The experiments are well-structured, covering a range of tasks that benefit from improved attention mechanisms. The paper provides:\n\n1. Comparisons against baseline models (Vanilla Transformer, Sinkformer, DiffTransformer) across various datasets.\n\n2. Analysis of computational efficiency via runtime evaluations.\n\n3. Fine-tuning experiments to validate the plug-and-play nature of ESPFormer.", "relation_to_broader_scientific_literature": "ESPFormer extends prior work on optimal transport-based attention mechanisms by addressing computational bottlenecks and using sliced optimal transport plans.", "essential_references_not_discussed": "While the paper discusses relevant prior work, additional references on non-iterative attention normalization techniques or alternative soft sorting methods might provide further context.", "other_strengths_and_weaknesses": "## Strengths\n\n1.  Introduces a novel, computationally efficient alternative to Sinkhorn-based attention mechanisms, eliminating the need for iterative normalization.\n\n2. The method is fully parallelizable, making it more scalable than Sinkhorn-based approaches, which require sequential iterations.\n\n3. Strong theoretical grounding connects ESP attention with optimal transport theory, ensuring the approach is well-principled.\n\n4. Extensive experiments demonstrate the effectiveness of ESPFormer across a variety of tasks, including image classification, point cloud classification, sentiment analysis, and neural machine translation.\n\n5. The paper provides clear empirical evidence that replacing self-attention mechanisms in pre-trained models with ESPFormer leads to performance improvements with minimal fine-tuning.\n\n6. The introduction of the inverse temperature hyperparameter allows fine-grained control over attention sparsity, adding adaptability across different tasks.\n\n## Weaknesses \n\n1. The reliance on soft sorting introduces an additional hyperparameter (temperature), which may require careful tuning for different tasks.\n\n2. Additional experiments on language model is needed to demonstrate further the performance of ESPFormer.", "other_comments_or_suggestions": "Since language model  is one of the most important application of Transformer, experiments on  language model (e.g., on WikiText-103) should be considered.", "questions_for_authors": "1. How does ESPFormer handle very long sequences compared to efficient self-attention variants like Linformer? Is it possible to reduce further the complexity of ESPFormer?\n\n2. I understand that the choice of Identity matrix is for computational convenience since it does not require projection. However, it is worth to try if a fixed set of directions $\\theta_1,\\ldots,\\theta_L$ which are not the standard basis can affect the performance (L can be larger than d). This set of directions can be randomly sampled from the uniform distribution or using quasi Monte Carlo methods to encourage repulsive structure [1].\n\n[1] Quasi-Monte Carlo for 3D Sliced Wasserstein, Nguyen et al.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "BRbW5BzJgb", "reviewer": "Reviewer_ZJyu", "summary": "This paper is about achieving doubly stochastic attention in Transformers through ESP (expected sliced transport plans). The softmax is known to be a notorious bottleneck for the expressivity and gradient flow in Transformers, thus is clearly a relevant research area. Moreover, prior art like the Sinkformer rely on iterative approximation procedures to obtain DSMs. The authors ensure differentiability of their ESP approach and and show empirically that performance improves across multiple datasets. Moreover, the authors introduce a flexible hyperparameter (tao) which allows to control the temperature of the attention while ensuring doubly stochasticity.\n\n## update after rebuttal: \nI had originally given a score of 4 to this paper due to high methodological novelty, sound experimental design and clear relevance for the field. The score was contingent on the author's response to some raised concerns and questions, which were adequately addressed. I keep my score.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the constructive feedback. Below, we provide our responses.\n\n**1. The Birkhoff Polytope.**\nWe thank the reviewer for raising this insightful question. \nIn the context of the Birkhoff polytope $B_N$, we will consider the transport plans between $\\mu=\\sum_{i=1}^N\\delta_{x_i}$ and $\\nu=\\sum_{j=1}^N\\delta_{y_j}$, i.e., two measures having the same number of supports $N$ and uniformly-distributed mass. To avoid confusion, we refer to the \"identity matrix\" in $B_N$ as the outer product plan and denote the $N\\times N$ matrix with all entries equal to $1$ as $\\mathbb{1}_{N\\times N}$. The following fact about $B_N$ is used in the following analysis [1]:\n* (Birkhoff-von Neumann Theorem) $B_N$ is a convex polytope whose extreme points, or corners, are the $N!$ permutation matrices.\n1. When $\\tau\\rightarrow\\infty$, a permutation matrix is recovered. In the original EST work [2], the authors have mentioned that the distribution of $\\theta$ will degenerate into a one-hot vector as $\\tau\\rightarrow\\infty$. The EST will recover min-SWGG [3], where the minimum is taken over all permutation matrices induced by the slices, resulting in a permutation matrix. By the Birkhoff-von Neumann Theorem, this will be one of the corners of $B_N$. \n2. When $\\tau=0$, the EST is not at the origin, i.e., the outer product measure $\\frac{1}{N^2}\\mathbb{1}_{N\\times N}$. In Figure 6 of [2], EST with varying $\\tau$ is compared with the Optimal Transport plan, the entropic Optimal Transport plan, and the outer product plan. EST is evidently different from the outer product measure. It is worth mentioning that entropic Optimal Transport plan interpolates between the outer product and the Optimal Transport plan, whereas EST does not.\n3. In [this figure](https://github.com/Anon4142/ESPFormer/blob/main/figures/dist_perm.png), we provide illustration for $\\tau=0$ and $\\tau\\rightarrow \\infty$ using a pair of distributions $\\mu, \\nu$, each supported on 3 points. The corners of $B_3$ consists of $S_0, S_1, \\cdots, S_5$. The pie chart in this figure shows the optimal transport plan (i.e., the optimal permutation) as a function of the slicing angle. As can be seen $S_1$ and $S_3$ are the highly voted transportation plans. When $\\tau=0$, the EST is a convex combination of the corners with weights indicated in the pie chart.  When $\\tau\\rightarrow \\infty$, the EST will recover the one permutation that gives the minimal transport cost, which is $S_1$ in this example. \n\n[1] Birkhoff, G., 1946.\n[2] Liu X, et al. 2024.\n[3] Mahey G, et al. 2023.\n\n**2. What $\\tau$ to choose?**\nThe optimal value of the hyperparameter $\\tau$ depends on the input probability measures (i.e., the task) and the number of slices. Smaller $\\tau$ values promote more diffuse attention, while larger values lead to sharper, more focused (i.e., near one-to-one) attention between tokens. We ran more ablative studies on the choice of $\\tau$ and slicing strategies. Please see responses to reviewer eE5t.\n\n**3. Source code availability.**\nWe apologize for the unclear placement of the code link, which was included as a footnote on page 7. In the camera-ready version, we will move it to the abstract for better visibility. Meanwhile, the anonymous GitHub repository is available [here](https://github.com/Anon4142/ESPFormer).\n\n**4. Exact/Approximate DSMs.**\nSoft sorting yields approximate DSMs during training, but annealing the temperature gradually transitions it toward hard sorting and exact DSMs. At inference, hard sorting can replace soft sorting, reducing complexity from $O(N^2)$ to $O(LN\\log N)$. We will include this discussion in the revised paper.\n\nTo demonstrate, we fine-tuned ESPFormer for 40 epochs on Cats vs. Dogs (1% and 10% data) and replaced soft sorting with hard sorting at inference. Results show improved accuracy and efficiency:\n\n| Data Fraction | Initial (Soft) | After Annealing | Hard Sort |\n|---------------|----------------|------------------|-----------|\n| 1%            | 59.87%         | 60.50%           | 61.02%    |\n| 10%           | 71.71%         | 72.16%           | 72.66%    |\n\nExperiments on 25% and 100% data are ongoing and will be included in the camera-ready version.\n\n**5. Updating Figure 3.**\nWe have updated Figure 3 to include the runtime for the Sinkhorn method with $S=3$. We also added the inference runtime for ESP via hardsorting (i.e., inference time). The revised Figure 3 is available [here](https://github.com/Anon4142/ESPFormer/blob/main/figures/esp_sink_runtime.png).\n\n**6. On the complexity of ESP, Sinkhorn and Vanilla Attention.**\nWe observe that with SoftSort, ESP achieves lower runtime complexity than Sinkhorn for $S \\geq 5$, though it remains slower when $S = 3$ (for training). While $S = 3$ or $5$ may be sufficient for some Sinkformer tasks, others, such as ModelNet40, require more iterations (e.g., $S = 21$). We will clarify this and include the complexity of vanilla attention in the main manuscript.\n\nThank you again for your consideration.", "claims_and_evidence": "The general claims of the paper are well supported by experimental evidence. \n- The ESPFormer ensures a more balanced distribution (than Softmax) and gives flexibility via inverse temperature $\\tau$.\n- Performance improvement over Transformer (and even Sinkformer and DiffFormer) is clear in all cases, even though the hyperparameter sweeps to reach these results are not made explicit, so it is unclear how the authors arrived e.g., at different values of $\\tau$.\n- ESPFormer can work as a drop-in replacement for models trained with standard attention of DiffAttention and can still boost performance during finetuning.\n\nOverall, I think that this is an exceptional  and very innovative paper. I still have some concerns that require clarification but I believe that this is a very valuable piece of research that will help shaping the future of attention and Transformers.", "methods_and_evaluation_criteria": "The authors empirically evaluate ESPFormers across a rich set of domain and show convincing results everywhere. So in general, evaluation criteria are absolutely convincing.\n\n\nHowever, at least one ablation study on the impact of $\\tau$ is necessary for the paper. From the Appendix, I understand that the authors tuned tau specifically for each dataset, yet they are not transparent on the conducted hyperparameter sweeps, raising some questions on how many attempts the ESPFormer was given to obtain the results finally shown in the paper.", "theoretical_claims": "(1) FIg1 suggests (visually speaking) that the slice-specific DSMs are actually permutation matrices rather than interior points of the Birkhoff polytope. Compare eq. (8), this seems to originate from the deltas being diracs (BTW, this is never said explicitly!), yet in Eq (8) neither the sigmas nor the Us seem to be binary to me. Please clarify! Moreover, if the slices of G are indeed permutation matrices, maybe you want to emphasize the Birkhoff von Neumann theorem which is that every DSM (every element in the birkhoff polytope) can be expressed as a convex combination of permutation matrices. This brings me to another question about empirical and theoretical expressivity. Can you guarantee that every DSM can be reached with your algorithm? The BvN theorem requires non-negative coefficients theta which are more flexible then the simple averaging you do to obtain the ESP Attention matrix. Ultimately, may the ESPFormer be complemented with an extra set of L parameters enforced to sum to 1 that allow a weighted average of the slice-specific DSMs in order to obtain full expressivity over the Birkhoff polytope?\n\n(2) \nWhether the runtime of $O(LN logn)$  for transport plan computation is lower than $O(N^2 logN)$ depends on the number of slices. Since you avoid adding extra parameters (which I agree is a good choice), you set $L=m$ which is called $d_k$ in standard attention. Indeed, typically $d_k << N$ in transformers in practice, still it should be emphasized that we are not talking about a theoretical speedup, just about a practical one, due to common design choices when deciding about transformer hyperparameters. Nobody prevents me from setting $d_k > N$ and in that case, IIUC, the runtime complexity would be lower.\n\n(3) Section 3.2 -- Dimensionalities requires clarification. It seems there is a mixup on the matrix dimension because they are swapped? I recommend to follow the notation from the original paper by Vaswani et al, where we start with X of size $N x d$. Instead you seem to start from X of shape $d x N$, so you are left-multiplying the inputs with the $WQ$ matrix to obtain Q and K of shape m x N. Standard attention runs $QK^T$ to get a matrix of shape $NxN$ whereas in your case you have to compute $K^TQ$ to get a $NxN$ matrix which is very unconventional. To ease readability, I suggest to update all notation. For the same reason in (12) you arrive at $VG$ rather than $GV$. \nAlso W_V needs to be quadratic according to your definition, whereas in general it does not have this constraint. Please clarify", "experimental_designs_or_analyses": "The experimental results are convincing and well designed, just ablation study and interpretations on the impact of $\\tau$ are missing, ideally on more than one dataset.\n\nFig 2: Effectively, the choice of tao seems to interpolate between the extremes of tao=100 where the attention matrix is a permutation matrix and tao=0 where the DSM is almost the identity matrix -- in other words we seem to move between the vertices and the origin of the Birkhoff polytope. Can you confirm (or even prove) this intuition? If yes, it migh be worth emphasizing this in the text to ease interpretability of the parameter tao.", "relation_to_broader_scientific_literature": "This paper suggests a clear improvement over the Sinkformer which has itself shown to improve over standard attention in Transformer. It removes the limitation of Sinkformer that only produces DSMs as iterative approximations and shows superior empirical results.", "essential_references_not_discussed": "The authors did a great job at revising prior literature about adaptations of Transformers, especially toward doubly stochastic attention. An alternative potential avenue for doubly stochastic Transformer could be quantum computing as was highlighted in the conclusion of this paper from ICML last year: https://proceedings.mlr.press/v235/mariella24a.html", "other_strengths_and_weaknesses": "Positive:\n- Strong theoretical motivation, clear novelty\n- Strong empirical results at only mild runtime increase\n\nNegative:\n- No source code is provided, hampering reproducibility\n- A small discussion on when to chose which tao would be helpful", "other_comments_or_suggestions": "- Fig 2: Effectively, the choice of tao seems to interpolate between the extremes of tao=100 where the attention matrix is a permutation matrix and tao=0 where the DSM is almost the identity matrix -- in other words we seem to move between the vertices and the origin of the Birkhoff polytope. Can you confirm (or even prove) this intuition? If yes, it migh be worth emphasizing this in the text to ease interpretability of the parameter tao.\n\n- Are your DSMs exact or approximate? In the Sinkformer they are only approximate due to the iterative nature of Sinkhorn's Algorithm.\n\n- For Figure 3, for completeness S=3 should be added since it is empirically often enough as shown in the Sinkformer.\n\nSection 3.3:\n- Please mention for completeness also the complexity of the standard attention following the same notation\n- Whether your complexity is better than the Sinkformer, in theory, depends on (1) whether m < S and (2) d < N. Regarding (1): Typicall no because the Sinkformer shows good performance with S being as small as 3 or 5 and observes saturation around S=20. Regarding (2): typically yes, your d should correspond to d_v in the standard Transformer). Even if this is more of a theoretical than a practical perspective, I feel amplifying it would improve clarity\n\n\n\nMinor: \nL412 (right side) - missing space", "questions_for_authors": "-", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "rpu05jTy88", "reviewer": "Reviewer_p6Sy", "summary": "This paper introduces ESPFormer, a novel Transformer architecture integrating a fast, doubly-stochastic attention mechanism based on Expected Sliced Transport Plans (ESP). By projecting high-dimensional queries/keys into 1D slices using axis-aligned directions (Θ = I), ESPFormer efficiently computes optimal transport plans via differentiable soft sorting, ensuring end-to-end trainability. The framework aggregates sliced plans using an inverse temperature parameter τ to balance attention sparsity and distribution, achieving O(mN²) runtime with full parallelization across m slices—significantly faster than Sinkhorn’s O(SN²) iterative approach. Experiments demonstrate consistent performance gains across diverse tasks: +0.5% accuracy on ModelNet40 (point clouds), +0.2 BLEU on IWSLT translation, and 6% improvement on Cats&Dogs with 1% data.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the constructive feedback. Below are our responses.\n\n**1. Axis-aligned slices**\n\nWe chose axis-aligned slices because the keys and queries are learned parameters. Thus, any potential optimization of slice orientations is implicitly captured by the query and key matrices, $W_Q$ and $W_K$. This choice enabled us to avoid introducing extra learnable parameters to ensure a fair comparison with baseline methods, such as vanilla attention and Sinkformer.\n\nTo investigate the effect of the number and type (learnable vs. frozen) of slices, we conducted additional experiments on the Cats vs. Dogs benchmark, varying the number of slices $L$, the inverse temperature parameter $\\tau$, and whether the slices were frozen or learnable. The results are presented below.\n\n| **$\\tau=0.1$**   | $L = 1$ | $L = 8$ | $L = 32$ | $L = 64$ | $L = 128$ |\n| ---------------- | ------- | ------- | -------- | -------- | --------- |\n| **Learnable**    | 74.45%  | 78.56%  | 79.33%   | 78.25%   | 76.15%    |\n| **Frozen**       | 66.61%  | 72.75%  | 78.41%   | 79.39%   | 79.66%    |\n| **Axis-Aligned** | —       | —       | —        | 79.47%   | —         |\n\n| **$\\tau=1.0$**   | $L = 1$ | $L = 8$ | $L = 32$ | $L = 64$ | $L = 128$ |\n| ---------------- | ------- | ------- | -------- | -------- | --------- |\n| **Learnable**    | 74.45%  | 79.07%  | 78.10%   | 77.64%   | 74.34%    |\n| **Frozen**       | 66.61%  | 73.09%  | 77.89%   | 78.88%   | 78.50%    |\n| **Axis-Aligned** | —       | —       | —        | 78.86%   | —         |\n\n| **$\\tau=10$**    | $L = 1$ | $L = 8$ | $L = 32$ | $L = 64$ | $L = 128$ |\n| ---------------- | ------- | ------- | -------- | -------- | --------- |\n| **Learnable**    | 74.45%  | 79.24%  | 78.06%   | 77.07%   | 74.20%    |\n| **Frozen**       | 66.61%  | 73.50%  | 76.79%   | 77.91%   | 78.13%    |\n| **Axis-Aligned** | —       | —       | —        | 77.78%   | —         |\n\nAs shown in the table, both learnable and frozen slicers can operate on the key and query projections. We chose a simpler, axis-aligned slicer for interpretability and fair comparison, avoiding extra parameters. Learnable slicers may help with fewer slices by focusing on informative directions, but their advantage diminishes as the slice count increases. In contrast, frozen slicers avoid added complexity but require many slices to capture distributional structure effectively, especially in high dimensions.\n\n**2. Consistency of optimal $\\tau$ values across tasks:**\n\nThe optimal inverse temperature values, $\\tau$, depend on: 1) the input measures, i.e., the task, and 2) the number of slices. Therefore, it is reasonable to expect that this hyperparameter should be task-dependent.  Smaller $\\tau$ values promote more diffuse attention, while larger values lead to sharper, more focused (i.e., near one-to-one) attention between tokens. Importantly, as shown in our experimental setup, we perform cross-validation only over a very coarse grid, i.e., $\\tau\\in\\{0.1,1,10,100\\}$.\n\n**3. Additional benchmarks:**\n\nWe conducted experiments on TweetEval. The table below summarizes the test performance of several models, including ESPFormer, on this benchmark. ESPFormer demonstrates competitive results compared to the baselines, outperforming both Sinkformer and the standard attention mechanism (denoted by \"Vanilla\"). For a controlled comparison, we replaced the attention module in a standard 6-layer Transformer (as proposed by Vaswani et al., 2017) with alternative components—Sinkformer, DiffAttention, and ESPFormer—while keeping the rest of the architecture unchanged. As illustrated, all variants outperform the vanilla baseline, with ESPFormer and DiffAttention achieving the highest accuracies on this benchmark.\n\n\n| Model    | Vanilla | Sink. | Diff.Att. | ESP. |\n| -------- | ------- | ----- | --------- | ---- |\n| Accuracy | 71.5    | 72.0  | 72.6      | 72.6 |\n\nNote that this evaluation is limited to a single architecture and run due to the rebuttal timeframe. We plan to extend it in the camera-ready version with plug-and-play experiments across diverse attention modules and pretrained language models.\n\n**4.  ESPFormer on ultra-long sequences (N > 10k)**\n\nDue to time and resource constraints, we couldn't evaluate ESPFormer on ultra-long sequences (N > 10k) in this submission. We agree this is important and plan to include full runtime and accuracy benchmarks in future work. Regarding efficiency, by switching from soft-sorting to hard sorting during inference, we can significantly reduce our latency.\n\nOnce again, we sincerely thank the reviewer for their thoughtful comments and insightful questions, which helped us refine the presentation and analysis of our contributions. \n\n**P.S.** Due to the limited rebuttal period, the results presented are based on a single run. We recognize the importance of reporting performance across multiple seeds and will include averaged results in the camera-ready version.", "claims_and_evidence": "The key claims (doubly-stochastic attention, efficiency gains, performance improvements) are supported by experiments across four benchmarks (ModelNet40, IMDb, IWSLT, Cats&Dogs) with consistent metrics (accuracy, BLEU). However, while axis-aligned slices (Θ=I) are justified by parameter efficiency, the paper lacks ablation studies comparing learned vs. fixed slices, leaving uncertainty about their optimality for non-axis-aligned data.", "methods_and_evaluation_criteria": "The proposed methods (sliced transport, soft sorting, τ-controlled aggregation) are theoretically sound and aligned with ESP’s properties (Liu et al., 2025). Axis-aligned slicing avoids extra parameters but may limit expressiveness for non-axis-separable distributions. Evaluation criteria (accuracy, BLEU, runtime) are standard and relevant.", "theoretical_claims": "The paper relies on prior work (Liu et al., 2025) for ESP’s theoretical foundation but does not provide new proofs. Assumptions about ESP’s equivalence to Wasserstein distance and τ’s role in sparsity are empirically validated.", "experimental_designs_or_analyses": "Experiments are generally controlled (same architectures across baselines, multiple runs) but have limitations:\n1.No analysis of τ’s impact beyond qualitative visualization (Fig. 2) or slice count (m).\n2.While the paper compares ESPFormer with Sinkformer and DiffTransformer on IMDB using accuracy, the narrow focus on a single dataset.", "relation_to_broader_scientific_literature": "Key Connections:\nSinkformer (Sander et al., 2022): Directly builds on this work by replacing Sinkhorn iterations with sliced OT for doubly-stochastic attention. The efficiency claim is framed as an improvement over Sinkformer’s O(SN²) complexity.\nSliced OT (Liu et al., 2025): Relies on ESP theory to construct transport plans from 1D slices, adapting it for attention via soft sorting.\nSoft Sorting (Prillo & Eisenschlos, 2020): Critical for differentiable sliced transport.", "essential_references_not_discussed": "No", "other_strengths_and_weaknesses": "Strengths:\n1.Introduces a temperature-controlled sparsity mechanism (via τ), enabling tunable attention patterns without iterative optimization.\n2.Parallelizable O(mN²) complexity addresses a critical bottleneck in OT-based attention, making doubly-stochastic constraints viable for long sequences.\nClarity:\n3.Well-structured with intuitive visualizations (e.g., Fig. 2: Sinkhorn vs. ESP attention patterns).\n4.Appendix details hyperparameters and runtime analysis, aiding reproducibility.\n\nWeaknesses:\n1.Axis-aligned slicing is presented as a default choice without justifying why it suffices (vs. learned/adaptive slices). Prior work on adaptive slicing (e.g., [Nguyen et al., 2023]) could strengthen motivation.\n2.In the task of image classification, the Cats and Dogs dataset and the MNIST dataset are relatively simple, and there is a lack of verification on more complex datasets. \n3.In the Sentiment Analysis experiment, the dataset is single. \n4.No analysis of ESPFormer’s scaling to ultra-long sequences (e.g., N=10k), a key use case for efficient attention.", "other_comments_or_suggestions": "no", "questions_for_authors": "1.Why choose axis-aligned slices (Θ=I) over learned or adaptive slicing (e.g., Nguyen et al., 2023)? Did you test adaptive slicing, and if so, how did it compare in terms of accuracy/efficiency?\n2.How do τ (inverse temperature) and m (slice count) affect performance? Are optimal τ values consistent across tasks?\n3.What the performance when ESPFormer handle ultra-long sequences (N > 10k)? Did you test runtime/accuracy on such data (e.g., text/document classification)?\n4.Could you supplement experiments on more diverse datasets? I feel that the improvements in your experimental section are not significant, and the single dataset used may introduce chance factors due to limited diversity.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: In this study, the authors propose a new variant of Transformer, called ESPFormer, which achieves doubly stochastic attention maps via simple operations like sorting. The authors analyze and interpret the model from the perspective of sliced optimal transport, and compare the model with the existing OT-based Transformer (Sinkformer) on the datasets. Experimental results show the feasibility of the proposed method to some extent. 

AC and reviewers admit the technical novelty of the proposed method. However, some highly relevant references [1, 2], although preprints, should be added as related work, and the discussion and comparison with these works should be added in the paper. In addition, AC believes that the experimental part should be enhanced because the datasets and the tasks considered in the paper are relatively simple. In summary, AC recommends accepting this work weakly. 

[1] Yuan, S., & Xu, H. (2024). Towards Better Multi-head Attention via Channel-wise Sample Permutation. arXiv preprint arXiv:2410.10914.

[2] Yuan, S., & Xu, H. (2023). Sliceformer: Make Multi-head Attention as Simple as Sorting in Discriminative Tasks. arXiv preprint arXiv:2310.17683.
Decision: Accept (poster)