ESPFormer: Doubly-Stochastic Attention with Expected Sliced Transport Plans

We appreciate the constructive feedback. Below are our responses.

1. Axis-aligned slices

We 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.

To 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.

As 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.

2. Consistency of optimal $\tau$ values across tasks:

The 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\}$.

3. Additional benchmarks:

We 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.

Note 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.

4. ESPFormer on ultra-long sequences (N > 10k)

Due 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.

Once again, we sincerely thank the reviewer for their thoughtful comments and insightful questions, which helped us refine the presentation and analysis of our contributions.

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.

Thank you for the constructive feedback. Below, we provide our responses.

1. The Birkhoff Polytope. We thank the reviewer for raising this insightful question. In 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]:

• (Birkhoff-von Neumann Theorem) $B_N$ is a convex polytope whose extreme points, or corners, are the $N!$ permutation matrices.

1. 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$.
2. 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.
3. In this figure, 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.

[1] Birkhoff, G., 1946. [2] Liu X, et al. 2024. [3] Mahey G, et al. 2023.

2. What $\tau$ to choose? The 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.

3. Source code availability. We 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.

4. Exact/Approximate DSMs. Soft 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.

To 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:

Experiments on 25% and 100% data are ongoing and will be included in the camera-ready version.

5. Updating Figure 3. We 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.

6. On the complexity of ESP, Sinkhorn and Vanilla Attention. We 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.

Thank you again for your consideration.

We thank the reviewer for their constructive feedback. Below are our responses.

1. Reliance on Soft sorting:

Soft 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. This improves runtime, reducing complexity from $O(N^2)$ to $O(N \log N)$, and reduces our dependency on soft-sorting. See Figure 3 for updated runtimes. The table below summarizes results on Cats vs. Dogs across training fractions.

Experiments for the 25% and 100% training splits are currently in progress and will be included in the camera-ready version of the paper.

2. Additional experiments:

We 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.

Note 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.

3. Axis-aligned slices:

We 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.

To 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.

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.

Once again, we thank the reviewer for their time and consideration,.

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.