Manifold-Preserving Transformers are Effective for Short-Long Range Encoding (2023)

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.

## introduction over the past few decades, deep neural networks have greatly improved the performance of various downstream applications. stacking multiple layers has been proven effective in extracting features at different levels of abstraction, thereby learning more complex patterns @xcite . since then, tremendous efforts have been made in building larger depth models and in making them faster @xcite . self-attention-based transformer model @xcite was proposed to parallelize the computation of longer sequences; it has achieved state-of-the-art performance in various sequence modeling tasks. following this, numerous efforts have been made to reduce computation and make the transformer model suitable even for longer sequences @xcite @xcite @xcite @xcite . however, very few of these studies discuss information propagation in large-depth models. a recent study @xcite characterized how transformer token representations change across layers for different training objectives. the dynamics of layer-wise learning are essential to understand different abstract levels a model could learn, preserve and forget to continue learning throughout the layers. however, recent studies need to shed more light on how transformers preserve contextual information across different layers @xcite . to understand how transformer encodes contextual similarities between tokens and preserves the similarities layer-wise and across different attention heads, we highlight an example in figure 1 . we select three pairs of semantically-related tokens. we observe the euclidean and cosine distances of the representations learned at different layers of the transformer trained on the imdb sentiment classification 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 representations 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 capabilities of deep neural models, which transformer evidently fails to do. neuroscientists have been working for years to understand how the human brain functions and simulates the behaviours of physical sciences @xcite . 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-connectedness (synapses) rather than the size of the brain, or the number of layers through which the information propagates @xcite . although a fully-grown human brain can have hundreds of billions of neurons, it has been found that at a time, only a tiny fraction of neurons fire @xcite . this sparse firing can help the brain sustain low entropy (energy). entropy, a measure that quantifies a state's randomness, has thus become an essential tool to understand the factors behind human intelligence @xcite and how the human brain operates. unfortunately, transformer and its variants are yet to be studied from these interdisciplinary viewpoints. our study explores the connection between model sparsity and entropy. towards this, we propose a complete redesign of self-attention-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 regularization through multi-head self-attention and dropout, transject does not require explicit regularizers and can be regularized implicitly due to its inherent injectivity. the backbone of transject is a non-normalized linear orthogonal attention and an injective residual connection; both ensure lipschitz continuity. by enforcing lipschitz continuity in euclidean and 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 justify 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 long-sequence classification tasks, respectively. our model performs best on the long sequence classification (lra) benchmark, achieving 0.2% better accuracy than the best baseline, skyformer @xcite . we further demonstrate transject on language modeling task on the penn treebank dataset, in which our model achieves 79% better test perplexity than the vanilla transformer. empirical analyses suggest a very low entropy of transject representations, indicating that transject captures sparse and more orderly representations compared to transformer. moreover, transject shows 13× lower inference runtime than transformer, indicating its efficiency in encoding long input sequences. foot_0 ## related works despite being a ubiquitous topic of study across different disciplines of deep learning, transformer @xcite models still require better mathematical formalization. on a recent development, @xcite formalized the inner workings of self-attention maps through the lens of measure theory and established the lipschitz continuity of self-attention under suitable assumptions. however, the lipschitz condition depends on the boundedness of the representation space and the lipschitz bound of the fully-connected feed-forward layer (ffn). a similar study @xcite also concluded that the dot-product self-attention is 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 always been an active area of study @xcite @xcite . a reversible network ensures better scaling in large depths and is more efficient than non-reversible structures. recently, @xcite designed 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. dynamical isometry is a property that mandates the singular values of the input-output jacobian matrix to be closer to one. @xcite showed that dynamical isometry could aid faster convergence and efficient learning. following their idea, @xcite showed that residual connections in transformers do not often satisfy dynamical isometry, leading to poor signal propagation 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 transformers. previously, @xcite enforced a stronger condition of isometry to develop deep convolution networks efficiently. however, isometric assumptions are not valid for sequential modeling due to different levels of abstraction within the input signals. moreover, isometry may not hold between contextually-dissimilar tokens. another essential aspect behind designing efficient large-depth models is ensuring model sparsity. several notable contributions have been made @xcite @xcite to enforce sparse activations within transformers to make them more efficient and scalable. @xcite argued that sparse networks often resemble the sparse activations by the human brain, bearing a similarity between artificial and biological networks. they empirically 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 inherently regularized and display lower entropy, projecting the representations onto a sparse subspace pushes them further from being lipschitz, preventing them from reversible. this work proposes an injective and lipschitz continuous alternative to vanilla transformer, namely, transject. with the enforced injectivity, our model establishes a theoretical guarantee for reversibility and thus can be scaled to larger depths. further, transject displays significantly lower entropy than transformer, indicating a lower energy footprint and more efficiency. ## designing injective transformer this section formally describes our proposed model, transject. it inherits the structure from the vanilla transformer and achieves a smoother activation plane by utilizing injective maps for transforming token representations across layers. for an l-layered stacked encoder, we aim to learn the representation of a sequence @xmath0, x 2 , • • • , x n } at each layer l 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. ## analysis to understand the connections behind model depth, activation bounds and entropies, we conduct detailed statistical analyses on our model and transformer. we use the imdb and charimdb classification tasks for these studies as part of short longrange 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 distances between tokens, we calculate the distribution of activation bounds at different encoding layers. as defined in section 3.1, we compute the activation factor for each encoding layer for transject and transformer. formally, for l th layer, we compute the activation factor here x (l) ∈ r n ×d is the hidden representation of the sequence x at l th layer. 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 calculate differential entropy as, the average entropy obtained by transject on imdb classification is -1.3 with a standard deviation of 0.12. on the other hand, transformer obtains an average entropy of 5.80 with a standard deviation of 0.13. with a higher activation factor, transformer projects the tokens onto much sparser and random subspaces, increasing the system's overall entropy. on the other hand, representations learned by transject are more orderly and thus have low entropy. moreover, the entropy of transject does not increase perpetually, which according to the second law of thermodynamics, suggests a reversible process. we compute spearman's rank and pearson correlation to understand the relationship between average activation bound and average entropy. we observe a spearman's rank correlation value of 1.0 on the imdb classification task, with a p-value of 0.001. the pearson's correlation also stands at 0.69. these analyses indicate the pos-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(see figure 6 in appendix c). it is noteworthy that dotproduct self-attention computes the query-key similarity between different tokens. contrarily, in our proposed attention mapping, we compute the similarities 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. additionally, we report the entropy of the representations at each attention head and expert level in figure 3b . although transject displays higher entropy at the expert level, it is still lesser than the corresponding attention heads of the transformer model. a sparse load-balancing capability of the expert model (see figure 7 in appendix c) ensures lower entropy throughout the model training. further, the changes in inter-quartile ranges in the later layers of the transformer show increasing randomness in larger depths, which can also be attributed to the higher activation factors. on the other hand, transject stabilises the entropy in the later layers. preserving distances between tokens. figure 4 shows representations obtained on tokens of a sample text at different encoder layers, projected onto 2-d. we use isometric mapping @xcite for projecting the high dimensional 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 charimdb classification task with different lengths of input sequences in ## conclusion in this work, we introduced transject, a new learning paradigm in language understanding by enforcing a distance-preserving criterion in the multi-layer encoder models. we derived that by enforcing orthogonal parameterization and utilizing smoother activation maps, transject can preserve layer-wise information propagation within a theoretical 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-based baselines. we observed lower entropy with transject with low variance, confirming the reversible process of statistical mechanics. these findings will encourage practitioners to explore the natural laws of science for building better, more intuitive, and cognitive ai.