Manifold-Preserving Transformers are Effective for Short-Long Range Encoding
Authors: Ayan Sengupta
Md Shad Akhtar
Tanmoy Chakraborty
Venue: EMNLP 2023 Findings (2023)
Keywords: Orthogonal attention
Lipschitz
Entropic Transformer

Abstract: Multi-head self-attention-based Transformers have shown promise in different learning tasks. Albeit these models exhibit significant improvement in understanding short-term and long-term contexts from sequences, encoders of Transformers and their variants fail to preserve layer-wise contextual information. Transformers usually project tokens onto sparse manifolds and fail to preserve mathematical equivalence among the token representations. In this work, we propose TransJect, an encoder model that guarantees a theoretical bound for layer-wise distance preservation between a pair of tokens. We propose a simple alternative to dot-product attention to ensure Lipschitz continuity. This allows TransJect to learn injective mappings to transform token representations to different manifolds with similar topology and preserve Euclidean distance between every pair of tokens in subsequent layers. Evaluations across multiple benchmark short- and long-sequence classification tasks show maximum improvements of 6.8% and 5.9%, respectively, over the variants of Transformers. Additionally, TransJect displays 79% better performance than Transformer on the language modeling task. We further highlight the shortcomings of multi-head self-attention from the statistical physics viewpoint. Although multi-head self-attention was incepted to learn different abstraction levels within the networks, our empirical analyses suggest that different attention heads learn randomly and unorderly. In contrast, TransJect adapts a mixture of experts for regularization; these experts are more orderly and balanced and learn different sparse representations from the input sequences. TransJect exhibits very low entropy and can be efficiently scaled to larger depths.

Manifold-Preserving Transformers are Effective for Short-Long Range Encoding
Ayan Sengupta ITT Delhi, India eez228655@ee iitd.ac.in ac. in
Md. Shad Akhtar HIT Delhi, India shad. akhtarr ac. in
Tanmoy Chakraborty ITT Delhi, India tanchak@iitd.as in
Abstract
2.0
Cognre Distance
Multi-head self-attention- based Transformers have shown promise in different learning tasks. Albeit these models exhibit significant improve- in understanding short-term and long- term contexts from sequences, encoders of Transformers and their variants fail to preserve layer-wise contextual information. Transform- ers usually project tokens onto sparse mani- folds and fail preserve mathematical equiv- alence among the token representations. In this work, we propose TransJect, an encoder model that guarantees theoretical bound for layer-wise distance preservation between pair tokens. We propose simple alternative to dot-product attention ensure Lipschitz conti- nuity. This allows TransJect to learn injective mappings transform token representations different manifolds with similar topology and preserve Euclidean distance between ev- ery pair of tokens in subsequent layers. Eval- uations across multiple benchmark short- and long-sequence classification tasks show maxi- mum improvements of 6.8% and 5.9%, respec- tively, over the variants of Transformers. Ad- ditionally, TransJect displays 79% better per- formance than Transformer the language modeling task. We further highlight the short- comings of multi-head self-attention from statistical physics viewpoint. Although multi- head self-attention was incepted to learn dif- ferent abstraction levels within the networks, our empirical analyses suggest that different at- tention heads learn randomly and unorderly. In contrast, TransJect adapts mixture of experts for regularization; these experts are more orderly and balanced and learn different sparse representations from the input sequences. JeansJect exhibits very low entropy and can be efficiently scaled larger depths.
0.5
Eurliten Disunce
8000
4000
2000
40
Layens
"novis" "taribre"
good" "effects"
"Ingvic" "8"
Figure Layer-wise distances between few selected tokens from the text "This movie is terrible but it has some good effects." We use a trained Transformer model extract token representations from different layers.
downstream applications. Stacking multiple lay- ers has been proven effective in extracting features at different levels of abstraction, thereby learning more complex patterns (Brightwell et al., 1996; Poole al., 2016). Since then, tremendous ef- forts have been made in building larger depth mod- els and in making them faster (Bachlechner et al., 2020; Xiao et al.). Self-attention-buse Trans- former model (Vaswani et al., 2017) was proposed to parallelize the computation of longer sequences; it has achieved satte-of-the-an performance in vari- ous sequence modeling tasks. Following this, nu- merous efforts have been made reduce computa- tion and make the Transformer model suitable even for longer sequences (Katharopoulos al., 2020; Peng et al., 2020; Kitaev al., 2020; Beltagy et al.. 2020; Press et al., 2021; Choromanski et al., 2021; Tay et al., 2021). However, very few of these stud- ies discuss information propagation in large-depth models. A recent study (Voita et al., 2019) char- acterized how Transformer token representations change across layers for different training objec- tives. The dynamics of layer-wise learning are essential understand different abstract levels a model could Iearn, preserve and forget to continue learning throughout the layers. However, recent
1 Introduction
Over the past few decades, Deep Neural Networks have greatly improved the performance various

---

studies need shed more light on how Transform- ers preserve contextual information across differ- ent layers (Wu et al.. 2020; Voita et al., 2019). To understand how Transformer encodes contextual similarities between tokens and preserves the sim- ilarities layer-wise and across different attention heads, we highlight example in Figure 1. We select three pairs of semantically-relate tokens. We observe the Euclidean and cosine distances of the representations learned at different layers of the Transformer trained on the IMDb sentiment clas- sification task. Although the trained Transformer model preserves the semantic similarity among the same tokens across layers, the Euclidean distance between their representations increases in the upper layers. This indicates that Transformer projects the representations to different and sparse subspaces (a subspace with low density), albeit preserving the angle between them. Additionally, we observe that the distances between different token repre- sentations vary across different attention heads at different encoder layers in a haphazard manner. Preserving distance and semantic similarity across layers is vital to ensure the continual learning capa- bilities of deep neural models, which Transformer evidently fails to do.
based Transformer with enforced injectivity, aka TransJect. With injectivity, our model imposes the constraint in which the representations of two distinct tokens always differ across all the layers. Unlike Transformer, which only injects regulariza- tion through multi-head self-attention and dropout, TransJect does not require explicit regularizers and can be regularized implicitly due to its inher- ent injectivity. The backbone of TransJect is non non-normalized linear orthogonal attention and an injective residual connection; both ensure Lipschitz continuity. By enforcing Lipschítz continuity in Euclidean and dot- dot-product space, the model can preserve the manifold structure between tokens and perform well on the final predictive task.
To validate our hypotheses and empirically jus- tify the superiority of our model, we use two short and five long sequence classification tasks. TransJect outperforms Transformer and other benchmark variants with an average margin of 3.4% and 2.2% on the short- and -equese classification tasks, respectively. Our model per- forms best on the long sequence classification (LRA) benchmark, achieving 0.2% better accu- racy than the best baseline, Skyformer (Chen al., 2021). We further demonstrate TransJect on lan- guage modeling task on the Penn TreeBank dataset, in which our model achieves 79% better test per- plexity than the vanilla Transformer. Empirical analyses suggest a very low entropy of TransJect representations, indicating that TransJect cap- tures sparse and more orderly representations compared to Transformer. Moreover, TransJect shows 13x lower inference runtime than Trans- former, indicating efficiency in encoding long input sequences.'
Neuroscientists have been working for years to understand how the human brain functions and sim- ulates the behaviours of physical scienees (Koch and Hepp, 2006; Vértes et al., 2012). Arguably, the human brain is more capable of 'associative learning' and 'behavioural formation', which can be attributed to the number of neurons and their inter- connectedenntss (synapses) rather than the size of the brain, or the number of layers through which the information propagates (Dicke and Roth, 2016). Although fully-grown human brain can have hun- dreds of billions of neurons, it has been found that at a time, only a tiny fraction of neurons fire (Ahmed et al., 2020; Poo and Isaacson, 2009). This sparse firing can help the brain sustain low en- tropy (energy). Entropy, a measure that quantifies a state's 's randomness, has thus become an essential tool to understand the factors behind human intelli- gence (Saxe et al., 2018; Keshmiri, 2020) and how the human brain operates.
2 Related Works
Despite being ubiquitous topic of study across different disciplines of deep learning. Trans- former (Vaswani et al., 2017) models still require better mathematical formalization. On recent development, Vuckovic et al. (2020) formalized the inner workings of self-attention maps through the lens of measure theory and established the Lipschitz continuity of self-attention under suit- able assumptions. However, the Lipschitz con- dition depends the boundedness of the repre- sentation space and the Lipschitz bound of the
Unfortunately, Transformer and its variants are yet to be studied from these interdisciplinary view- points. Our study explores the connection be- tween model sparsity and entropy. Towards this, we propose a complete redesign of attennnte
'The source codes of TransJect can be found at https: thub wwchhuu.....vviiiieeeeeeeeee.eea
git.

---

fully-connected feed-forward layer (FFN). A simi- lar study (Kim al., 2021) also concluded that the dot-product self-attention neither Lipschitz nor injective under the standard conditions. Injectivity of the transformation map is essential to ensure that the function is bijective and, therefore, reversible. Reversibility within deep neural networks has al- ways been an active area of study (Gomez al., 2017; Arora et al., 2015; Chang al., 2018). A reversible network ensures better scaling large depths and is more efficient than non-reversible structures. Recently, Mangalam et al. (2022) de- signed a reversible Transformer and empirically highlighted its effectiveness in several image and video classification tasks. However, any similar development has yet to be made for developing scalable reversible sequential models.
inherently regularized and display lower entropy. projecting the representations onto sparse sub- space pushes them further from being Lipschitz, preventing them from reversible.
This work proposes injective and Lips- chitz continuous alternative to vanilla Transformer, namely, TransJect. With the enforeed injectivity, our model establishes a theoretical guarantee for reversibility and thus can be scaled to larger depths. Further, TransJect displays significantly lower en- tropy than Transformer, indicating a lower energy footprint and more efficiency.
3 Designing Injective Transformer
This section formally describes our proposed model, TransJect. It inherits the structure from the vanilla Transformer and achieves a smoother ac- tivation plane by utilizing injective maps for trans- forming token representations across layers. For an L-layered stacked encoder, we aim to learn the rep- resentation of sequence X {Iı, #2, :.. (N) at each layer / that preserves the pairwise distance between every pair of words within a theoretical bound. We illustrate the components of TransJect in Figure 2. All the proofs presented in the paper are supplied in Appendix A.
Dynamical isometry is property that mandates the singular values of the input-output Jacobian matrix to be closer one. Pennington et al. (2017) showed that dynamical isometry could aid faster convergence and efficient learning. Following their idea, Bachlechner et al. (2020) showed that residual connections in Transformers do not often satisfy dynamical isometry, leading to poor signal propa- gation through the models. To overcome this, they proposed residual with zero initialization (ReZero) and claimed dynamical isometry for developing faster and more efficient large-depth Transform- ers. Previously, Qi et (2020) enforced stronger condition of isometry to develop deep convolution networks efficiently. However, isometric assump- tions are not valid for sequential modeling due to different levels of abstraction within the input sig- nals. Moreover, isometry may not hold between mentextuualy-dissimiee tokens.
3.1 Background
Activation bound. For any function f R Rm, we define the activation bound Ky as [f(=))]b> for a suitable integer p. A linear sup ||l|| map M equals the induced matrix norm || ||p. Intuitively, this is the maximum scale factor by which a mapping expands a vector æ. In Euclidean space, we usually choose = 2.
Another essential aspect behind designing effi- cient large-depth models is ensuring model spar- sity. Several notable contributions have been made (Baykal al., 2022; Jaszezur al., 2021; Li et al., 2023; Tay et al., 2020a) to enforce sparse ac- tivations within Transformers to make them more efficient and scalable. Li et al. (2023) argued that sparse networks often resemble the sparse activa- tions by the human brain, bearing similarity be- tween artificial and biological networks. They em- pirically showed that the trained Transformers are inherently sparse, and the sparsity emerges from all the layers. As discussed in the previous section, Transformers project the representations sparsely onto sparse subspaces. Although sparse models are
Lipschitz Continuity. A function :: Rm under |1-IIp norm is called Lipschitz continuous' if there exists a real number K 2 such that
[If()) - f())ll S K K| yllp-
(1)
for any x, y E R Lipschitz continuity can be defined over any metric space, but in this paper, we restrict its definition to only Euclidean space with 2. K is called Lipschitz bound.
3.2 Space-Preserving Orthogonal Attention
The backbone TransJect is the space- preserving orthogonal attention.

---

Input Sequences
Oethoganol
(b)
Embedding
(c)
(a)
%
U
Xg
Pos ilonnal Enocding
9 xxdd
8xdxa
Q1
Σ
V22
Cigervauue Decomppositer
E
*
x 2"
Rxdxd
dxd
(a))
Orthogonal Feecl FeeecForvard onward
Orthogonal
Attention with
EXPO
df d
Encoder Layer
Elementiwise acdiion
Orthogamal Mairs
Encoder Layer
x
(d)
(WJ
Flemen--ane dingorss matti
Orthogonal Residual
Figure 2: Internals of TransJect with (a) an I--layered encoder containing a Mixture of Experts (MOE) with experts, (b) approximated eigenvalue computation, and (c) an orthogonal attention with injective residual. (d) The outputs of the intermediate encoder are fed to the orthogonal residual FFN. The hidden dimension used for representing each token is denoted by d.
Theorem -Space-Preserving Orthogonal Atten- tion). Replacing W W WV with real square orthogonal matrices in non-normalized linear self- attention reduces the activation bound to of (X), where 01 (X) is the largest singular value of X. This reduces the attention operation Attn
the batch size and d being the hidden dimension. This work uses simple approximation to com- pute Σ, the eigenvalues of x Formally, we compute U arg miny ||x X - -UEU") and E arg ming || UEUFII To learn the approximate eigenvalues, we can minimize the re- construction loss || UEUT for a learn- able orthogonal eigenvector U. We use standard- ization to enforce the stochasticity constraint on Σ. Further, instead of deriving the cigenvalues, we can also initialize a random diagonal matrix Σ, without any approximation that optimizes the only task-specific training objective, without forcing the reconstruction. We denote this version of the model as Random-TransJect We compute Σ once, only on the initial token embeddings.
XUEV, for learnable orthogonal matrices and and the singular values Σ.
Notice that the activation bound of the modified attention mechanism does not depend on any learn- able parameters; instead can be bounded by the largest eigenvalue of X X. Therefore, we assume stochastic X X to ensure that the largest eigen- value is always 1 (see Corollary 2 in Appendix A. and the attention operator preserves the pairwise distance between any two tokens. We learn orthog- onal projection matrices in each layer, U and V. In contrast, the diagonal matrix containing eigenval- ues Σ is learned on the initial embedding obtained from the initial embedding layer defined in Sec- 0. Therefore, in our
3.3 Injective Residual (IR)
We fine-tune the hidden representation for every layer by learning a new attention projection on the hidden state learned in the previous layer. Formally, we define,
tion 3.5, also denoted proposed non-normalized orthogonal linear atten- tion, we compute the attention matrix (encoded in Σ) only once and learn different projections of in different layers.
(1) (1-1) ai F(X(-.)..
(2)
Here, F the self-attention operator, followed by a suitable non-linear activation function, and a; E (0, 1) is the residual weight. We use a learn- able residual weight and a sigmoid activation to
Approximating eigenvalues. Eigenvalue de- composition computationally expensive with a runtime complexity of O (Bdd), with being

---

scale it in (0,1). In the previous studies, ReLU and GELU (Hendrycks and Gimpel, 2016) have been popular choices for the activation function. In this work, we choose ELU (Clevert et al., 2015), non- linear C (continuous and differentiable) activation function with Lipschitz bound of 1. Although ReLU is Lipschitz function with K 1, it is not everywhere differentiable and injective. Fol- lowing Bachlechner et al. (2020), we adopt ReZero (residual with zero initialization) enforce dynam- ical isometry and stable convergence.
3.5 Injective Token Embedding
Transformer introduced conditional encoding to inject tokens' relative and absolute positional in- formation into the self-attention layer. It leverages sinusoidal representations of every position and adds to the original token embeddings to infuse the positional information. To ensure injectivity at each layer, we need to ensure that the initializa- tion of the token embeddings is also injective, i.e., no two tokens should have the same embedding. Unfortunately, the addition operator is not injec- tive. Therefore, we compute the initial embedding of token 2; as X, (0) = Concat Emb(a;), PE..). Here PE is defined similarly to the positional en- coding proposed by Vaswani et al. (2017). The embedding matrix is orthogonally parameterized. Concatenation ensures the injectivity of embed- dings. However, to maintain the dimensionality, we learn the initial embedding and positional en- coding lower dimensional space, IR, in which is the hidden size in the encoder. We define the final encoder mapping for each sublayer I com- posite mapping defined by,
Lemma (Residual contractions are injective). f: + X + 7 F(X) is injective for L 2 3. To maintain the dimensionality, Transformer projects the representations to lower-dimensional space, which reduces the number of synapses among the neurons by a factor of H, the number of heads. As opposed to this, we devise a Mixture Expert (MOE) attention (motivated by Shazeer et al. (2017)). With this, we compute (I.e) for each expert E 21, E} in each layer us- ing Equation 2, learnable expert weights A,"s, and use a convex combination of them to compute,
E
SubLayer (1) (X (I-1)))-
x Σ 11 (1c),
s.t. Σ x =1 (3)
ORF o MOE IR(X(-1)). (4)
Note that X is computed for each sample, and the same expert weights are used for all tokens within a sample.
Theorem 2. The composite map defined in Equa- tion 4 is an injective Lipschitz with fixed (data- independent) upper bound.
A bounded activation bound ensures that the in- cremental learning of our encoder model reduces with large depth, which makes our model scalable to larger depths. It further enforces the impor- tance of learning better embeddings in the initial embedding layer, which drives the entire encod- ing. The runtime complexity of our orthogonal non-normalized attention is O(Nd), whereas dot- product self-attention has runtime complexity of O(N2d d + Ndf). In comparable setting where N >> d, TransJect should have a lower runtime complexity than Transformer.
Corollary 1 (Injectivity of MOE). The mapping function defined in Equation is injective.
3.4 Orthogonal Residual FEN (ORF)
We reformulate the position-wis FFNs with or- thogonal parameterization. FFN layers in Trans- former emulate key-value memory (Geva al., 2021). We enforce Lipschitz continuity the feed- forward sublayer to preserve the layer-wise mem- ory. Formally, we define,
ORF(X(1) 00 oi ELU (XLU(X( (FLUX L
Experimental Setup
W,) bz).
4.1 Tasks and Datasets
With both Wi and W2 being square orthogonal matrices.
We evaluate TransJect and its variants seven short- and long-sequence classification tasks and one language modeling task. We choose the IMDb movie review sentiment classification (Maas et al., 2011) and the AGnews topic classification (Zhang et al., 2015) datasets for short text classification
Corollary 2 (Injectivity of ORF). Orthogonal residual FFNs are injective.
Proof. Using the Lipschitz continuity of ELU, we can prove the corollary directly using Lemma 1.

---

context when the context itself is limited. To high- light the effectiveness of the MOE module, we evaluate an ablation of our model by dropping the MOE module (reducing the number experts to Orthogonal Transformer
Model
IMDb
AGnews
Transformer
81.3
88.8
TransoormerrRReZeero
83.4
89.6
85.1
86.3
Linformer (Wang et al., i
82.8
86.5
1). Dropping the MOE module reduces the val- idation accuracy on the IMDb classification task by 1.7%. On the other hand, using only a single head in the original Transformer model reduces its performance the IMDb classification task by 2.1% %.
Synthesizer (Tay al., 2021)f
84.6
89.1
TeansJect
88.1
88.8
Tranos--T
86.5
90.2
Table 1: Text classification accuracy on IMDb and AG- news (results highlighted with i are taken from Tay al. (2021)).
Long sequence classification. We evaluate TransJect against Transformer along with sev- eral of its recent variants and report the test accu- racy in Table 2. Similar to short sequence clas- sification tasks, TransJect very effective for long sequences and consistently outperforms all the baselines. Out of five tasks, TransJect achieves the best performance in three and beats the best baseline, Skyformer by 0.2% on the leaderboard. TransJect achieves 2.3% and 0.2% better test accuracies on ListOps and byte-level text classi- fication tasks than the corresponding best base- lines, Big Bird and Skyformer, respectively. In- terestingly, Random-TransJJed achieves the best performance on the Pathfinder task, with wide margin of 1.9%. The ListOps task evaluates the ability to learn long-range hierarchical dependen- cies, whereas the Pathfinder task evaluates the ability learn spatial dependencies. As argued by Chen et al. (2021), learning both these dimen- sions poses difficulty for self-attention based meth- ods. With a superior performance on both tasks, TransJect showcases its effectiveness in learning long-tern term patterns from both temporal sequences and sequences with different hierarchical dimen- sions. Moreover, Random-TransJee performs better than TransJect on both Pathfinder and Im- age classification tasks with a margin of 1.1% and 1 1.3%, respectively. We argue that the sparsity in inputs in these two visual tasks is difficult to be approximated by the eigenvalues, which costs Trans Ject in these tasks.
the former one is binary classification task, whereas the latter one contains four classes. To further highlight the effectiveness of TransJect on longer sequences,
evaluate our model on the LRA benchmark (Tay et al., 2020b). LRA benchmark consists of five long sequence classifica- tion tasks - ListOps (Nangia and Bowman, 2018), Byte-level text classification on IMDb review (CharIMDb) dataset (Maas et al., 2011), Byte-level document retrieval on AAN dataset (Radev et al., 2013), Pathfinder (Linsley al., 2018), and Image classification the CIFAR-10 dataset (Krizhevsky and Hinton, 2010). For language modeling, we use Penn TreeBank (PTB) (Mikolov al., 2010) dataset, containing 100M tokens. We provide the details of hyperparameter in Appendix B.I.
4.2 Results
Short sequence classification. Table 1 shows the performance of the competing models. On IMDb classification, TransJect outperforms Transformer with a 6.8% margin. TransJect achieves 3.5% better accuracy than Synthesizer, the best baseline. With zero residual initializa- tion, Transformer can achieve 2.1% better accu- racy in the IMDb classification task. We use an additional ablation of Transformer, in which all the learnable weight matrices are parameterized orthogonally. This improves accuracy on IMDb classification by 3.8%. On the AGnews topic classi- fication task, Bandom-TransJec. achieves 90.2% accuracy, 1.1% better than the best baseline. Inter- estingly,
Language modeling. We report validation and test perplexity on PTB in Table 3 for TransJect and other baselines. Our model achieves 79% lower test perplexity than the vanilla Transformer. As argued by Bachlechner al. (2020), loss-free information propagation is required for training
performs better than TransJect on the AGnews classification task; in- jecting randomness through randomly initialized eigenvalues aids in 1.1% performance improve- ment. Limited contextual information can create difficulty reconstructing X X from the approx- imate eigenvalues. Therefore, having randomly- initialized eigenvalues can aid in learning better
large depth models. Due to this, Transformer with ReZero initialization achieves 75% better perfor- mance than the vanilla Transformer that uses uni-

---

Model
ListOps
Retrieval
Pathfinder
Image
Avg.
Transformer
38.4
61.9
80.7
65.2
40.6
57.4
Transformer+ReZero
38.3
59.1
79.2
68.3
38.6
56.7
Orthogonal Transformer
39.6
58.1
81.4
68.9
42.2
58.0
Reformer (Kitaev et 2020)f
37.7
62.9
79.0
66.5
48.9
59.0
Rig Bird (Zaheer al., 2020)T
39.3
63.9
80.3
68.7
43.2
59.1
Linformer (Wang et al., 2020)f
37.4
58.9
78.2
60.9
38.0
54.7
Informer (Zhou et al., 2021)F
32.5
62.6
77.6
57.8
38.1
53.7
Nystromformer (Xiong et al., 2021)|
38.5
64.8
80.5
69.5
41.3
58.9
Performets (Choromanski et al., 2021))
38.0
64.2
80.0
66.3
41.4
58.0
Skyformer (Chen et al., 2021)f
38.7
64.7
82.1
70.7
40.8
59.4
TransJect
42.2
64.9
80.3
71.5
38.9
59,6
Random-Trand TransJect
40.1
64.6
80.2
72.6
40.2
59.5
Table 2: Test accuracy on the LRA benchmark (results highlighted with are taken from Chen et al. (2021)).
Model
Val
Test
Transformer
636.55
598.93
|
X
Teeeeoo
161.52
147.80
A = ExEi
(5)
|| 1
Orthogonal Transformer
179.98
165.78
TransJect
134.82
127.59
Here X(0 E R is the hidden representation of the sequence X 1th layer.
Random-Trans/ln TransJect
215.68
199.72
Table Perplexity (lower value is better) calculated for language modeling on PTB dataset on validation ("Val") and testing ('Test') splits.
Similarly, we compare the differential entropy (aka entropy) of the hidden representations learned by TransJect at each layer to understand how the entropy state changes across the layers. We calcu- late differential entropy
form residual weights. However, worth not- ing that TransJect achieves 14% better perfor- mance than ReZero initialization, which could be attributed to the inherent injectivity, that ReZero fails to ensure. On the other hand, TransJect with random eigenvalues performs poorly due to its in- ability to encode the inter-dependence between to- kens.
entropy (x(1)) BBi (- log P(X!I))
= -BB P(X 3.h log P(X(i)dn)
(6)
where, P is the empirical probability density func- tion of X. Note that x(I)
x(!)
hith, J.hy leads the
entropy to -O. Therefore, sparser states always have lower entropy and are more deterministic. At the same time, unorderly, random and stochas- tic states have higher entropy. We highlight the distribution of activation factors and entropy of token embeddings at every layer of TransJect and Transformer in Figure 3a. Under Fuclidean and dot-product, TransJect displays an empiri- cal activation bound of Unlike TransJect, Transformer has much higher empirical activation bounds. Although orthogonal parameterization and ReZero lead to much lower activation bounds, they are still higher than our model. Interestingly, Trans- former aims to preserve the semantic similarity at the later layers at the expense of distance; however, frat can preserve both of them with a tighter bound, leading more robust representation for each token. We hypothesize that restricting the distance between a pair of tokens acts as a regu- larization, improving the encoder's final predictive
5 Analysis
To understand the connections behind model depth, activation bounds and entropies, we conduct de- tailed statistical analyses on our model and Trans- former. We use the IMDb and CharlMDb classifi- cation tasks for these studies as part of short long- range classification tasks, respectively. We use the outputs inferred by our models on a subsample of the test data for these analyses.
Activation bounds and entropy. Continuing our initial discussion on preserving layer-wise dis- tances between tokens, we calculate the distribution of activation bounds at different encoding layers. As defined in Section 3.1, compute the activa- tion factor for each encoding layer for TransJect and Transformer. Formally, for 4h layer, we com- pute the activation factor

---

INDB
ChartWob
INDB
CharlMDb
Activatton rockor
Entrapy
Tarsoooe
Entrary
Entrany
Vansject Layer
TrandecLayd
Transformer
Orthoganal Tanddome
(a) Distribution of activation and entropy of models on IMDb and CharlMDb classification tasks.
(b) Attention head/expert level distribution of entropy.
Figure 3: We observe higher activation factors Transformer, whereas the median empirical activation factor for TransJect is & 1. Lower entropy for TransJect indicates that at the neuron level as well as at the expert level, the representations are more orderly.
performance. We computed the Pearson's 's corre- lation between the activation bound of different encoder models (TransJect, Transformer, Orthog- onal Transformer and Transformer-Reieerr) and the final predictive performance and obtained cor- relation score of -0.81 with pvalue 0.09. A sim- ilar correlation score of -..68 is obtained on the CharlMDb classification task. These results high- light the negative linear relationship between the final predictive performance and the average activa- tion bound of the encoders. Therefore, we conclude that a tighter activation bound could lead to better predictive performances on short- and long-range encoding tasks.
itive relationships between the two measures, i.e. having a lower activation bound lowers the overall entropy of the model representations. The same argument can be used to explain the higher entropy in later layers of the Transformer. Our analyses
also confirm the positive correlation between model depth, activation and entropy for Transformer(sse Figure 6 in Appendix C). It is noteworthy that dot- product self-attention computes the query-key sim- ilarity between different tokens. Contrarily, in our proposed attention mapping, we compute the simi- larities between neurons i.e., similarities between different projection vectors, which in turn, enforces our model to learn different projection dimensions and reduces the overall entropy of the system. Ad- report the entropy of the represen- tations at each attention head and expert level in Figure 3b. Although TransJect displays higher entropy at the expert level, is still lesser than the corresponding attention heads of the Transformer model. A sparse badd-balancing capability of the expert model (see Figure in Appendix C) ensures lower entropy throughout the model training. Fur- ther, the changes in inter-quartile ranges in the later layers of the Transformer show increasing random- ness in larger depths, which can also be attributed to the higher activation factors. On the other hand, TransJect stabilises the entropy the later layers.
The average entropy obtained by TransJect on IMDb classification -13 with a standard de- viation of 0.12. On the other hand, Transformer obtains average entropy of 5.80 with stan- dard deviation of 0.13. With higher activation factor, Transformer projects the tokens onto much sparser and random subspaces, increasing the sys- tem's overall entropy. On the other hand, represen- tations learned by TransJect are more orderly and thus have low entropy. Moreover, the entropy of TransJect does not increase perpetually, which ac- cording to the second law of thermodynamics, sug- gests reversible process. We compute Spearman's 's rank and Pearson correlation to understand the rela- tionship between average activation bound and av- erage entropy. We observe Spearman's 's rank corre- lation value of 1.0 on the IMDb classification task, with value of 0.001. The Pearson's correlation also stands at 0.69. These analyses indicate the pos-
ditionally,
Preserving distances between tokens. Figure 4 shows representations obtained on tokens of sample text at different encoder layers, projected onto 2-D. We use isometric mapping (Tenenbaum et al., 2000) for projecting the high dimensional

---

Dranejoat Inmnal Rinbedding
Phaneffcct Larper
Tranejocr Loper
(CA
2RM
1921
Race
120
(a) TransJect
Transtomner hrttal Eibbedding
Transtormer Laayer
Transforer Laper
122
5000
tk
32
04]
hab
(92)
1900
(E27)
brx lestll
330
5000
[CU)
-13X
-1004
-300 -30
-1000
1003
-2000
-1030
-1000
-200
Transformer
Figure 4: Isomap plot of layer-wise token embeddings learned by (a) TransJect and (b) Transformer the "This movie is terrible, but it has some good effects." We highlight four cemantically-releaan words and visualize their positions to understand how different models project tokens onto different subspaces in different encoding layers. TeansJect preserves the relative distances between the tokens and their interconnectedness over the layers, indicating that the projection manifolds are topologically similar.
Speed
model with high activation bound is more irre- versible and, therefore, be less efficient. On the other hand, TransJect exhibits lower entropy, and has higher available energy (from principles ther- modynamics), leading to more efficiency.
Model
1k
2k
3k
4k
Transformer
1.0
1.0
1.0
1.0
TransformerRRZeero
1.0
1.I
Orthogonal Transfommer
1.0
1.0
1.1
1.4
TransJect
5.0
9.6
12.8
26.3
Table 4: Test-time inference speed on long-range text classification tasks various text lengths. The reported numbers are speedups w.r.t. Transformers.
Conclusion
In this work, introduced TransJect, new learning paradigm in language understanding by enforcing a distance prreervepres criterion in the multi-layer encoder models. We derived that by en- forcing orthogonal parameterization and utilizing smoother activation maps, TransJect can preserve layer-wise information propagation within the- oretical bound, allowing the models to regularize inherently. We further argued in favor of injectivity and Lipschitz continuity for better generalization and efficiency. Our empirical analyses suggested a superior performance of TransJect over other self-attention-hant baselines. We observed lower entropy with TransJect with low variance, con- firming the reversible process of statistical mechan- ics. These findings will encourage practitioners to explore the natural laws of science for building better. more intuitive, and cognitive Al.
vectors to the 2-D space. TransJect maintains the initial embedding space throughout the layers, showing robustness in learning initial embeddings. On the other hand, Transformer expands the projection subspaces to more sparse subspaces, even though they project semantically similar tokens closer.
Efficiency comparison. We report the test-time speed on the CharlMDb classification task with different lengths of input sequences in Table 4. Al- beit having 50% more parameters (see Table 5 in Appendix C) than vanilla Transformer, average, we observe 13x speedup for TransJect, which increases to 26x for longer sequences. From the definitions of thermodynamics, a higher entropy leads an irreversible process. This means that a

---

7 Ethical Considerations
Fast Convergence at Large Depth. Number: arXiv:2003.0488 arXiv:2003.0488 [es, stat].
Limitations. As discussed previously, TransJect uses approximated eigenvalues to approximate the gram matrix X X. Therefore, our pro- posed model could be ineffective for sequences with limited context. Additionally, our model usually exhibits higher forward pass time than non-parameterized models due to orthogonal parametrization.
Cenk Baykal, Nishanth Dikkala, Rina Panigrahy, Cyrus Rashtchian, and Xin Wang. 2022. A theoretical view on sparsely activated networks. arXiv preprint arXiv:2208.0440
Iz Beltagy, Matthew E. Peters, and Arman Coham. 2020. Longformer: The Long-Dooe bong-Document Transformer. Number: arXiv:2004.05155 arXiv:2004.05150 [cs).
Graham Brightwell, Claire Kenyon, and Hélène Paugam-Moisy. 1996. Multilayer neural networks: one or hidden layers? Advances Neural Infor mation Processing Systems, 9.
Intended Use and Broader Impact. The empir- ical analyses and insights gathered in this work encourage to explore the dependence between sparsity, entropy and model representations. Our study can be further utilized in developing more efficient larger-depth language models.
Bo Chang, Lili Meng, Eldad Haber, Lars Ruthotto, David Begert, and Elliot Holtham. 2018. Reversible architectures for arbitrarily deep residual neural net- works. In Proceedings of the AAAI conference on artificial intelligence, volume 32.
User Privacy. There is no known user privacy concern with this work.
Yifan Chen, Qi Zeng, Heng Ji, and Yun Yang. 2021. Skyformer: Remodel self-attention with gaussian kernel and nystrom method. Advances in Neural Information Processing Systems, 34:2122-212
Potential Misuse. There no known potential misuse of this work.
Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sar- los, Peter Hawkins, Jared Davis, Afroz Mohiud- din, Lukasz Kaiser, David Belanger, Lucy Col- well, and Adrian Weller. 2021. Rethinking Atten- tion with Performers. Number: arXiv:2009.1477 arXiv:2009.1 1119..14999 14794 [cs, stat].
Reproducibility. We furnish the experi- mental details Appendix B.1 for reproduc- ing the results reported in this paper. We have open-sourced the source code at https: git.
//github.
Djork-Amé Clevert, Thomas Unterthiner, and Sepp Hochreiter. 2015. Fast and accurate deep network learning by exponential linear units (elus). arXiv preprint 1..02215
Environmental Impact. On the language model- ing task, each training iteration takes ~ 0.6 seconds on a single Tesla T4 GPU. Total CO2 emissions are estimated to be 0.06 kg. Estimations were con- ducted using the MachineLearring Impact calcula- tor presented in Lacoste et al. (2019).
Ursula Dicke and Gerhard Roth. 2016. Neuronal factors determining high intelligence. Philosophical Trans- actions of the Royal Society B: Biological Sciences. 37 571(166 (71(565))20055018
Mor Geva, Roei Schuster, Jonathan Berant, and Omer Levy. 2021. Transformer feed-forward layers are key-value memories. Proceedings of the 2021 Conference on Empirical Methods in Natural Lan- guage Processing, pages 5484-5495.
References
Mohsin Ahmed, James Priestley, Angel Castro, Fabio Stefanini, Ana Sofia Solis Canales, Eliza- beth M Balough, Erin Lavoie, Luca Mazzucato. Ste- fano Fusi, and Autia Losonezy. 2020. Hippocampal network reorganization underlies the formation of a temporal association memory. Neuron, 107(2):283- 291.
Aidan N Gomez, Mengye Ren, Raquel Urtasun, and Roger B Grosse. 2017. The reversible residual net- work: Backpropagatton without storing activations. Advances in neural information processing systems, 30.
Dan Hendrycks and Kevin Gimpel. 2016. sian error linear units (gelus). arXiv preprint 17X6::6606.8411
Gaus-
Sanjeev Arora, Yingyu Liang, and Tengyu Ma. 2015. Why deep nets reversible: A simple theory, with implications for training. arXiv preprint arXiv:1511.0565:
Sebastian Jaszezur, Aakanksha Chowdhery, Afroz Mo- hiuddin, Lukasz Kaiser, Wojciech Gajewski, Henryk Michalewski, and Jonni Kanerva. 2021. Sparse is enough scaling transformers. Advances Newral Information Processing Systems, 34:9895-9907.
Thomas Bachlechner, Bodhisatiwa Prasad Majumder, Huanru Henry Mao, Garrison Cottrell, and Ju- lian McAuley. 2020. ReZero is All You Need:

---

Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pap- pas, and François Fleuret. 2020. Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention. Number: arXiv:2006.162
Nikita Nangia and Samuel R Bowman. 2018. Listops: A diagnostic dataset for latent tree learning. arXiv preprint arXiv: 1804.06028.
arXiv:2006. 16236 [cs, stat].
Hao Peng, Nikolaos Pappas, Dami Yogatama, Roy Schwartz, Noah Smith, and Lingpeng Kong. 2020. Random feature attention. In International Confer- encee Learning Representations.
Soheil Keshmiri. 2020. Entropy and the brain: An overview. Entropy, 22(9):917.
Hyunjik Kim, George Papamakarios, and Andriy Mnih. 2021. The lipschitz constant of self-attention. In In- ternational Conference on Machine Learning, pages 5562-5577 PMLR.
Jeffrey Pennington, Samuel S. Schoenholz, and Surya Ganguli. 2017.
Resurrecting sigmoid in deep learning through dynamical isometry: the- ory and practice. Number: arXiv:1711.0433 arXiv: 171111111144555 [cs, stat].
Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. 2020. Reformer: The Efficient Transformer. Num- ber: arXiv:2001.0445 arXiv:2001.0444 14440011144411111
Cindy Poo and Jeffry Isaacson. 2009. Odor repre- sentations in olfactory cortex""sparse" coding, global inhibition, and oscillations. Neuron, 62(6):850-2 861.
[cs, stat].
Christof Koch and Klaus Hepp. 2006. Quantum me- chanics in the brain. Nature, 440(7084):611-61) 611.
Ben Poole, Subhaneil Lahiri, Maithra Raghu, Jascha Sohl-Dickstein, and Surya Ganguli. 2016. Expo- nential expressivity in deep neural networks through transient chaos. Advances in neural information pro- cessing systems, 29.
Alex Krizhevsky and Geoff Hinton. 2010. Convolu- tional deep belief networks on 10. Unpublished manuscript, 40(7):1-9 9.
Ofir Press, Noah A. Smith, and Mike Lewis. 2021. Shortformer: Better Language Modeling using Shorter Inputs. Number: arXiv:2012.15832 arXiv:2012.15832 [cs].
Alexandre Lacoste, Alexandra Luccioni, Victor Schmidt, and Thomas Dandres. 2019. Quantifying the carbon emissions machine learning. arXiv preprint arXiv:1910.09700.
Haozhi Qi, Chong You, Xiaolong Wang. Yi Ma, and Jitendra Malik. 2020. Deep isometric learning for visual recognition. International Conference on Machine Learning, pages 7824-7835. PMLR.
Mario Lezcano Casado. 2019.
Trivializations for gradient-based optimization on manifolds. Advances in Neural Information Processing Systems, 32.
Zonglin Li, Chong You, Srinadh Bhojanapalli, Daliang Li, Ankit Singh Rawat, Sashank Reddi, Ke Ye, Felix Chern, Felix Yu, Ruiqi Guo, et al. 2023. On emer- gence of activation sparsity in trained transformers. International Conference on Learning Representa- tions.
Dragomir R Radev, Pradeep Muthukrishnan, Vahed
Qazvinian, and Amjad Abu-Jbara. 2013. The acl anthology network corpus. Language Rexources and Evaluation, 47:919 19-4449944
Glenn N Saxe. Daniel Calderone, and Leah J Morales. 2018. Brain entropy and human intelligence: A resting-state fimi study. PloS one, 13(2)e0111582.
Drew Linsley, Junkyung Kim. Vijay Veerabadran, Charles Windolf, and Thomas Serre. 2018. Learn- ing long-range spatial dependencies with horizontal gated recurrent units. Advances in neural informa- tion processing systems, 31.
Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoe Le. Geoffrey Hinton, and Jeff Dean. 2017. Outrageously large neural networks: The sparsely-gated mixture-of-expen layer. arXiv preprint arXiv:
Andrew Maas, Raymond E Daly, Peter T Pham, Dan Huang, Andrew Y Ng. and Christopher Potts. 2011. Learning word vectors for sentiment analysis. In Proceedings of the 49th annual meeting of the associ- ation for computational linguistics: Human language technologies, pages 142-150.
Yi Tay. Dara Bahri, Donald Metzler, Da-Cheng Juan, Zhe Zhao, and Che Zheng. 2021. Synthesizer: Re- thinking self-attention for transformer models. In International conference machine learning, pages 10183-10192. PMLR.
Karttikeya Mangalam, Haoqi Fan, Yanghao Li, Chao- Yuan Wu, Bo Xiong, Christoph Feichtenhofer, and Jitendra Malik. 2022. Reversible vision transform- ers. In Proceedings of the IEEF//VE Conference on Computer Vision and Pattern Recognition, pages 10830-10840.
Yi Tay. Dara Bahri, Liu Yang. Donald Metzler, and Da-Cheng Juan. 2020a. Sparse sinkhorn attention. In International Conference on Machine Learning. pages 9438-9447, PMLR.
Yi Tay, Mostafa Dehghani, Samira Abnar, Yikang Shen. Dara Bahri, Philip Pham, Jinfeng Rao, Liu Yang, Sebastian Ruder, and Donald Metzler. 2020b. Long range arena: benchmark for efficient transformers. arXiv preprint 201.0000..
Tomas Mikolov, Martin Karafiát, Lukas Burget, Jan Cer- nockj, and Sanjeev Khudanpur. 2010. Recurrent neu- ral network based language model. Interspeech, volume 2, pages 1045-1048. Makuhari.

---

Joshua B Tenenbaum, Vin de Silva, and John Langford. 2000. A global geometric framework for nonlinear dimensionality reduction. science, 29005500):2319-2323.
long sequence time-series forecasting. In Proceed- ings of the AAAI conference on artificial intelligence. volume 35, pages 11 11100-11111
A Theoretical Results
Ashish Vaswami, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. 2017. Attention All You Need. Number: arXiv: 1706.037222 arXiv:1706.03760 [es].
A.1 Background
We furnish the background materials and proofs of all the theoretical results presented in the main text.
Petra E Vértes, Aaron F Alexander-Bloch, Nitin Gog- tay, Jay N Giedd, Judith L. Rapoport, and Edward T Bullmore. 2012. Simple models of human brain functional networks. Proceedings of the National Academy of Sciences, 109(15):5868-588)
Lemma (Activation bound of linear maps) For matrix M E Ruxm KM is same as the largest absolute singular value, under 1|-1/2.
Voita, Rico Sennrich, and Ivan Titov. 2019. Elena
The bottom-up evolution of representations in the transformer: A study with machine translation and language modeling objectives. arXiv preprint 1xXiv:9009001880
Proof.
|| ||ا||
KM = sup z40
sup ||Mx|| ||]|-=
James Vuckovic, Aristide Baratin, and Remi Tachet des Combes. 2020. A mathematical theory of attention. arXiv preprint arXiv:2007.02876
Squaring both the side, we decompose M as UEVT. where U and V are orthogonal matrices, and E is the diagonal matrix containing the singu- lar values (square root of eigenvalues of M M), Σ1 2 Eg 2 Eg. 2 0. For an orthogonal ma- trix U. |U2||2 ||~ UTUlle = || æ|12 = || This leads to
Sinong Wang. Belinda Z Li, Madian Khabsa, Han Fang. and Hao Ma. 2020. Linformer: Self attentotnn with linear complexity. arXiv preprint arXiv:2006.0 047668
John Wu, Yonatan Belinkov, Hassan Sajjad, Nadir Dur- rami, Fahim Dalvi, and James Glass. 2020. Similarity analysis of contextual word representation models. In Proceedings of the 58th Annual Meeting of the As- sociation for Computational Linguistics, pages 4638- 4655.
| - -IUVVV
x
2
EEx 1
n
Σ
Σ x?
Lechao Xiao, Yasaman Bahri, Jascha Sohl-Dickstein, Samuel S Schoenholz, and Jeffrey Pennington. Dy- namical Isometry and Mean Field Theory of CNNs: How to Train 10.000-Layer Vanilla Convolutional Neural Networks. page 10.
1=1
Hence,
K2
sup En xf=1 i=1
Yunyang Xiong. Zhanpeng Zeng, Rudrasis Chakraborty, Mingxing Tan, Glenn Fung, Yin Li, and Vikas Singh. 2021. Nyströmformer: nyström-bar based algorithm for approximating self-attention, In Proceedings of the AAAI Conference on Artificial Intelligence, vol- ume 35, pages 14138-14144
Being a convex sum, K3 < Ex, which completes the proof.
Corollary (Largest eigenvalue of a square stochastic matrix) The largest absolute value of any eigenvalue of square stochastic matrix is equal to 1.
Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago On- tanon, Philip Pham, Anirudh Ravula, Qifan Wang. Li Yang, et al. 2020. Big bird: Transformers for longer sequences. Advances in Neural Information Processing Systems, 33: 17283
Proof. For any square stochastic matrix M guxn
Xiang Zhang. Junbo Zhao, and Yann LeCun. 2015. Character-level convolutional networks for classi- fication. Advances in neural information processing systems, 28.
1=1 M₁1
1
M,2
Haoyi Zhou, Shanghang Zhang, Jieqi Peng, Shuai Zhang, Jianxin Li, Hui Xiong, and Wancai Zhang. 2021. Informer: Beyond efficient transformer for
MI
= I.
Σド=1 Myn

---

Hence, is an eigenvalue of M. Next, we prove that is the largest eigenvalue of M. For any eigenvalue À and its corresponding eigenvector v. Au Mu. Without loss of generality, we assume arg max; |ua| = 1. Hence, for the 1st entry in this column vector λo₁ = ンド= My. Using trian- gle inequality we get.
A.2 Main Results
A.3 Proof of Theorem
Here we use the linear attention (Katharopoulos et al., 2020) with the dot-product similarity be- tween Q and K. Being real symmetric matrix, X X can be decomposed into QEQ", which leads
Alleer S laaa
Σ j=1
M1 S Σ
Misllei
Attention( (iriion((,K x wowks k 2 E
/=1
$ |mil-
@
(9)
orthogonal
diagonal orthogonal
Hence, for any eigenvalue |^| - 1. Hence, proves our corollary that the largest eigenvalue is 1.
As the product of two orthogonal matrices is orthogonal, Equation 9 reduces to
Lemma 3 (Lipschitz bound for continuously differentiable functions). Any C function f IR" + IR' with bounded derivative has Lipschitz bound as supæ vvll].
Attention Q, K.K. V) - XUEV
(10)
with suitable set of learnable orthogonal matrices U and V, and Σ being the matrix containing the eigenvalues of X X.
Proof. For any v. y C R", we define g : [0, 1]
A.4 Proof of Lemma
R" as
(7) Let us assume Eæ * y such that f()) = f(y). which implies [x) s(y)ll = O. Using triangle inequality and Equation 2 we get
s(t) = f(x + t(y x)).
It is easy to verify that g(0) f (xx and g(1) f(y).
a (F(x)) F(y)) L
x y f f
- - (x))) = - (ty(t) (8)
This implies.
Using chain rule of differentiation on Equation 7, we get,
x - y/l
04 1 (1F((((x))
$ [lf(x)) f(0)]++11
-F(y)l
Vig(t) Vef Vef(I t(y z)) (y x).
Therefore,
Using this in Equation 8, we get
x vl1 [1F(z) - F(v)l
(1) - f(x)ll
I [(()) F(y)l-
tly - x) (y (y- xd = =11/ /' Vel vel(20 x
Here F(X) = EXU X XW for a suit- able choice of orthogonal weights W and WV. Although it is difficult to analytically calculate the matrix derivative of XW X X XWV w.r.t. X, in- tuitively,
$ //' V t(y 2)) 1-- )dl
As f is Ct, the supremum of its derivative exists. allowing us to set supt V., x t(y - x)) K and deduce
calculate that X W W O X x and thus the norm of the input- output Jacobian is bounded by 3 ||x|| Here we can ignore the weight matrices, as they are orthogo nally parameterized and therefore has ||-I| = 1. As x" X is stochastic, the norm of the input-output
[[f()))(()]I S K11y-=|1 /' dt K K11y-z||

---

Jacobian is bounded by Using Lemma 3, we can calculate the Lipschitz bound of ELU activation as 1. Therefore, using the chain rule of calculus and Lemma 3, we can deduce the Lipschitz bound for as 3. Using this, we get
A.6 Proof of Theorem
The original input space of the composite mapping the concatenated token embedding that lies [-1, 1] as both orthogonal embedding and posi- tion encoding functions have the function image (codomain) in -], As the input space is com- pact, any continuously differentiable function is Lipschitz. Therefore, the composite function is au- tomatically Lipschitz continuous. Using Lemma and the fact that the composition of multiple in- jective functions is injective, we can deduce that the encoding function is injective for any encoder with 7. 2 3. The dynamical isometry property can also prove the Lipschitz continuity of the encoding function irrespective of the number of encoding lay- ers. Although ReZero only guarantees dynamical isometry during initialization, we can observe the distribution of residual weights throughout model training. We highlight the distribution of residual weights of TransJect for IMDb and CharlMDb classification tasks in Figure 5. Residual weight values remain 0.33, highlighting that the residual connections remain injective.
3 - yl < L
|| - y/1.
This contradicts that < 1 for an encoder with L 2 3. Hence, we prove the lemma by contradic- tion.
A.5 Proof of Corollary 1
Let assume ((()) + ( such that
Σ Aex (11)
ai
Σ
de Fo(all111
L
E
Σ Aey 1) +
01
Fe V).
L
Using Σ de- 1 obtain,
Next, we compute the Lipschitz bound for f. For the sake of simplicity, let us assume + = 4 % = t. We first expand f as
|| 1) n 11l
01 =11, L
E
de (z (1-1))
04 L
Σ
E
de (y
E
E
f(x)
Σ
A Σ
Xa ELU(
e=1
6=1
a ELU(T).
Therefore,
|| y (1-1))
E == ELU
A ((-1)))))
Σ
(t
X
Σ Xa FIU + bi W2 + b2.
Using Lipschitz bound of Fo(xo((-1) 3, we obtain
This implics,
(1-11-1)) 1 (11-1) <
3
Σ
Aellaz (0-1 "1-
-- ELU xWi+
I
6=1
3
[1
(-1)
-y
y 11-1)
Aea EIU + b1 Wa + b2.
L
This contradicts that fact that < 1. Hence, we prove the corollary by contradiction.
Using the Lipschitz bound of F(X) from the
proof of Lemma and the fact Σ X = 1, we

---

JMDb
CharlMDb
0.00
-0.01
Transject Layer
Transject Layer
Figure 5: Distribution of Residual Weights in different layers of TransJect
get
layer. We utilize the BERT pre-trained tokenizer to tokenize the texts for these two tasks, We follow the evaluation methodologies followed by Chen et al. (2021). In all these classification tasks, we keep the configuration of our models with 2, E 4, and 512. Except for the text classifica- tion task, we use max pooling on the hidden repre- sentation obtained from the final encoder layer be- fore feeding to the final classification feed-forward layer. For the PTB language modeling task, we maximum sequence length of 35 in both the input and the output. For all the models, we set L 6, and d 512. All IMDb and AGNews classification experiments are run for epochs. We use an early stopping based on the test loss with the patience 4 epochs to terminate learn- ing on platcaus. For both these experiments, we use Adam optimizer with a learning rate of 0.0005, 31 0.9, 82 0.98 and 10-9. We use both training and test batches of size 32. We follow the implementation details shared by Chen al. (2021) for the LRA benchmark. All these experiments are run for 50k steps. We use Adam optimizer with a learning rate of 0.01 for the language modeling task and trained the models for 10 epochs. One Tesla T4 and one Tesla V100 GPU are used for conducting all the experiments. For each task, we report the average performance across three different runs.
| 7/
< (1 y|1 ||ill + 3|a|| -1 yll) [[
<
(1 3|a|| yll.
Finally,
[lf(x)) f(y)l
$ || vll 3/0| || y/
+3/a + 3/a)))) - yl
<1 +/al)>²l yl
3
<1 + L 311 y/l-
Here, / is layer-wise operator. Hence, for all the I encoder layers,
[15(x) f(y))l < (1+ + ) 2 yll-
Using lim (1 + 1)" e, get fx:) (w)ll < < - 3/1. Although the theoretical Lip- schitz bound is very high, the empirical Lipschitz bound much smaller due to dynamical isometry. The Lipschitz bound also independent of the in- put data, making theoretically valid. However, as demonstrated in Figure 5, the residual weights are empirically closer to 0, even during the model train- ing, due to which the empirical Lipschitz bound is 1 6
enforce orthogonality during the forward pass and after backpropagation, we use PyTorch pa- rameterization This implementation uses differ- ent orthogonal maps (e.g. houscholder or Cay- ley mapping) to ensure orthogonality. It further the framework of dynamic trivialization (Lez- cano Casado, 2019) ensure orthogonality after gradient descent. We use the implementation by
B Experimental Setup
B.1 Hyperparameter settings
For IMDb and AGnews classifications, we choose maximum text length of 512. In all these two classification tasks, keep the configuration of our models with L 6. E 4, and d 512. For these two tasks, we use mean pooling on the hidden representation obtained by the final encoder layer before passing it to the final classification
/tgp//bes//stalll//eeeeaetee torch. nn. utils. pww.metrizationnaee. orthogonal hhml

---

(a) TransJect on IMDb classification task
(b) TeansJect on CharlMDb classification task
Drepout
Nam Japas
Ioteen Soe
Pur Jrpa
middin Siee
(c) Transformer on IMDb classification task
(d) Transformer on CharlMDb classification task
Figure 6: Distribution of activation factors and entropy with different model configurations of TransJect and Transformer on IMDb and CharlMDb classification tasks.
125
100
J0
Donsity
75
50
0.0
0.2
0.4
0.6
0.8
1.0
8.248
0.249
6 250
0.251
0.252
0.253
0.254
Transject Expert Weights
Trangject Expert We ghiss
Expert
Expert
Expert
Expert
(a) IMDb classification
(b) CharlMDb classification
Figure 7: Expert weights on (a) IMDb and (b) CharlMDb classification tasks. The mode around 0.25 highlights the inherent load- hud-bblanning capabilities among the experts.
Scikit-learn' that uses cosine_sinin sissinesimmlarity as the cosine distance. Cosine similarity is the co- sine the angle between two vectors which lies in [-1, 1]. Thus the distance lies in [0, 2].
C Analysis
Connections between entropy and model configuration. We highlight the connections between activation bounds and total entropy under different model configurations in Figure 6 for TransJect and Transformer, respectively. We observe that more experts in MOE increase activation at the final layer. On the other hand, with a larger dropout for Transformer, the model regularizes more and creates more sparsity, leading to lower activation factors and entropy. Contrary to TransJect, increasing the number of layers of the Transformer increases activation. This behaviour indicates the layer-agnostic property of our model.
Model
#Parameters
Transformer
190.1M
Transformer with dęf 512
95.6M
TransJect (E 1)
66.2 2M
TransJect (E 8)
286.8M
Table 5: Comparison of number of learnable pa- rameters and inference speed on the CharlMDb classi- fication task. The number of parameters is computed for encoders with L. and 512. Transformer uses hidden size dæf 2048 in the feedforward layer, whereas TransJect uses the original encoder hidden dimension maintain invertibility.
Effectiveness of the expert model. We observe the individual importance of experts and how they interplay in the mixture model. We illustrate the expert weight distribution in Figure 7, confirming
'https:/ www.:/laww..rrr/saabee/mmuuueeee generated/skleegt metrics. pairwiseeee cosine. distances.htie html

---

that the experts are well-balanced and that our model does not require an enforeed load balancing. Further, we calculate the entropy of each expert representation to understand how it constituents the final representation at every layer. The expert entropy value of TransJect is 2.2. Computing equívalent entropy of different attention heads for the Transformer gives us an entropy of 0.67. The lower entropy displayed by the mixture of experts highlights the generalization capabilities of the experts. Different experts learn differently but orderly, unlike random and sparse learning by the attention heads of Transformer models.
Parameter comparison. We compare the mod- els in terms of total learnable parameters in Table 5. In equivalent setup, TransJect has 50% more learnable parameters than Transformers.


Reviews:
[{"review_id": "1rXHNpEFIL", "reviewer": "Reviewer_p61L", "summary": "", "questions": "A. Why do we need to keep the value of distance consistent? This paper does not cite related papers or conduct its own studies to prove that.\n\nB. Why is the FFN module called Orthogonal FFN? It seems not to use orthogonal.\n\nC. How much improvement does the MOE bring? Is the MOE specific to your model design?", "limitations": "", "rating": 3, "confidence": 4, "soundness": 3, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank you for the valuable feedback. We address the raised concerns below.\n\n**R1. (weak baselines; s4 not used as a baseline)** We would like to reiterate that the primary motivation of our work is to demonstrate the importance of Lipschitz continuity and injectivity within the Transformer architecture. As the state-space models (S4) do not fall in the Transformer family, we do not compare them with our model. S4 and other state-space models show different inductive biases than TransJect. Therefore, they are not comparable and are excluded from this study. \n\n**R2.** As explained in Section 1 (lines 115-134) of our paper, preserving distance between a pair of tokens within a certain analytical bound ensures that the manifold structure between different tokens is not distorted in larger depths. Our empirical analysis highlights the importance of injectivity and Lipschitz continuity, and the significant reduction in test-time inference time justifies the need for these mathematical properties. Our empirical analysis also highlights the advantages of having low activation bound on the final predictive task. We have computed the Pearson's correlation between the activation bound of different encoder models (TransJect, Transformer, Orthogonal Transformer and Transformer+ReZero) and the final predictive performance and obtained a correlation score of $-0.81$ with $p$-value $0.09$. A similar correlation score of $-0.68$ is obtained on the CharIMDb classification task. These results highlight the negative linear relationship between final predictive performance and the average activation bound of the encoders. Therefore, we conclude that tighter (lower) activation bounds can lead to better predictive performance on short- and long-range encoding tasks.\n\n**Novelty of the paper** We do not claim the novelties with MOE, orthogonal attention or learnable residuals. Our novelty is that we analytically show the Lipschitz continuity of our modified Transformer architecture by unifying these components. Furthermore, reformulating self-attention as $XU\\Sigma V$ is novel and essential to ensure the Lipschitz condition. To our knowledge, ours is the first work ensuring injectivity within the Transformer family. Although a few earlier works explored the Lipschitz continuity of Transformers, ours is one of the few studies that have computed the theoretical Lipschitz bound independent of the model configuration. \n\n**Q A.** As explained in Section 1 (lines 115-134) of our paper, preserving distance between a pair of tokens within a certain analytical bound ensures that the manifold structure between different tokens is not distorted in larger depths. Our empirical analysis also highlights the advantages of having low activation bound on the final predictive task. We have computed the Pearson's correlation between the activation bound of different encoder models (TransJect, Transformer, Orthogonal Transformer and Transformer+ReZero) and the final predictive performance and obtained a correlation score of $-0.81$ with $p$-value $0.09$. A similar correlation score of $-0.68$ is obtained on the CharIMDb classification task. These results highlight the negative linear relationship between final predictive performance and the average activation bound of the encoders. Therefore, we conclude that tighter (lower) activation bounds can lead to better predictive performance on short- and long-range encoding tasks.\n\n**Q B.** Orthogonal parameterization ensures a Lipschitz bound of $1$, which is required to ensure that the TransJect Lipschitz bound is layer-independent. Therefore, we use an Orthogonal FFN instead of the vanilla non-parameterized FFN.\n\n**Q C.** MOE is very specific to our proposed model. A convex sum of different experts is required to ensure the injectivity of our encoder. The concatenation operation usually used in the multi-head attention is not injective. Dropping the MOE (i.e., the number of experts is 1) drops the validation accuracy on the IMDb classification task by $4.7%$%. On the CharIMDb classification, MOE contributes an improvement of $1$%.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "This paper connects common-used modules and the manifold-preserving abilities.", "reasons_to_reject": "1. The baselines are weak. This paper does not compare s4 in the LRA benchmark, which achieves much higher results. On text classification and language modeling, this paper also only compares with a few weak baselines.  Moreover, the proposed method uses learnable residual, elu,  and MOE.\n\n2. The contribution brought by this paper is limited.  First, it does not prove that we need to keep the value of distance between elements consistent. Second, orthogonal attention, moe, and residual with learnable weight are both known; what is the unique contribution of this paper?", "questions_to_authors": "A. Why do we need to keep the value of distance consistent? This paper does not cite related papers or conduct its own studies to prove that.\n\nB. Why is the FFN module called Orthogonal FFN? It seems not to use orthogonal.\n\nC. How much improvement does the MOE bring? Is the MOE specific to your model design?", "paper_topic_and_main_contributions": "This paper proposes a Transformer variant that keeps the distance between elements consistent among layers and heads. This paper achieves its goal by combining orthogonal attention and residual with learnable weight. For experiments, this paper validates the proposed method on a long-range arena, two small-range text classification tasks, and a small language modeling task.", "excitement": 3, "reproducibility": 3, "ethical_concerns": "No", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "7gn1qff8eI", "reviewer": "Reviewer_9LuQ", "summary": "", "questions": "As stated above, when do the authors think the assumption (topology shall remain similar across layers) will be true, and when it might fail?\n\nAlso, how does the activation bounds relate to eventual performance?", "limitations": "", "rating": 4, "confidence": 3, "soundness": 4, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank you for the valuable feedback. We address the raised concerns. \n\n**R1.**  We would like to clarify that our intention is not to restrict the distances between a pair of tokens across layers but to bind them within a theoretic bound. Therefore, the model still learns different semantic relationships and topological properties but preserves the overall manifold structure of the embedding space within a specific theoretical bound. This also results in a lower activation bound for our model.\n\nWe have computed the Pearson's correlation between the activation bound of different encoder models (TransJect, Transformer, Orthogonal Transformer and Transformer+ReZero) and the final predictive performance and obtained a correlation score of $-0.81$ with $p$-value $0.09$. A similar correlation score of $-0.68$ is obtained on the CharIMDb classification task. These results highlight the negative linear relationship between final predictive performance and the average activation bound of the encoders. Therefore, we conclude that tighter (lower) activation bounds can lead to better predictive performance on short- and long-range encoding tasks. \n\n**Q1.** We think the hypothesis of bounded activation bound should hold in any encoding task. However, the activation bound might change in different semantic and pragmatic tasks, depending upon the necessity of learning different topological properties within texts. For instance, we observe an average activation bound of $1.04$ for TransJect in the IMDb classification task but an activation bound of $1.0$ in CharIMDb classification.\n\n**Q2.** We have computed the Pearson's correlation between the activation bound of different encoder models (TransJect, Transformer, Orthogonal Transformer and Transformer+ReZero) and the final predictive performance and obtained a correlation score of $-0.81$ with $p$-value $0.09$. A similar correlation score of $-0.68$ is obtained on the CharIMDb classification task. These results highlight the negative linear relationship between the final predictive performance and the average activation bound of the encoders. Therefore, we conclude that a tighter activation bound could lead to better predictive performances on short- and long-range encoding tasks.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "Many analysis in this paper are interesting, and the proposed method to preserve the similarity topology across layer is also interesting and promising. If this method indeed works generally for all the tasks, it will be come a good alternative to current Transformer building blocks.", "reasons_to_reject": "The key assumption of this paper \"the token similarity across layers shall remain the same\", might not be correct for all NLP tasks. I think in many situtations, the attention for different layers are trying to learn different semantic relationship and meaning. Some papers also study extracting some syntax tree from attention of different layers: (e.g., https://aclanthology.org/W18-5444/)\n\n\nI'm interested to hear more from the authors, when their suggested assumption shall work, and at which cases (and tasks) the assumption as well as the proposed method might fail.", "questions_to_authors": "As stated above, when do the authors think the assumption (topology shall remain similar across layers) will be true, and when it might fail?\n\nAlso, how does the activation bounds relate to eventual performance?", "paper_topic_and_main_contributions": "This paper tries to improve Transformer attention. They have a assumption that the pairwise similarity relationship shall remain similar across layers. Based on this assumption, they propose to replace the dot-product based attention with a orthogonal decomposition (Attn(X) = XU E V).\n\nThis ensures injective mappings of token representation across layers.\n\nThey show higher performance on several langauge model benchmarks, and give a analysis of multi-head self-attention from statistical physics viewpoint (they have higher empirical activation bounds compared with the proposed TransJect)", "excitement": 4, "reproducibility": 3, "ethical_concerns": "No", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"Missing_References": "NA", "Typos_Grammar_Style_and_Presentation_Improvements": "NA"}, "extra_text": {}}, {"review_id": "pwKC9nmG7v", "reviewer": "Reviewer_Vm5s", "summary": "", "questions": "- What is the relationship between the preservation of layer-wise distance between a pair of tokens and the classification of sequences into short and long lengths? Does the proposed method gain a better representation? There is still a gap between representation across tokens and sentence representation. \n- In line 330, the authors argue that a transformer layer project representations to a lower-dimensional space, does it refer to the multi-head attention operation? Does MoE perform better than single head attention without a lower dimension?  The incorporation of Mixture of Experts (MoE) appears to primarily focus on dimension preservation, whereas there exist numerous alternative techniques that also prioritize the preservation of dimensions.\n- How the learnable weight, such as $\\alpha$ in equ. 2 and $\\lambda$ in equ. 3 show? Could you give more analysis for those weights?\n- How to ensure the $W_1$ and $W_2$ matrices to be square orthogonal? \n- How the pair-wise distance across token connect to entropy is also not clear.", "limitations": "", "rating": 3, "confidence": 3, "soundness": 4, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank you for the valuable feedback. We address the raised concerns. \n\n**R1.** We request the reviewer to elaborate on the confusion, as it is difficult for us to address this comment without much clarity. Also, lines 95-114 of our paper emphasize the motivation of using entropy to evaluate different encoder models analytically. As stated, entropy is a key metric widely used in cognitive and physical sciences to quantify a state's randomness. We use entropy to evaluate the randomness within different token representations obtained by different encoder models.\n\n**R2.** In this work, we propose a non-normalized orthogonal linear attention and devise an injective variant of the Transformer. Our proposed model TransJect is shown to be Lipschitz continuous with a fixed Lipschitz bound. We observe low activation bound and low entropy for our model, highlighting its effectiveness in encoding short- and long sequences. Our model exhibits significantly lower inference time due to low entropy, highlighting its efficiency in encoding longer sequences. \n\nWe have already highlighted the ablation study in Section 4.2 of our paper. Table 1, 2 and 3 highlights the ablation results. We want to clarify that the models -- Orthogonal Transformer and Transformer+ReZero are two ablations we have used to highlight the importance of different components. Orthogonal Transformer uses Orthogonal parameterization but does not use learnable residual. On the other hand, Transformer+ReZero uses learnable residuals but does not use Orthogonal parameterization in any attention or FFN layer. Additionally, we have computed another ablation with only a single expert in the MOE module. This ablation of TransJect highlights the effectiveness of the MOE module. Dropping the MOE module reduces the validation accuracy on the IMDb classification task by $4.7$%. On the CharIMDb classification, MOE contributes an improvement of $1$%. \n\n**Q1.** We hypothesize that restricting the distance between a pair of tokens acts as a regularization, improving the encoder's final predictive performance. We computed the Pearson's correlation between the activation bound of different encoder models (TransJect, Transformer, Orthogonal Transformer and Transformer+ReZero) and the final predictive performance and obtained a correlation score of $-0.81$ with pvalue $0.09$. A similar correlation score of $-0.68$ is obtained on the CharIMDb classification task. These results highlight the negative linear relationship between the final predictive performance and the average activation bound of the encoders. Therefore, we conclude that a tighter activation bound could lead to better predictive performances on short- and long-range encoding tasks.\n\n**Q2.** We would like to clarify that the primary motivation behind introducing MOE is not to preserve dimension but to ensure injectivity within our encoder model. Concatenation operator, usually used in multi-head self-attention, is not an injective operation. Whereas Corollary 1 of our paper shows that the convex sum of different experts leads to an injective mapping. \n\nTo highlight the effectiveness of the MOE module, we evaluate an ablation of our model by dropping the MOE module (reducing the number of experts to 1). Dropping the MOE module reduces the validation accuracy on the IMDb classification task by $4.7$% (accuracy $84$%). On the other hand, using only a single head in the original Transformer model reduces its performance on the IMDb classification task by $2.1$% (accuracy $79$%).\n\n**Q3.** We have already provided the analysis on $\\alpha$, the residual weight and $\\lambda$, the expert weights in Figure 4 and Figure 6 in the Appendix of our paper. Section 5 of the main manuscript also describes the statistical properties of these variables.\n\n**Q4.**  Appendix B.1 of our paper already describes the mechanism we adopted to ensure orthogonality. We use the PyTorch parameterization library to ensure the orthogonality of these learnable parameters. The implementation uses orthogonal maps (e.g. householder mapping) and dynamic trivialization to ensure orthogonality during the forward and backward pass.\n\n**Q5.** We compute Spearman's rank and Pearson correlation to understand the relationship between average activation bound and average entropy. On the IMDb classification task, we observe a Spearman's rank correlation value of $1.0$, with pvalue of $0.001$. The Pearson's correlation also stands at $0.69$. These analyses indicate the positive relationships between the two measures, *i.e.* having a lower activation bound lowers the overall entropy of the model representations.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "- The paper is well motivated.\n- The paper study an interesting question and conduct extensive experiments.", "reasons_to_reject": "- The paper is loosely organized  and the correspondence from their theory to their methods is not clear. The paper introduce too many redundant materials, such as line 95 to line 114.  \n- The paper combine a series of techniques and the truly contribution is not clear. The ablation study for each component should be given to demonstrate their effectiveness.", "questions_to_authors": "- What is the relationship between the preservation of layer-wise distance between a pair of tokens and the classification of sequences into short and long lengths? Does the proposed method gain a better representation? There is still a gap between representation across tokens and sentence representation. \n- In line 330, the authors argue that a transformer layer project representations to a lower-dimensional space, does it refer to the multi-head attention operation? Does MoE perform better than single head attention without a lower dimension?  The incorporation of Mixture of Experts (MoE) appears to primarily focus on dimension preservation, whereas there exist numerous alternative techniques that also prioritize the preservation of dimensions.\n- How the learnable weight, such as $\\alpha$ in equ. 2 and $\\lambda$ in equ. 3 show? Could you give more analysis for those weights?\n- How to ensure the $W_1$ and $W_2$ matrices to be square orthogonal? \n- How the pair-wise distance across token connect to entropy is also not clear.", "paper_topic_and_main_contributions": "This paper study how to preserve similarities in transformer and propose a new method for attention. The authors replace the attention mechanism with an orthogonal decomposition $XU\\Sigma V$for the input $X$.\n The authors argue that the representation will be injective and low entropy.", "excitement": 3, "reproducibility": 4, "ethical_concerns": "No", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper introduces a Transformer variant that maintains the distance between elements consistent across layers and heads. The paper achieves this by combining orthogonal attention, residual, and learnable weight. The method is validated on LRA and few small datasets.

Pros:
- The paper proposes to combine orthogonal attention with residual networks and analytically show the Lipschitz continuity of their modified Transformer architecture. Which is an interesting idea.

Cons:
- It is unclear why “preserving distance between a pair of tokens within a certain analytical bound” is actually useful or needed. This makes the objective of the paper unclear. 
- In the absence of clearly motivated objectives, it is unclear as to why Lipschitz continuity is required (as shown by the authors for their model).
- Actual contribution of the paper is unclear
- The reason behind a >5x improvement of speed while 50% more parameter is not clear and not analyzed in the paper.
Decision: Accept-Findings