Sinkformers: Transformers with Doubly Stochastic Attention

Introduction

The Transformer (Vaswani et al.,, 2017), an architecture that relies entirely on attention mechanisms (Bahdanau et al.,, 2014), has achieved state of the art empirical success in natural language processing (NLP) (Brown et al.,, 2020; Radford et al.,, 2019; Wolf et al.,, 2019) as well as in computer vision (Dosovitskiy et al.,, 2020; Zhao et al.,, 2020; Zhai et al.,, 2021; Lee et al.,, 2019). As the key building block of the Transformer, the self-attention mechanism takes the following residual form (Yun et al.,, 2019) given a $n$-sequence $(x_{1},x_{2},...,x_{n})$, embedded in dimension $d$:

In this work, we propose to take the normalization process further by successively normalizing the rows and columns of $K^{0}$. This process is known to provably converge to a doubly stochastic matrix (i.e., whose rows and columns both sum to $1$) and is called Sinkhorn’s algorithm (Sinkhorn,, 1964; Cuturi,, 2013; Peyré et al.,, 2019). We denote the resulting doubly stochastic matrix $K^{\infty}$. Intuitively, such a normalization relies on a democratic principle where all points are matched one to another with different degrees of intensity, so that more interactions are considered than with the SoftMax normalization, as shown in Figure 1.

We call our Transformer variant where the SoftMax is replaced by Sinkhorn a Sinkformer. Since Sinkhorn’s first iteration coincides exactly with the SoftMax, Sinkformers include Transformers as a special case. Our modification is differentiable, easy to implement using deep learning libraries, and can be executed on GPUs for fast computation. Because the set of row-wise stochastic matrices contains the set of doubly stochastic matrices, the use of doubly stochastic matrices can be interpreted as a prior. On the experimental side, we confirm that doubly stochastic attention leads to better accuracy in several learning tasks. On the theoretical side, doubly stochastic matrices also give a better understanding of the mathematical properties of self-attention maps.

To summarize, we make the following contributions.

We show empirically that row-wise stochastic matrices seem to converge to doubly stochastic matrices during the learning process in several classical Transformers (Figure 2). Motivated by this finding, we then introduce the Sinkformer, an extension of the Transformer in which the SoftMax is replaced by the output of Sinkhorn’s algorithm. In practice, our model is parametrized by the number of iterations in the algorithm, therefore interpolating between the Transformer and the Sinkformer.

On the theoretical side, we show that Transformers and Sinkformers can be viewed as models acting on discrete distributions, and we show under a symmetry assumption that Sinkformers can be seen in the infinite depth limit as a Wasserstein gradient flow for an energy minimization (Proposition 2). We also show that the classical Transformer with the SoftMax operator cannot be interpreted as such a flow (Proposition 3). To the best of our knowledge, this is the first time such a connection is established. We also prove that in the infinite number of particles limit (when $n$ goes to infinity), the iterations of Sinkformers converge to the heat equation (Theorem 1), while the corresponding equation for Transformers is nonlinear and nonlocal (Proposition 4).

On the experimental side, we show that Sinkformers lead to a significant accuracy gain compared to Transformers on the ModelNet 40 3D shapes classification task. We then demonstrate better performance of Sinkformers on the NLP IMDb dataset for sentiment analysis and IWSLT’14 German to English neural machine translation tasks. Sinkformers also achieve a better accuracy than Vision Transformers on image classification tasks. Therefore, the proposed method is capable of enhancing the performance of transformers in a wide range of applications.

Background and related work

Proposed by Vaswani et al., (2017), the Transformer is a fully attention-based architecture. Originally designed to process sequences for natural language processing (NLP), many variants have since been developed such as Vision Transformers (Dosovitskiy et al.,, 2020; Zhai et al.,, 2021), Set Transformers (Lee et al.,, 2019) or Point Cloud Transformers (Zhao et al.,, 2020). The Transformer and its variants are based on an encoder-decoder structure, where the decoder can have a more or less complex form. The encoder is fully self-attention based. After embedding and concatenating with positional encoding the original input sequence, the encoder uses a series of residual blocks that iterates relation (1) followed by a feed forward neural network applied to each $x_{i}$ independently. In its most complex form such as in neural machine translation, the decoder combines a self-attention based mechanisms and a cross attention one, meaning that it is given access to the encoder via another multi-head attention block.

Sinkhorn and Attention.

Impact of bi-normalization.

Theoretical properties of kernels $\mathcal{K}$, which attention is an instance of, can also be studied through the operator $f\mapsto f-\mathcal{K}f$. Bi-normalization of kernels over manifolds have already been studied in the literature, on uniform measures (Singer,, 2006), weighted measures (Hein et al.,, 2007) and in a more general setup with associated diffusion operators (Ting et al.,, 2011). Milanfar, (2013) proposes to approximate smoothing operators by doubly stochastic matrices using Sinkhorn’s updates, leading to better performance in data analysis and signal processing. Importantly, the works of Marshall and Coifman, (2019) and Wormell and Reich, (2021) exactly introduce a normalization that is based on Sinkhorn’s algorithm. They prove that this method models a Langevin diffusion and leads to the approximation of a symmetric operator. They also show that convergence to this operator is faster with Sinkhorn normalization than with the SoftMax normalization. In section 5, we adopt a similar point of view with a parametrized cost and show that different normalizations result in different partial differential equations (PDEs) in the infinite number of particles limit.

Infinite depth limit.

Studying deep residual neural networks (ResNets) (He et al.,, 2016) in the infinitesimal step-size regime (or infinite depth limit) has recently emerged as a new framework for analyzing their theoretical properties. The ResNet equation

can indeed be seen as a discretized Euler scheme with unit step size of the ordinary differential equation (ODE) $\dot{x_{i}}=T(x_{i})$ (Weinan,, 2017; Chen et al.,, 2018; Teh et al.,, 2019; Sun et al.,, 2018; Weinan et al.,, 2019; Lu et al.,, 2018; Ruthotto and Haber,, 2019; Sander et al.,, 2021). In section 4, we adopt this point of view on residual attention layers in order to get a better theoretical understanding of attention mechanisms. This is justified by the fact that, for instance, GPT-3 (Brown et al.,, 2020) has 96 layers.

Neural networks on measures.

Sinkformers

We now introduce Sinkformers, a modification of any Transformer by replacing the SoftMax operator in the attention modules by Sinkhorn’s algorithm.

In Transformers, attention matrices are row-wise stochastic. A natural question is how the sum over columns evolve during training. On $3$ different models and $3$ different learning tasks, we calculated the sum over columns of attention matrices in Transformers. We find out that the learning process makes the attention matrices more and more doubly stochastic, as shown in Figure 2.

Thus, row-wise stochastic attention matrices seem to approach doubly stochastic matrices during the learning process in classical Transformers. Therefore, it seems natural to impose double stochasticity as a prior and study theoretically and experimentally the resulting model. A process to obtain such matrices which extends the SoftMax is Sinkhorn’s algorithm.

Sinkhorn’s algorithm.

Sinkformers.

For simplicity, we consider a one head attention block that iterates equation (1). Note that $K^{1}\coloneqq\texttt{SoftMax}(C)$ is precisely the output of Sinkhorn’s algorithm (3) after $1$ iteration. In this paper, we propose to take Sinkhorn’s algorithm several steps further until it approximately converges to a doubly stochastic matrix $K^{\infty}$. This process can be easily implemented in practice, simply by plugging Sinkhorn’s algorithm into self-attention modules in existing architectures, without changing the overall structure of the network. We call the resulting drop-in replacement of a Transformer a Sinkformer. It iterates

In the next two sections 4 and 5, we investigate the theoretical properties of Sinkformers. We exhibit connections with energy minimization in the space of measures and the heat equation, thereby proposing a new framework for understanding attention mechanisms. All our experiments are described in Section 6 and show the benefits of using Sinkformers in a wide variety of applications.

Computational cost and differentiation.

Turning a Transformer into a Sinkformer simply relies on replacing the SoftMax by Sinkhorn, i.e., substituting $K^{1}$ with $K^{\infty}$. In practice, we use a finite number of Sinkhorn iterations and therefore use $K^{l}$, where $l$ is large enough so that $K^{l}$ is almost doubly stochastic. Doing $l$ iterations of Sinkhorn takes $l$ times longer than the SoftMax. However, this is not a problem in practice because Sinkhorn is not the main computational bottleneck and because only a few iterations of Sinkhorn are sufficient (typically $3$ to $5$) to converge to a doubly stochastic matrix. As a result, the practical training time of Sinkformers is comparable to regular Transformers, as detailed in our experiments.

Sinkhorn is perfectly suited for backpropagation (automatic differentiation), by differentiating through the operations of (3). The Jacobian of an optimization problem solution can also be computed using the implicit function theorem (Griewank and Walther,, 2008; Krantz and Parks,, 2012; Blondel et al.,, 2021) instead of backpropagation if the number of iterations becomes a memory bottleneck. Together with Sinkhorn, implicit differentiation has been used by Luise et al., (2018) and Cuturi et al., (2020).

Invariance to the cost function.

Recall that in practice one has $C_{i,j}=({W_{Q}x_{i}})^{\top}{W_{K}x_{j}}$. An important aspect of Sinkformers is that their output is unchanged if the cost is modified with non interacting terms, as the next proposition shows.

Attention and gradient flows

We denote $k^{\infty}\coloneqq\texttt{Sinkhorn}(c)$ the resulting limit. Note that if $\mu$ is a discrete measure supported on a $n$ sequence of particles $(x_{1},x_{2},...,x_{n})$, $\mu=\frac{1}{n}\sum_{i=1}^{n}\delta_{x_{i}}$, then for all $(i,j)$, $k^{0}(x_{i},x_{j})=K^{0}_{i,j}$, $k^{1}(x_{i},x_{j})=K^{1}_{i,j}$ and $k^{\infty}(x_{i},x_{j})=K^{\infty}_{i,j}$, so that $k^{0}$, $k^{1}$ and $k^{\infty}$ are indeed the continuous equivalent of the matrices $K^{0}$, $K^{1}$ and $K^{\infty}$ respectively.

Infinitesimal step-size regime.

In order to better understand the theoretical properties of attention matrices in Transformers and Sinkformers, we omit the feed forward neural networks acting after each attention block. We consider a succession of attention blocks with tied weights between layers and study the infinite depth limit where the output is given by solving a neural ODE (Chen et al.,, 2018). In this framework, iterating the Transformer equation (1), the ResNet equation (2) and the Sinkformer equation (4) corresponds to a Euler discretization with step-size $1$ of the ODEs

When $T_{\mu}$ is defined by the ResNet equation (2), $T_{\mu}=T$ does not depend on $\mu$. It defines an advection equation where the particles do not interact and evolve independently. When $T_{\mu}$ is defined by the Transformer equation (1) or Sinkformer equation (4), $T_{\mu}$ has a dependency in $\mu$ and the particles interact: the local vector field depends on the position of the other particles. More precisely we have in this case $T^{1}_{\mu}(x)=\int k^{1}(x,x^{\prime})W_{V}x^{\prime}d\mu(x^{\prime})$ for the Transformer and $T^{\infty}_{\mu}(x)=\int k^{\infty}(x,x^{\prime})W_{V}x^{\prime}d\mu(x^{\prime})$ for the Sinkformer. It is easily seen that when $\mu$ is discrete we recover the operators in equation (1) and (4).

Wasserstein gradient flows.

Particular case.

Flows for attention.

Our goal is to determine the PDEs (7) defined by the proposed attention maps. We consider the symmetric case, summarized by the following assumption:

${W_{K}}^{\top}W_{Q}={W_{Q}}^{\top}W_{K}=-W_{V}$

Assumption 1 means we consider symmetric kernels (by imposing ${W_{K}}^{\top}W_{Q}={W_{Q}}^{\top}W_{K}$), and that when differentiating $x\mapsto\exp(c(x,x^{\prime}))$, we obtain $-\exp(c)W_{V}$. We show that, under this assumption, the PDEs defined by $k^{0}$ and $k^{\infty}$ correspond to Wasserstein gradient flows, whereas it is not the case for $k^{1}$. A particular case of imposing ${W_{K}}^{\top}W_{Q}={W_{Q}}^{\top}W_{K}$ is when $W_{Q}=W_{K}$. This equality setting is studied by Kim et al., (2021), where the authors show that it leads to similar performance for Transformers. Since imposing ${W_{K}}^{\top}W_{Q}={W_{Q}}^{\top}W_{K}$ is less restrictive, it seems to be a natural assumption. Imposing $W_{Q}^{\top}W_{K}=-W_{V}$ is more restrictive, and we detail the expressions for the PDEs associated to $k^{0},k^{1},k^{\infty}$ without this assumption in Appendix A. We have the following result.

A proof is given in Appendix A. Proposition 2 shows that $k^{0}$ and $k^{\infty}$ correspond to Wasserstein gradient flows. In addition, the PDE defined by $k^{1}$ does not correspond to such a flow. More precisely, we have the following result.

One has that $T_{\mu}^{1}=-\nabla[\log(\int k^{0}(\cdot,x^{\prime})d\mu(x^{\prime}))]$ is not a Wasserstein gradient.

A proof is given in Appendix A, based on the lack of symmetry of $T_{\mu}^{1}$. As a consequence of these results, we believe this variational formulation of attention mechanisms for Sinkformers (Proposition 2) provides a perspective for analyzing the theoretical properties of attention-based mechanisms in light of Wasserstein gradient flow theory (Santambrogio,, 2017). Moreover, it makes it possible to interpret Sinkformers as argmin layers, which is promising in terms of theoretical and experimental investigations, and which is not possible for Transformers, according to Proposition 3.

Our results are complementary to the one of Dong et al., (2021), where the authors show that, with no skip connections and without the feed forward neural network acting after each attention block, the output of a Transformer converges doubly exponentially with depth to a rank-1 matrix. On the contrary, we propose a complementary analysis by taking skip-connections into account, as is standard in Transformers. Precisely because we consider such connections, we end up with very different behaviors. Indeed, as shown in the next section, our analysis reveals that the relative signs for $W_{K}$, $W_{Q}$ and $W_{V}$ imply very different behavior, such as aggregation or diffusion. The dynamics obtained when considering skip connections are therefore richer than a rank collapse phenomenon.

Attention and diffusion

A proof is available in Appendix A, making use of Theorem 1 from Marshall and Coifman, (2019). We recover in Equation (8) the well-known heat equation.

A proof is given in Appendix A. While equation (8) corresponds to the heat equation, equation (9) is different. First, it is nonlinear in $\rho$. Second, it is nonlocal since the evolution of the density at $x$ depends on the value of this density at location $W^{-1}_{K}W_{Q}x$. Note that the linear and local aspect of Sinkformer’s PDE on the one hand, and the nonlinear and nonlocal aspect of Transformer’s PDE on the other hand, remain true without assuming ${W_{Q}}^{\top}W_{K}=-W_{V}$ (details in Appendix A).

Experiments

We now demonstrate the applicability of Sinkformers on a large variety of experiments with different modalities. We use Pytorch (Paszke et al.,, 2017) and Nvidia Tesla V100 GPUs. Our code is open-sourced and is available at this address: https://github.com/michaelsdr/sinkformers. All the experimental details are given in Appendix C.

In all our experiments, we use existing Transformer architectures and modify the SoftMax operator in attention modules with Sinkhorn’s algorithm, which we implement in $\log$ domain for stability (details in Appendix B).

1 ModelNet 40 classification

The ModelNet 40 dataset (Wu et al.,, 2015) is composed of 40 popular object categories in 3D. Transformers for point clouds and sets have been applied to the ModelNet 40 classification in several works, such as Set Transformers (Lee et al.,, 2019) or Point Cloud Transformers (Guo et al.,, 2021).

Set Transformers (Lee et al.,, 2019) also have an encoder decoder structure with different possibilities for defining attention-based set operations. We propose to focus on the architecture that uses Induced Self Attention Block (ISAB), which bypasses the quadratic time complexity of Self Attention Blocks (SAB). More details about this architecture can be found in (Lee et al.,, 2019). We reproduce the ModelNet 40 classification experiment using $5000$ uniformly sampled points for each shape and use a Set Transformer and a Set Sinkformer with two ISAB layers in the encoder and a decoder composed of a SAB and a Pooling by Multihead Attention (PMA) module. While the reported test accuracy is of $87.8\%$ using a Set Transformer, we obtain as our best accuracy when performing $21$ iterations of Sinkhorn algorithm within our Sinkformer of $89.1\%$. Results are summarized in Table 1.

Moreover, we show in Figure 3 the learning curves corresponding to this experiment. Interestingly, the number of iterations within Sinkhorn’s algorithm increases the accuracy of the model. Note that we only consider an odd number of iterations since we always want to have row-wise stochastic attention matrices to be consistent with the properties of the SoftMax.

Point Cloud Transformers.

We also train Point Cloud Transformers (Guo et al.,, 2021) on ModelNet 40. This architecture achieves accuracy comparable to the state of the art on this dataset. We compare best and median test accuracy over $4$ runs. Results are reported in Table 1, where we see that while the best test-accuracy is narrowly achieved for the Transformer, the Sinkformer has a slightly better median accuracy.

2 Sentiment Analysis

We train a Transformer (composed of an attention-based encoder followed by a max-pooling layer) and a Sinkformer on the IMDb movie review dataset (Maas et al.,, 2011) for sentiment analysis. This text classification task consists of predicting whether a movie review is positive or negative. The learning curves are shown in Figure 4, with a gain in accuracy when using a Sinkformer. In this experiment, Sinkhorn’s algorithm converges perfectly in 3 iterations (the resulting attention matrices are doubly stochastic), which corresponds to the green curve. The Sinkformer only adds a small computational overhead, since the training time per epoch is $4$m $02$s for the Transformer against $4$m $22$s for the Sinkformer.

3 Neural Machine Translation

We train a Transformer and its Sinkformer counterpart using the fairseq (Ott et al.,, 2019) sequence modeling toolkit on the IWSLT’14 German to English dataset (Cettolo et al.,, 2014). The architecture used is composed of an encoder and a decoder, both of depth $6$. We plug Sinkhorn’s algorithm only into the encoder part. Indeed, in the decoder, we can only pay attention to previous positions in the output sequence. For this reason, we need a mask that prevents a straightforward application of Sinkhorn’s algorithm. We demonstrate that even when using the hyper-parameters used to optimally train the Transformer, we achieve a similar BLEU (Papineni et al.,, 2002) over $6$ runs. We first train a Transformer for $30$ epochs. On the evaluation set, we obtain a BLEU of $34.43$. We then consider a Sinkformer with the weights of the trained Transformer. Interestingly, even this un-adapted Sinkformer provides a median BLEU score of $33.81$. We then divide the learning rate by $10$ and retrain for $5$ additional epochs both the Transformer and the Sinkformer to obtain a median BLEU of respectively $34.68$ and $34.73$ (Table 2). Importantly, the runtime for one training epoch is almost the same for both models: $2$m $48$s (Transformer) against $2$m $52$s (Sinkformer).

4 Vision Transformers

Vision Transformers (ViT) (Dosovitskiy et al.,, 2020) have recently emerged as a promising architecture for achieving state of the art performance on computer vision tasks (Zhai et al.,, 2021), using only attention based mechanisms by selecting patches of fixed size in images and feeding them into an attention mechanism.

We train a ViT and its Sinkformer counterpart on a binary cats and dogs image classification task. The evolution of the train and test accuracy is displayed in Figure 5. The median test accuracy is $79.0\%$ for the Transformer against $79.5\%$ for the Sinkformer, whereas the maximum test accuracy is $80.0\%$ for the Transformer against $80.5\%$ for the Sinkformer. We also use $3$ iterations in Sinkhorn’s algorithm which leads to a negligible computational overhead (training time per epoch of 3m 25s for the Sinkformer against 3m 20s for the Transformer).

Impact of the patch size on the final accuracy.

We consider a one-layer and one-head self-attention module on the MNIST dataset, with no additional layer. The purpose is to isolate the self-attention module and study how its accuracy is affected by the choice of the patch size. Results are displayed in Figure 6. We recall that a MNIST image is of size $28\times 28$. When taking only one patch of size $28$, both models are equivalent because the attention matrix is of size $1$. However, when the patch size gets smaller, the two models are different and the Sinkformer outperforms the Transformer.

Conclusion

In this paper, we presented the Sinkformer, a variant of the Transformer in which the SoftMax, which leads to row-wise stochastic attention, is replaced by Sinkhorn’s algorithm, which leads to doubly stochastic attention. This new model is motivated by the empirical finding that attention matrices in Transformers get closer and closer to doubly stochastic matrices during the training process. This modification is easily implemented in practice by simply replacing the SoftMax in the attention modules of existing Transformers without changing any parameter in the network. It also provides a new framework for theoretically studying attention-based mechanisms, such as the interpretation of Sinkformers as Wasserstein gradient flows in the infinitesimal step size regime or as diffusion operators in the mean-field limit. On the experimental side, Sinkformers lead to better accuracy in a variety of experiments: classification of 3D shapes, sentiment analysis, neural machine translation, and image classification.

Acknowledgments

This work was granted access to the HPC resources of IDRIS under the allocation 2020-[AD011012073] made by GENCI. This work was supported in part by the French government under management of Agence Nationale de la Recherche as part of the “Investissements d’avenir” program, reference ANR19-P3IA-0001 (PRAIRIE 3IA Institute). This work was supported in part by the European Research Council (ERC project NORIA). We thank Marco Cuturi and D. Sculley for their comments on a draft of the paper. We thank Scott Pesme, Pierre Rizkallah, Othmane Sebbouh, Thibault Séjourné and the anonymous reviewers for helpful feedbacks.

Appendix

In Section A we give the proofs of all the Propositions and the Theorem. In Section B we present the implementation details of Sinkformers. Section C gives details for the experiments in the paper.

Appendix A Proofs

We use the variational formulation for Sinkhorn (Peyré et al.,, 2019):

Recall that for $p\in\{0,1,\infty\}$, we have $T_{\mu}^{p}(x)=\int k^{p}(x,x^{\prime})W_{V}x^{\prime}d\mu(x^{\prime})$.

We can now derive the different gradient expressions for $T_{\mu}^{0}$, $T_{\mu}^{1}$ and $T_{\mu}^{\infty}$.

For $T_{\mu}^{0}$: under Assumption 1, we have that $f(x,x^{\prime})=e^{C(x,x^{\prime})}$ is symmetric. This gives

and by differentiation under the integral, under sufficient regularity assumptions on $\mu$, this gives

Since $\nabla_{x}c(x,x^{\prime})=-W_{V}x^{\prime}$, we get

For $\mu=\frac{1}{n}\sum_{i=1}^{n}\delta_{x_{i}}$ this is exactly

For $T_{\mu}^{2}$: one has the dual formulation for $\mathcal{F}^{\infty}$ (Peyré et al.,, 2019):

where we denote the soft $c$ transform as

which actually depends on $\mu$ and $c$. One has for an optimal pair $f=f^{c}$ (Peyré et al.,, 2019). In addition, one has $k^{\infty}(x,x^{\prime})=e^{c(x,x^{\prime})+f(x)+f(x^{\prime})}$. The Wasserstein gradient of $\mathcal{F}^{\infty}$ is then

where $f$ is an optimal solution of (10) (which is unique up to a constant). The gradient of $f$ can be obtained using (11) and the fact that $f=f^{c}$:

A.3 The SoftMax normalization does not correspond to a gradient flow - Proof of Proposition 3

Taking $\mu=\delta_{y}$ gives $k^{0}(x,y)=k^{0}(x^{\prime},y)$, which by symmetry implies that $k^{0}$ is a constant.

This is a contradiction since $k^{0}(x,x^{\prime})=\exp(x^{\top}W_{Q}^{\top}W_{K}x^{\prime})$. ∎

A.4 Sinkformer’s PDE - Proof of Theorem 1

We perform the change of variable $y=Ax^{\prime}$. This gives

where $C_{A}$ depends only on $A$. We then apply Theorem 1 from Marshall and Coifman, (2019) with $f=W_{V}A^{-1}$, $q(x)=\rho(A^{-1}x)$ and $w=\frac{1}{C_{A}}$, to obtain that

in $L^{2}$ norm. Since $q(x)=\rho(A^{-1}x)$ we have obtained that $\frac{\nabla q}{q}=A^{-1}\frac{\nabla\rho}{\rho}(A^{-1}\cdot)$ so that

which is exactly what we wanted to show. Note that when $W_{V}=-W_{Q}^{\top}W_{K}$ this gives the expected result. The general form for the PDE is then

A.5 Transformer’s PDE - Proof of Proposition 4

We perform the change of variable $y=W_{K}x^{\prime}$. This gives:

Using the Laplace expansion result from Singer, (2006), we obtain that

By doing a Taylor expansion for the denominator, we find

Since $\overline{T}^{1}_{\mu,\varepsilon}={T}^{1}_{\mu,\varepsilon}+\frac{1}{\varepsilon}W_{Q}^{\top}W_{Q}=\frac{1}{\varepsilon}(g_{\varepsilon}(W_{Q}x)+W_{Q}^{\top}W_{Q}x)$ and because $W_{V}W^{-1}_{K}=-{W_{Q}}^{\top}W_{K}W^{-1}_{K}=-W^{\top}_{Q}$ we have

Appendix B Implementation details

This allows for fast and accurate computations, where $\log({Ke^{g^{l}}})$ and $\log({K^{\top}e^{f^{l}}})$ are computed using log-sum-exp.

Appendix C Experimental details

For our experiments on ModelNet using Set Transformers, we first prepossess the ModelNet 40 dataset. We then uniformly sample $5000$ points from each element in the dataset. Our architecture is composed of two ISAB layers in the encoder and a decoder composed of a SAB and a Pooling by Multihead Attention (PMA) module. For the training, we use a batch-size of $64$ and we use Adam (Kingma and Ba,, 2014). The training is done over $300$ epochs. The initial learning rate is $10^{-3}$ and is decayed by a factor $10$ after $200$ epochs.

Point Cloud Transformers.

For our experiments on ModelNet using Point Clouds Transformers, we uniformly sample $1024$ points from each element in the dataset. For the training, we use a batch-size of $32$ and we use SGD (Ruder,, 2016). The training is done over $300$ epochs. The initial learning rate is $10^{-4}$ and is decayed by a factor $10$ after $250$ epochs.

C.2 Sentiment Analysis

We use the code available at the repository nlp-turorialhttps://github.com/lyeoni/nlp-tutorial/tree/master/text-classification-transformer, where a pretrained Transformer is fine-tuned on the IMDb dataset. In our experiment, we reset the parameters of the pretrained Transformer and train it from scratch on the IMDb dataset. We use an architecture of depth $6$, with $8$ heads. For the training, we use a batch-size of $32$ and we use Adam. The training is done over $15$ epochs. The initial learning rate is $10^{-4}$ and is decayed by a factor $10$ after $12$ epochs.

C.3 Neural Machine Translation

We use the Transformer from fairseq and the command for training it on the IWSLT’14https://github.com/pytorch/fairseq/blob/main/examples/translation/README.md dataset. When fine-tuning a Sinkformer, we simply divide the original learning rate by $10$.

C.4 Vision Transformers

This experiment is done on the cats and dogshttps://www.kaggle.com/c/dogs-vs-cats/data dataset. For this experiment, we use a batch-size of $64$ and Adam. We use an architecture of depth $6$, with $8$ heads, and select a patch-size of $16$. The training is done over $300$ epochs. The initial learning rate is $5\times 10^{-5}$ and divided by $10$ after $250$ epochs.

Impact of the patch size on the final accuracy.

For this experiment, we use a batch-size of $100$ and Adam. We use an architecture of depth $1$, with $1$ heads, without non-linearity, and select different values for the patch-size. The training is done over $45$ epochs. The initial learning rate is $1\times 10^{-3}$ (resp. $2\times 10^{-3}$) for the Transformer (resp. Sinkformer) and divided by $10$ after $35$ epochs and again by $10$ after $41$ epochs.