---
title: SortCut Algorithm for Efficient Sparse Attention
url: https://www.emergentmind.com/topics/sortcut-algorithm
type: topic
---

# SortCut Algorithm for Efficient Sparse Attention

SortCut is an algorithmic enhancement to the Sparse Sinkhorn Attention mechanism, designed to enable efficient and scalable self-attention for long sequences by dynamically truncating attention to the most salient regions. Its core principle is the explicit selection, via learned sorting and hard truncation, of a fixed budget of sequence blocks before applying local attention. This process yields near-linear time and space complexity in sequence length, while maintaining empirical performance comparable to full attention on a variety of natural language and sequence modeling tasks [2002.11296].

## 1. Conceptual Motivation and Design Principles

Standard Transformer self-attention incurs quadratic memory and computational costs in the sequence length $\ell$. Local attention reduces this but loses the ability to aggregate global context. Sparse Sinkhorn Attention addresses this by learning a permutation of sequence blocks, allowing local attention to span non-contiguous parts of the original sequence. SortCut introduces an additional compression step: after learning block importance through permutation, it truncates the permuted sequence by retaining only the highest-scoring $n$ blocks. This approach exploits input-dependent structure, focusing memory and compute resources on the most relevant content per example.

SortCut is especially effective when applied in encoder architectures, where a single learned sort suffices per layer. For decoder/autoregressive applications, the sorting must be performed at each timestep with special causal constraints, increasing computational overhead.

## 2. Formal Definition and Mathematical Specification

Consider sequence embeddings $X \in \mathbb{R}^{\ell \times d}$, partitioned into $N_B = \ell / B$ non-overlapping blocks of size $B$. A block-pooling operation $\psi_P(X) = B(X) \in \mathbb{R}^{N_B \times d}$ (typically mean or sum pooling) aggregates each block's embeddings.

A small block-sorting network $P(\cdot): \mathbb{R}^d \rightarrow \mathbb{R}^{N_B}$ assigns each block a vector of sorting logits $R_{\text{logits}} \in \mathbb{R}^{N_B \times N_B}$. These logits are perturbed (optionally) with Gumbel noise, scaled by a temperature parameter $\tau$, then normalized using $N_{\text{iter}}$ Sinkhorn iterations to yield a soft permutation matrix $R$.

Block-wise key and value representations $K_B$, $V_B$ are linearly projected and block-pooled, then permuted:

\[
K_S = R K_B, \qquad V_S = R V_B
\]

SortCut truncates the sorted blocks, retaining only the first $n$:

\[
K_{\text{cut}} = K_S[1:n, :], \qquad V_{\text{cut}} = V_S[1:n, :]
\]

Attention is computed using queries $Q$ against the unfolded truncated blocks:

\[
Y = \mathrm{Softmax}\left(Q [\psi_S(K)]_{:n}^\top \right) \ [\psi_S(V)]_{:n}
\]

where $\psi_S(X) = U(R \cdot B(X))$ denotes block-sorting followed by unpooling.

## 3. Algorithmic Workflow and Pseudocode

The SortCut variant of Sparse Sinkhorn Attention involves the following pipeline for each attention head:

1. Compute $Q$, $K$, $V$ by projecting the input $X$.
2. Pool $K$ and $V$ into blocks ($K_B$, $V_B$).
3. Obtain sorting logits via the feed-forward network $P(\cdot)$.
4. Optionally add Gumbel noise, normalize via Sinkhorn iterations to soft permutation matrix $R$.
5. Apply $R$ to block-pooled $K_B$ and $V_B$ for permutation.
6. Truncate to top-$n$ blocks ($K_{\text{cut}}$, $V_{\text{cut}}$).
7. Expand truncated blocks back to token-level representations.
8. Compute scaled dot-product attention using $Q$ and the expanded $K_{\text{cut}}$, yielding outputs $Y$.

Block pooling can be adapted (sum, mean) as necessary. For multi-head attention, each head maintains independent sorting networks and budgets.

| Step                | Input(s)         | Output                |
|---------------------|------------------|-----------------------|
| Block pooling       | $K, V$           | $K_B, V_B$            |
| Sorting logits      | $K_B$            | $R_{\text{logits}}$    |
| Sinkhorn norm.      | $R_{\text{logits}}$, $\tau$ | $R$            |
| Permutation         | $K_B, V_B, R$    | $K_S, V_S$             |
| Truncation          | $K_S, V_S, n$    | $K_{\text{cut}}, V_{\text{cut}}$ |

This structured pipeline underpins SortCut’s efficiency and flexibility.

## 4. Computational Complexity and Efficiency

For sequence length $\ell$, block size $B$, number of blocks $N_B = \ell / B$, and SortCut budget $n$:

- **Full attention:** $O(\ell^2 d)$ time, $O(\ell^2)$ memory
- **Block-local attention:** $O(B^2 d + (\ell/B)^2 d)$ time, $O(B^2 + (\ell/B)^2)$ memory
- **Sinkhorn block-sort attention:** Adds $O(N_{\text{iter}} N_B^2)$ Sinkhorn cost, no truncation
- **SortCut (with truncation):** $O(\ell n d + N_{\text{iter}} N_B^2)$ time, $O(\ell n + N_B^2)$ memory

Since $n \ll N_B \ll \ell$ typically, SortCut achieves near-linear time and space in $\ell$, aside from the soft permutation (Sinkhorn) cost—amortized over training. Unlike masking-based sparse architectures, SortCut requires only standard matrix operations, simplifying implementation.

## 5. Hyperparameterization and Influence

Key SortCut hyperparameters include:

- **Block size $B$:** Governs granularity; larger $B$ yields fewer blocks and cheaper Sinkhorn sorting, at potential loss of block-level detail.
- **Sinkhorn temperature $\tau$ and iteration count $N_{\text{iter}}$:** Lower $\tau$ and higher $N_{\text{iter}}$ make the learned permutation more discrete. However, excessively sharp (hard) sorting impairs gradient flow. Empirical ablative analysis shows $\tau \in [0.5, 0.75]$ and $5 \leq N_{\text{iter}} \leq 10$ are typically optimal; $N_{\text{iter}} = 0$ (no Sinkhorn) severely degrades performance.
- **Truncation budget $n$:** The number of blocks retained post-sort. Lower $n$ saves memory but potentially omits context; in practice, even $n=8$ (on long NLP sequences) can nearly match full attention using $1/32$ of the memory.
- **Block sorting network $P(\cdot)$:** Its architecture and capacity affect permutation quality.

Ablations confirm that small $n$ can suffice and that the balance between Sinkhorn softness and truncation aggressiveness is crucial to optimal trade-offs.

## 6. Integration with Meta-Sorting Networks and Causal Sinkhorn

The meta-sorting network $P(\cdot)$ emits per-block logits guiding the permutation. In encoder architectures, a single sort suffices per layer, maximizing efficiency. For causal (decoder) attention, the block pooling and Sinkhorn normalization must enforce causality:

- Block pooling up to position $i$ for block $i$, ensuring no future context access.
- Masked Sinkhorn normalization prohibits permutation entries for indices $j > i$.

Causal SortCut thus requires recomputation at every decoding step, making it better suited to encoder use cases.

## 7. Empirical Results and Comparative Analysis

On standard NLP and sequence tasks (IMDb, SST, SNLI, MultiNLI, up to 2k tokens), SortCut encoders with $n=8$ blocks demonstrate accuracy within $1$–$2\%$ of full attention while reducing memory usage by $30\times$–$50\times$.

Representative experimental results:

| Model                         | IMDb Accuracy (%) | SNLI Accuracy (%)  | Memory Use             |
|-------------------------------|------------------|--------------------|------------------------|
| Vanilla Transformer           | 85.12            | 78.87              | Baseline (O($\ell^2$)) |
| SortCut ($B=128, n=8$)        | 84.32            | –                  | $\sim1/32$             |
| SortCut ($B=128, n=32$)       | 84.43            | –                  | $\sim1/8$              |
| SortCut ($n=16$)              | –                | 80.30              | Reduced                |

Examinations reveal that Sinkhorn is essential for performance; setting $N_{\text{iter}} = 0$ collapses accuracy. Optimal results arise from a soft—but not too hard—sorting regime and moderate block budgets.

Empirically, SortCut sometimes outperforms full attention models on NLI tasks, suggesting that focused content selection can act as a regularizer. The practical utility of hard truncation and learned permutation is consistently confirmed by ablations [2002.11296].

---

In summary, SortCut generalizes Sparse Sinkhorn Attention with post-sorting truncation, dynamically selecting the most pertinent subsequence blocks for attention computation. This results in a simple, efficient, and highly scalable mechanism—well-suited for long-sequence modeling without requiring specialized hardware or custom kernels, achieving substantial savings in computational resources while closely matching standard attention's empirical performance.

Source: https://www.emergentmind.com/topics/sortcut-algorithm