PaperlayerView exact revision on arXiv

Source-linked AI summary

Sinkformers: Transformers with Doubly Stochastic Attention

Michael E. Sander, Pierre Ablin, Mathieu Blondel, Gabriel Peyré

arXiv:2110.11773v2cs.LGstat.ML

TL;DR

Transformers use row-wise SoftMax attention, while the paper asks whether doubly stochastic normalization provides a better-founded alternative. It replaces SoftMax with Sinkhorn to create Sinkformers, showing theoretical connections to Wasserstein flows and heat diffusion and improved accuracy across several tasks. The theoretical interpretation relies on stated assumptions.

  • Problem

    Transformer attention is row-wise stochastic, motivating investigation of whether doubly stochastic attention offers a useful alternative.

  • Method

    The paper replaces SoftMax with iterative Sinkhorn row-and-column normalization, producing a differentiable Sinkformer without changing the network’s parameters.

  • Results

    Sinkformers improve accuracy across vision and natural language tasks, while their iterations admit Wasserstein gradient-flow and heat-equation interpretations under the paper’s settings.

  • Takeaways & Limitations

    Doubly stochastic attention provides a practical Transformer replacement and a framework for analyzing self-attention through measure-based dynamics.

  • Takeaways & Limitations

    The Wasserstein gradient-flow result is established under a symmetry assumption and a more restrictive parameter condition is discussed for a particular case.

Abstract

from arXiv · show

Attention based models such as Transformers involve pairwise interactions between data points, modeled with a learnable attention matrix. Importantly, this attention matrix is normalized with the SoftMax operator, which makes it row-wise stochastic. In this paper, we propose instead to use Sinkhorn's algorithm to make attention matrices doubly stochastic. We call the resulting model a Sinkformer. We show that the row-wise stochastic attention matrices in classical Transformers get close to doubly stochastic matrices as the number of epochs increases, justifying the use of Sinkhorn normalization as an informative prior. On the theoretical side, we show that, unlike the SoftMax operation, this normalization makes it possible to understand the iterations of self-attention modules as a discretized gradient-flow for the Wasserstein metric. We also show in the infinite number of samples limit that, when rescaling both attention matrices and depth, Sinkformers operate a heat diffusion. On the experimental side, we show that Sinkformers enhance model accuracy in vision and natural language processing tasks. In particular, on 3D shapes classification, Sinkformers lead to a significant improvement.

1 Introduction

Sinkformers replace SoftMax’s row-wise stochastic attention with Sinkhorn-normalized doubly stochastic attention. The paper motivates this change empirically and develops theoretical and experimental evidence for its benefits.

  • Motivation and method: SoftMax normalizes attention matrices row-wise, whereas Sinkhorn’s algorithm successively normalizes rows and columns to produce a doubly stochastic matrix.Doubly stochastic matrices have rows and columns that both sum to 1.
  • Motivation and method: Sinkformers replace SoftMax with Sinkhorn normalization, retaining the Transformer architecture while using more distributed interactions among points.The modification is presented as a prior favoring attention matrices in which all points participate with varying interaction intensities.
  • Empirical motivation: Training experiments suggest that row-wise stochastic attention matrices in classical Transformers move closer to doubly stochastic matrices as learning proceeds.This observation motivates using Sinkhorn normalization as an informative prior.
  • Theoretical contributions: Sinkformers admit a Wasserstein gradient-flow interpretation under the paper’s stated theoretical setting, unlike classical Transformers with SoftMax.The paper also proves that, as the number of particles tends to infinity, Sinkformer iterations converge to the heat equation.
  • Experimental results: Sinkformers improve reported performance across 3D shape classification, sentiment analysis, machine translation, and image classification tasks.The paper highlights a significant accuracy gain over Transformers on ModelNet40 3D shape classification.

2 Background and related work

The background situates Transformers as attention-based models and connects their analysis to kernel normalization, continuous-depth dynamics, and neural networks acting on measures. Prior work provides related Sinkhorn, diffusion, and operator perspectives.

  • Transformers: Transformers use encoder-decoder architectures whose encoders iterate self-attention residual blocks followed by pointwise feed-forward networks.Variants include Vision, Set, and Point Cloud Transformers, while neural machine translation may also use cross-attention in the decoder.
  • Sinkhorn and attention: Prior work has used Sinkhorn-generated doubly stochastic matrices for differentiable sorting and rearranging of input sequences.That related approach treats the matrix as a relaxed permutation matrix rather than as the normalization proposed here.
  • Kernel normalization: Bi-normalization of kernels has been studied for uniform and weighted measures, with applications to smoothing operators and signal processing.Related studies report Langevin-diffusion and symmetric-operator interpretations, with faster convergence than SoftMax normalization.
  • Infinite depth limit: Residual neural networks can be analyzed in the infinitesimal-step-size regime as discretized Euler schemes for ordinary differential equations.The paper applies this continuous-depth viewpoint to residual attention layers.
  • Neural networks on measures: Self-attention on unordered point sets can be modeled as an operation on discrete probability measures, enabling analysis of varying cardinalities and large-sample limits.A collection of points is represented as an empirical measure, and the limit n →∞ supports mean-field analysis.

3 Sinkformers

Sinkformers replace SoftMax attention normalization with Sinkhorn normalization, producing approximately doubly stochastic attention matrices. The design is motivated by observed training dynamics and retains practical compatibility with Transformer architectures.

  • Sinkformers: Sinkformers replace the SoftMax operator in attention modules with Sinkhorn’s algorithm.The resulting model is a drop-in Transformer modification that can be implemented in existing architectures.
  • Motivation: Classical Transformer attention matrices become increasingly close to doubly stochastic during training across three models and three learning tasks.The observation was made by tracking column sums during learning.
  • Sinkhorn normalization: Sinkhorn alternates row-wise and columnwise normalizations until the resulting matrix is doubly stochastic.The limit K∞ has rows and columns that both sum to 1.
  • Sinkformers: SoftMax is exactly the first Sinkhorn iteration, so finite iteration counts interpolate between the Transformer and Sinkformer.This provides a parameterized transition between row-wise normalization and approximate double stochasticity.
  • Implementation: Sinkformers preserve differentiability, GPU execution, and the overall Transformer structure while adding iteration-dependent normalization cost.Using l Sinkhorn iterations takes l times longer than SoftMax, although Sinkhorn is not the main computational bottleneck in practice.
  • Cost invariance: Sinkhorn normalization is invariant to adding separate row and column terms to the attention cost.Thus, the cost can be rewritten using squared L2 distances without changing K∞.

4 Attention and gradient flows

The paper analyzes attention as interacting particle dynamics and connects Sinkformer updates to Wasserstein gradient flows under symmetry assumptions. This framework distinguishes Sinkformer dynamics from SoftMax dynamics and supports further continuous-limit interpretations.

  • Measure-valued dynamics: Self-attention and Sinkformer blocks can be represented as dynamics of probability measures whose particles interact through a measure-dependent vector field.For discrete measures, the continuous operators recover the corresponding Transformer and Sinkformer updates.
  • Continuous limits: The analysis studies continuous limits by omitting feed-forward networks, tying layer weights, and considering infinite depth through a neural ODE.This setup is used to determine the PDEs generated by different attention normalizations.
  • Wasserstein gradient flows: Under the stated symmetry assumption, Sinkformer dynamics correspond to a Wasserstein gradient flow for an energy functional.The associated evolution is a steepest descent in Wasserstein space that minimizes a global energy.
  • PDE interpretation: Different attention kernels generate different PDEs, and the Sinkformer formulation provides a variational perspective that supports interpreting Sinkformers as argmin layers.The paper presents this as a framework for analyzing theoretical properties of attention mechanisms.
  • Wasserstein gradient flows: The SoftMax-normalized attention dynamics do not correspond to a Wasserstein gradient flow.The paper attributes this distinction to the lack of symmetry in the relevant vector field.
  • Skip connections: With skip connections, the dynamics can exhibit aggregation or diffusion rather than only rank collapse.The paper contrasts this richer behavior with prior analysis without skip connections, where Transformer outputs converge to a rank-1 matrix with depth.

5 Attention and diffusion

The paper analyzes Sinkformers in continuous and asymptotic regimes, showing that their dynamics can yield heat diffusion, while Transformer dynamics remain nonlinear and nonlocal.

  • Sinkformer dynamics: Sinkformer dynamics use the attention kernel defined by Sinkhorn normalization and rescale the query-key interaction by the bandwidth parameter.The continuous mapping is studied as the bandwidth ε approaches zero.
  • Mean-field limit: In the mean-field limit, the Sinkformer’s rescaled dynamics recover the heat equation.The result is established under compact-support, density-smoothness, and positive-semidefiniteness assumptions.
  • Comparison with Transformers: The theoretical analysis compares Sinkformer and SoftMax attention through limiting PDEs derived using asymptotic expansions.The SoftMax analysis uses a modified cost to enable a Laplace expansion.
  • Mean-field limit: The Transformer’s corresponding PDE is nonlinear and nonlocal in the density.Its evolution at a location depends on density values at transformed locations.
  • Comparison with Transformers: The resulting contrast is linear and local behavior for Sinkformers versus nonlinear and nonlocal behavior for Transformers.This distinction persists without assuming that WQ and WK have the stated invertibility properties.

6 Experiments

Experiments evaluate Sinkformers across 3D shape classification, sentiment analysis, translation, and vision tasks, generally reporting improved accuracy with modest computational overhead.

  • ModelNet 40 classification: Sinkhorn iterations increase accuracy in the Set Sinkformer ModelNet 40 experiment.The study varies the number of iterations while retaining row-wise stochastic attention through odd iteration counts.
  • ModelNet 40 classification: On Point Cloud Transformers, the Sinkformer achieves slightly better median accuracy, while the Transformer narrowly achieves the best test accuracy.The comparison uses four runs for each model.
  • Natural language processing: The Sinkformer improves accuracy on IMDb sentiment analysis, with Sinkhorn converging perfectly in three iterations.The resulting attention matrices are doubly stochastic, and the added computation is small.
  • Natural language processing: On IWSLT’14 German-to-English translation, median BLEU is 34.68 for the Transformer and 34.73 for the Sinkformer.The Sinkformer modifies only the encoder, and training time per epoch is nearly unchanged.
  • Vision Transformers: On cats-and-dogs classification, median test accuracy is 79.0% for the Transformer and 79.5% for the Sinkformer.Maximum test accuracy is 80.0% versus 80.5%, with three Sinkhorn iterations adding about five seconds per epoch.
  • Vision Transformers: For a minimal MNIST self-attention model, the Sinkformer outperforms the Transformer when patch size becomes smaller than the full image.With one 28×28 patch, both models are equivalent because the attention matrix has size 1.

Conclusion

The paper concludes that replacing SoftMax with Sinkhorn yields a Transformer variant with doubly stochastic attention and improved accuracy across several application domains.

  • Conclusion: The Sinkformer replaces row-wise stochastic SoftMax attention with doubly stochastic Sinkhorn attention.The replacement does not change the network’s parameters or overall architecture.
  • Conclusion: Sinkformers report better accuracy across 3D shape classification, sentiment analysis, neural machine translation, and image classification.The conclusion summarizes improvements across multiple modalities and tasks.

A.1 Invariance to the cost function - Proof of Proposition 1

The proof establishes that adding separable row and column potentials to the cost leaves Sinkhorn normalization unchanged. The section also develops gradient and PDE expressions associated with the relevant normalized kernels.

  • Invariance to the cost function: Adding row and column potentials to the cost leaves the argmin over doubly stochastic matrices unchanged, so Sinkhorn(C) = Sinkhorn(C̃).The proof uses the variational formulation of Sinkhorn normalization.
  • Gradient expressions: Under symmetry and regularity assumptions, the section derives gradient expressions for the normalized kernels using differentiation under the integral.The derivation includes the symmetric-kernel assumption and Wasserstein-gradient formulations.
  • Auxiliary result: A contradiction argument shows that a kernel satisfying the stated relation must be constant under symmetry.The proof specializes the measure to a Dirac mass to establish equality across arguments.
  • PDE derivation: The PDE analysis applies changes of variables, Laplace expansion, and Taylor expansion to obtain the expected limiting form.The general PDE form is stated after handling the transformed variables and asymptotic expansion.

B Implementation details

The implementation uses log-domain Sinkhorn iterations for numerical stability and computes normalization terms with log-sum-exp. The experiments specify separate sampling, optimization, and training schedules for Set Transformers and Point Cloud Transformers.

  • Sinkhorn implementation: Sinkhorn’s algorithm is implemented in the log domain, starting from g0 = 0_n and iterating toward scaling vectors that produce a doubly stochastic limit.The limiting matrix is represented as diag(e^f∞)K0diag(e^g∞).
  • Numerical stability: Log-sum-exp is used to compute log(K e^g_l) and log(K^T e^f_l), enabling fast and accurate computations.These operations stabilize the matrix-vector calculations in the log-domain iterations.
  • Set Transformers: Set Transformer experiments sample 5000 points per ModelNet example and train a two-ISAB encoder with SAB and PMA decoder for 300 epochs.Training uses batch-size 64, Adam, and an initial learning rate of 10^-3.
  • Point Cloud Transformers: Point Cloud Transformer experiments sample 1024 points per ModelNet example and train for 300 epochs with batch-size 32 and SGD.The initial learning rate is 10^-4 and is reduced by a factor of 10 after 250 epochs.

C.2 Sentiment Analysis

The experiments describe training configurations for IMDb sentiment analysis, IWSLT’14 fine-tuning, cats-and-dogs classification, and patch-size comparisons. These settings vary model depth, heads, optimizer, learning rate, and training schedule.

  • Sentiment analysis: IMDb experiments train a depth-6, 8-head Transformer from scratch for 15 epochs using Adam and batch-size 32.The initial learning rate is 10^-4 and is reduced by a factor of 10 after 12 epochs.
  • Machine translation: IWSLT’14 Sinkformer fine-tuning uses the fairseq Transformer training command with the original learning rate divided by 10.The passage specifies the learning-rate adjustment but not an outcome metric.
  • Image classification: Cats-and-dogs classification uses a depth-6, 8-head model with patch-size 16, Adam, batch-size 64, and 300 training epochs.The initial learning rate is 5 × 10^-5 and is divided by 10 after 250 epochs.
  • Patch-size analysis: Patch-size experiments use a depth-1, single-head model without non-linearity and train for 45 epochs across different patch sizes.The initial learning rates are 1 × 10^-3 for the Transformer and 2 × 10^-3 for the Sinkformer.
Loading 2110.11773v2…