PaperlayerView exact revision on arXiv

Source-linked AI summary

Sparse Sinkhorn Attention

Yi Tay, Dara Bahri, Liu Yang, Donald Metzler, Da-Cheng Juan

arXiv:2002.11296v1cs.LGcs.CL

TL;DR

Dense attention has quadratic memory costs and may be suboptimal for long, noisy sequences. Sparse Sinkhorn Attention learns differentiable permutations, enabling local windows to produce quasi-global attention, with causal balancing and SortCut variants. Across large-scale modeling and classification tasks, it remains competitive with vanilla Transformers and outperforms sparse Transformer alternatives.

  • Problem

    Dense attention incurs quadratic memory costs and may suffer when sequence length is large and noisy, motivating sparse attention outputs.

  • Method

    Sparse Sinkhorn Attention uses a meta sorting network with differentiable Sinkhorn balancing to reorder sequences before local attention, with causal and truncation variants.

  • Results

    Across large-scale generative modeling and classification tasks, Sinkhorn Transformers outperform or remain competitive with vanilla and sparse Transformer models.

  • Takeaways & Limitations

    Learned neural sorting can provide memory-efficient sparse attention while retaining competitive performance across multiple applications.

  • Takeaways & Limitations

    SORTCUT is limited to Transformer encoding unless autoregressive decoding recomputes self-attention at every time step, while the mixture mode retains quadratic complexity.

Abstract

from arXiv · show

We propose Sparse Sinkhorn Attention, a new efficient and sparse method for learning to attend. Our method is based on differentiable sorting of internal representations. Concretely, we introduce a meta sorting network that learns to generate latent permutations over sequences. Given sorted sequences, we are then able to compute quasi-global attention with only local windows, improving the memory efficiency of the attention module. To this end, we propose new algorithmic innovations such as Causal Sinkhorn Balancing and SortCut, a dynamic sequence truncation method for tailoring Sinkhorn Attention for encoding and/or decoding purposes. Via extensive experiments on algorithmic seq2seq sorting, language modeling, pixel-wise image generation, document classification and natural language inference, we demonstrate that our memory efficient Sinkhorn Attention method is competitive with vanilla attention and consistently outperforms recently proposed efficient Transformer models such as Sparse Transformers.

1. Introduction

Sparse Sinkhorn Attention reduces attention memory by dynamically sorting representations and then applying local attention to the reordered sequences. The method adds causal balancing and dynamic truncation variants, and experiments report competitiveness with vanilla attention while outperforming local and Sparse Transformer baselines.

  • Motivation: The work targets dense attention’s quadratic memory cost and potential weakness on long, noisy sequences by learning sparse attentive outputs.Sparse outputs are presented as an inductive bias resembling hard attention.
  • Method: Sparse Sinkhorn Attention learns sparse attention outputs by dynamically sorting internal representations with differentiable Sinkhorn balancing.A meta sorting network generates block-wise permutations, which are passed to a local attention module.
  • Method: Sorted sequences let local attention access beyond the default neighborhood without additional attention computation costs.The reordered representation creates a quasi-global effect using local context windows.
  • Efficiency: O(B^2 + N^2_B) replaces O(ℓ^2) attention memory, and SORTCUT further reduces complexity to O(ℓN_k).Here, B is the block size and N_k is a user-defined budget with N_k <<< ℓ.
  • Extensions: Causal Sinkhorn balancing supports autoregressive decoding, while SORTCUT dynamically truncates sequences for more efficient encoding.The paper also presents a mixture of Sparse Sinkhorn and vanilla attention for further performance improvements.
  • Experiments: Across large-scale generative modeling tasks, Sinkhorn Transformers match or outperform vanilla Transformers and consistently outperform Sparse and Local Attention Transformers.The evaluation includes language modeling, pixel-wise image generation, and document classification.

2. Related Work

Prior efficient-attention work reduces computation or sparsifies outputs, but approaches trade global context, portability, or memory efficiency in different ways. Sparse Sinkhorn Attention builds on differentiable neural sorting to learn sparse attention structure.

  • Local attention limits tokens to fixed windows, restricting global receptive fields and long-term dependency modeling.
  • Sparse Transformer factorizes attention into local and strided operations, assigning different heads to different sparse patterns.
  • Reformer uses locality-sensitive hashing, Transformer-XL caches hidden states recurrently, and Star Transformer converts dense graphs into star-shaped topologies.
  • Sparse attention outputs can focus on limited items, providing a denoising inductive bias for long or noisy sequences and potentially improving inference speed.
  • Sparse Sinkhorn Attention applies neural sorting to learn differentiable rankings that sparsify attention and improve efficiency.
  • Sparse Attention may require specialized GPU kernels, reducing portability to platforms such as TPU pods.

3. Sparse Sinkhorn Attention

Sparse Sinkhorn Attention neural-sorts blocks before local attention, allowing tokens to attend across blocks that may be far apart in the original sequence. This preserves local computation while expanding the effective receptive field.

  • Block-based local attention normally restricts tokens to their own block, limiting the global receptive field and long-term dependency modeling.
  • A Meta Sorting Network learns to sort sequences so local windows can support efficient quasi-global attention.
  • Sparse Sinkhorn Attention sorts blocks and neighborhoods so each token attends within a newly sorted block that may be distant in the original sequence.

3.1. Learning to Sort

The learning-to-sort component partitions representations into blocks, predicts relaxed block permutations, and applies Sinkhorn normalization to obtain differentiable sorting matrices. The learned matrix then reorders block representations end to end.

  • SortNet accepts ℓ vectors of dimension d, partitions them into blocks, and learns relaxed permutation matrices end to end.
  • Blockwise pooling maps R^ℓ×d to R^N_B×d, producing one representation for each of N_B blocks.
  • Each block learns a projection to N_B other blocks, effectively learning the position to which it should be shifted or permuted.
  • A doubly stochastic matrix is nonnegative with every row and column summing to 1, serving here as a relaxed permutation matrix.
  • Sinkhorn normalization consecutively normalizes the rows and columns of the sorting matrix, with N_k iterations controlled by a user-defined hyperparameter.
  • The normalization converges to a doubly stochastic limit when the nonnegative matrix has support, and its gradients enable end-to-end training.
  • Multiplying the doubly stochastic matrix by blocked inputs sorts the sequence, with block reshaping operators converting between token-wise and block-wise forms.

3.2. Sparse Sinkhorn Attention

Sparse Sinkhorn Attention computes attention on block-sorted sequences, combining sorted-key/value interactions with local attention. Gumbel noise makes the sorting operator differentiable, while a mixture variant adds vanilla attention but restores quadratic complexity.

  • The attention mechanism is reformulated to operate on block-sorted sequences rather than the original ordering.
  • The neural sorting function reorders keys and values so local computation can select information beyond the default neighborhood.
  • Standard local attention is added as a secondary term, with attention weights computed only for positions sharing a local block.
  • Gumbel Noise: Gumbel noise and a temperature parameter are injected into sorting; lowering the temperature makes the operator closer to a discrete permutation matrix.
  • Multi-Head Attention: The multi-head implementation learns a separate sorting network for each attention head rather than sharing one permutation matrix.
  • Mixture Model: The mixture variant combines Sinkhorn and vanilla dot-product attention, but it regresses to vanilla self-attention’s quadratic complexity.

3.3. Causal Sparse Sinkhorn Attention

Causal Sparse Sinkhorn Attention preserves autoregressive causality while learning content-based permutations from past context. It masks future information during balancing and excludes blocks sorted into earlier positions.

  • 3.3. Causal Sparse Sinkhorn Attention: Blocks sorted into an earlier position p < i are masked out so current time steps cannot access future time steps.The sorting network therefore also determines which blocks are masked.
  • 3.3. Causal Sparse Sinkhorn Attention: Permutation matrices are generated from cumulative sums of embeddings rather than global sequence information, preserving causal conditioning.The cumulative sum conditions sorting on preceding context up to the current block.
  • 3.3. Causal Sparse Sinkhorn Attention: Causal Sinkhorn Balancing masks future positions during iterative normalization because standard Sinkhorn balancing requires future tokens.This adapts Sinkhorn balancing for causal self-attention.
  • 3.3. Causal Sparse Sinkhorn Attention: The masking function M(.) is used within the causal normalization procedure.

3.4. SORTCUT Sinkhorn Attention

SORTCUT extends Sparse Sinkhorn Attention with post-sorting truncation of input blocks. It dynamically keeps a budgeted subset, improving encoding efficiency while remaining distinct from autoregressive decoding use.

  • 3.4. SORTCUT Sinkhorn Attention: SORTCUT performs a hard top-k operation on sorted input blocks within the computational graph.Unlike attention that merely re-weights tokens, it explicitly truncates the sequence.
  • 3.4. SORTCUT Sinkhorn Attention: SORTCUT may be used on the Transformer encoder, but autoregressive decoding requires recomputing self-attention at every time step.This is the method's stated operational boundary.
  • 3.4. SORTCUT Sinkhorn Attention: Figure 2 presents the proposed SortCut Encoding Scheme.

4. Complexity Analysis

Sparse Sinkhorn Attention reduces the quadratic memory burden of vanilla self-attention by combining block sorting with local computation. SORTCUT further reduces encoder complexity to linear time under a small budget.

  • 4. Complexity Analysis: O(ℓ^2) is the vanilla Transformer self-attention memory complexity for input length ℓ.
  • 4. Complexity Analysis: The proposed Sinkhorn model reduces memory complexity to O(B^2 + (ℓ/N_B)^2), equivalent to local-attention memory complexity.Here B = ℓ/N_B.
  • 4. Complexity Analysis: SORTCUT encoder memory complexity is O(ℓN_k + (N/B)^2), where N_k is a user-defined budget hyperparameter.
  • 4. Complexity Analysis: O(ℓ) complexity is attainable for the SORTCUT encoder when the budget is much smaller than the sequence length.The passage states this condition as ℓ/B << ℓ.

5. Experiments

Across algorithmic sorting, language modeling, image generation, document classification, and natural language inference, Sinkhorn variants generally outperform local and sparse attention baselines while retaining memory-efficiency advantages. Results vary by task and configuration, with mixture models strongest in some language-modeling settings and SORTCUT matching vanilla performance on sentiment classification.

  • 5.1. Algorithmic Tasks: Sinkhorn Transformers outperform all other Transformer variants, including vanilla attention, on algorithmic sorting.The task trains on length 256 and evaluates generalization to length 512 using exact match and edit distance.
  • 5.2. Language Modeling: 2−3 perplexity points is the Sinkhorn gain over local attention across block-size settings in subword language modeling.On the base setting, Sinkhorn also outperforms vanilla Transformer at B = 32 and B = 64.
  • 5.2. Language Modeling: The Sinkhorn Mixture model outperforms all models in subword language modeling, while big Sinkhorn does not outperform vanilla Transformers.The big model remains reasonably competitive despite being more memory efficient.
  • 5.2. Language Modeling: The Sinkhorn Mixture model achieves the best performance for both base and big character-level language-modeling parameterizations.Ordinary Sinkhorn outperforms local and Sparse Transformer, but vanilla full attention outperforms efficient methods in this experiment.
  • 5.2. Language Modeling: Sinkhorn Transformer remains competitive with High Budget MoE and Evolved Transformer models but does not outperform the 5-billion-parameter Mesh Tensorflow setup.The comparison uses per-word perplexity on LM1B.
  • 5.3. Pixel-wise Image Generation: Sinkhorn Transformer outperforms all baselines on pixel-wise CIFAR-10 image generation, including Sparse Transformer with the same local window.The Sinkhorn Mixture performs worse than ordinary Sinkhorn Transformer in this task.
  • 5.4. Text Classification: SORTCUT achieves performance similar to standard Sinkhorn and vanilla Transformer across all sentiment-analysis settings.The comparison covers word- and character-level IMDb and SST classification.
  • 5.4. Text Classification: Both Sinkhorn and SORTCUT Sinkhorn outperform vanilla Transformer on SNLI and MNLI.

6. Analysis

The analysis examines how Sinkhorn sorting choices affect LM1B perplexity. Sinkhorn normalization and moderate sorting settings are important, while excessive iterations and hard sorting degrade performance.

  • Ablations: Sinkhorn normalization is essential: removing it causes substantial performance degradation and the worst ablation result.The ablation sets N_k = 0, equivalent to omitting Sinkhorn normalization on R.
  • Ablations: A linear layer is the best sorting-network model among the tested configurations.The analysis also finds that sharing key and value weights more often hurts performance.
  • Temperature: τ = 0.75 is optimal on LM1B, with higher temperatures favoring soft sorting over more discrete sorting.The reported trend indicates that soft sorting performs better than hard sorting when other variables are fixed.
  • Sorting iterations: 5 −10 sorting iterations appear optimal, whereas no sorting performs extremely badly and more than 20 iterations hurt perplexity.A small number of iterations is sufficient for good performance.

7. Conclusion

The paper concludes that Sparse Sinkhorn Attention provides efficient, sparse attention through neural sorting of internal representations. Across large-scale generative modeling and classification tasks, the Sinkhorn Transformer is memory efficient while outperforming or matching vanilla and sparse Transformer models.

  • 7. Conclusion: Sparse Sinkhorn Attention applies neural sorting of internal representations to efficient and sparse attention computation.The conclusion positions this mechanism as useful across multiple large-scale generative modeling and classification tasks.
  • 7. Conclusion: The Sinkhorn Transformer outperforms or remains competitive with vanilla Transformer and sparse Transformer models while being memory efficient.This conclusion is reported across a multitude of applications.
Loading 2002.11296v1…