A Neural ODE Interpretation of Transformer Layers
Abstract
Transformer layers, which use an alternating pattern of multi-head attention and multi-layer perceptron (MLP) layers, provide an effective tool for a variety of machine learning problems.
As the transformer layers use residual connections to avoid the problem of vanishing gradients, they can be viewed as the numerical integration of a differential equation.
In this extended abstract, we build upon this connection and propose a modiﬁcation of the internal architecture of a transformer layer.
The proposed model places the multi-head attention sublayer and the MLP sublayer parallel to each other.
Our experiments show that this simple modiﬁcation improves the performance of transformer networks in multiple tasks.
Moreover, for the image classiﬁcation task, we show that using neural ODE solvers with a sophisticated integration scheme further improves performance.
Introduction
Over the last few years, the transformer layer introduced by Vaswani et al.
(2017) has become a key component in deep learning models used in natural language processing (Devlin et al., 2019; Brown et al., 2020; Shoeybi et al., 2019), image and video processing (Dosovitskiy et al., 2021;
Arnab et al., 2021; Han et al., 2022), and audio and speech processing (Dong et al., 2018).
Current state-of-the-art techniques in language processing (e g., machine translation, natural language understanding, and information/knowledge extraction) rely heavily on the use of transformer layers to encode information in its word vector about the relevant context of a given word.
This allows the model to focus on relevant contexts at different length scales.
Although the transformer layers were originally introduced as a sequence-to-sequence transduction model, they have also demonstrated superior performance in various computer vision tasks beyond image classiﬁcation, e.g., semantic segmentation (Zheng et al., 2021; Strudel et al., 2021;
Ding et al., 2022), object detection (Carion et al., 2020; Song et al., 2022), and view synthesis (Kulhánek et al., 2022; Lin et al., 2022).
Inside a multilayer transformer network, each transformer layer consists of a multi-head attention sublayer followed by an MLP sublayer, creating an alternating pattern of these sublayers throughout the network.
Moreover, both of these sublayers use residual connections to avoid the vanishing gradient problem and facilitate the training of very deep transformer networks.
Prior work (Haber and
Ruthotto, 2017; Haber et al., 2018; Lu et al., 2018) has shown that the forward propagation through residual connections can be viewed as Euler discretization of a time-varying ordinary differential equation (ODE).
This insight suggests a connection between transformer networks and differential equations that can potentially be exploited to further improve the performance of transformer networks (e.g., higher accuracy, fewer parameters).
Indeed, the residual connections alongside the alternating pattern of multi-head attention and MLP sublayers can be interpreted as numerical integration via the Lie-Trotter splitting scheme (Lu et al., 2020; Dutta et al., 2021).
In this work, we show that the connection between transformers and ODEs can be leveraged to design a new architecture wherein the multi-head attention and MLP sublayers are placed parallelly, not sequentially, inside the individual transformer layers.
Our experiments show that the proposed model performs better than the original transformer layer when tested on image classiﬁcation, machine translation, and language modeling tasks.
A growing body of work has focused on improving transformer networks by reorganizing the sublayers and leveraging the connection between transformers and ODEs.
Macaron Net (Lu et al., 2020) draws inspiration from numerical integration techniques to use a multi-head attention sublayer sandwiched between two MLP sublayers.
Dutta et al.
(2021) use a temporal evolution scheme to avoid the computationally expensive step of calculating dot-product attention at each transformer layer; instead, it computes the dot-product attention at the initial step and then time-evolves it through the layers.
The proposed model is similar in spirit to this line of work - it modiﬁes the internal architecture of the transformer layer and places the multi-head attention sublayer and the MLP sublayer side-by-side.
On the other hand, Press et al.
(2020) have explored the effect of reordering the individual sublayers and changing their numbers while keeping the total number of model parameters ﬁxed.
Their work shows that a transformer network can improve its performance by concentrating the multi-head attention and MLP layers in the lower and upper stages of the network, respectively; however, this performance improvement is not uniform across all tasks.
The main contributions of this work are as follows:
• By leveraging the connection between transformer layers and ODEs, we propose a novel variant of the transformer layer wherein the multi-head attention and MLP sublayers are placed side-by-side.
• Through numerical experiments, we demonstrate that the proposed model performs better than the original transformer layer across multiple tasks (image classiﬁcation on CIFAR-100, machine translation on WMT-2014 English-German dataset, and language modeling on WikiText-103).
Our experiments have been carried out with small models due to resource constraints.
• We also demonstrate that the performance of the proposed model can be further improved by using neural ODE solvers with sophisticated integration schemes (e.g., RK4).
Proposed Architecture
By letting X m :=
[ x m 1 , x m 2 , · · · , x mL ] denote the input to the m -th transformer layer, the operation carried out by the multi-head attention sublayer can be expressed as
[Equation]
where L is the length of the input sequence and the function G represents the multi-head dot-product attention.
ˆ x mi , i.e., the output from this sublayer, is then fed to the MLP sublayer and undergoes the following transformation to yield X m +
=
[ x m +
[Equation]
m +
[Equation]
where the function F represents the sequence of linear mappings and activation functions.
As Lu et al.
(2020) and Dutta et al.
(2021) have highlighted, (1)-(2) can viewed as the numerical integration (over the time interval [ m, m + 1] ) of the following ODE via the Lie-Trotter splitting scheme
[Equation]
[Figure]
where X :=
[ x 1 , x 2 , · · · , x L ] .
This interpretation paves the way for multiple approaches to realize transformer layers.
As neural
ODE solvers and their variants (Chen et al., 2018; Massaroli et al., 2020; Kidger et al., 2021) provide a means to run backpropagation through any black-box ODE solver, a transformer layer can be implemented using neural ODE networks.
Alternatively, the time integration of (3) over the interval [ m, m + 1] can also be approximated as
[Equation]
Our proposed model implements (4) by placing the multi-head attention and MLP sublayers side-byside (Figure 1).
Moreover, if they share their weights, a D -layer deep stack of transformer layers can be viewed as the numerical integration of (3) over the time interval [0 , D ] .
This perspective offers a means to introduce weight sharing into a transformer network in a gradual way for understanding the trade-off between model performance and the number of model parameters.
For example, a 12-layer transformer network can be replaced by a sequence of six 2-layer, weight sharing transformers (which can be implemented by either stacking two layers of the proposed model or via integration over a longer horizon using a neural ODE).
Experiments
In this section, we investigate the performance of the proposed model for both computer vision and language processing tasks.
However, due to limited computing resources, we use smaller models for comparing the performance between the baseline transformer architecture and the proposed model.
Therefore, the reported results of baselines may not be comparable to the state-of-the-art results.
3.1 Image Classiﬁcation
First, we investigate the proposed model’s performance in image classiﬁcation tasks.
In particular, we take the DeiT-Ti model proposed by Touvron et al.
and modify the transformer layer architecture to implement our proposed model.
Then, to compare the performance of DeiT-Ti and variants of our proposed models on CIFAR-100, we use top-1 accuracy as the metric.
Table 1 demonstrates the performance of the side-by-side sublayers with different levels of weight sharing.
We observe that by only changing the attention block and MLP block from sequential to parallel, the top-1 accuracy increases from 66.02% to 70.92%.
This architecture does not use any weight sharing.
Increasing the amount of weight sharing reduces the number of independent layers and results in fewer trainable parameters.
However, as the number of parameters decreases, the model performance also deteriorates.
By leveraging the connection between differential equations and the transformer layers, we also replace the Euler scheme with a more sophisticated Runge-Kutta (RK4) integration scheme.
This further increases the top-1 accuracy from 70.92% (Euler) to 72.66% (RK4).
Furthermore, we study the effects of dropout (Srivastava et al., 2014) and stochastic depth (Huang et al., 2016) in Table 2.
We conclude that neither dropout nor stochastic depth is helpful in the proposed architecture.
[Table]
[Figure] [Figure]
[Table]
We also conducted a small ablation study on the placement of the normalization layer in our proposed model.
Figure 2 shows that all three variations perform better than the original DeiT-Ti model.
Moreover, our proposed normalization approach (Figure 2a) yields the best performance.
3.2 Natural Language Processing
We leverage the open source toolkit Fairseq (Ott et al., 2019) to investigate the performance of the proposed model in neural machine translation and language modeling.
Additional details about the size of the architectures used in this section are provided in the Appendix.
3.2.1 Neural Machine Translation
[Table]
To train small transformer models on the WMT-2014 English-German translation dataset, we follow the training procedure speciﬁed by Ott et al.
After training the models forepochs, we compare their performance by computing the validation loss and detokenized BLEU score with SacreBLEU (Post, 2018) as described by Ott et al.
As shown in Table 3, the proposed model outperforms the baseline.
3.2.2 Neural Language Modeling
[Table]
We follow the training and evaluation procedure speciﬁed by Baevski and Auli (2019) and Ott et al.
(2019) to train a small language model on the WikiText-103 dataset.
Table 4 reports the metric of perplexity for the baseline and proposed model for different dropout rates.
We can notice that for these small models, dropout doesn’t lead to improved performance for both the baseline and the proposed architecture.
With zero dropout, our proposed architecture performs better than the baseline.
Conclusion
This work has proposed a new variant of the transformer layer by leveraging its connection with ODEs and has shown that the proposed model outperforms the original transformer layer.
Furthermore, as shown by our initial results from an RK4-based neural ODE solver, one can extend this work to investigate the potential of using a time-dependent neural ODE to implement transformer networks.
Appendices A Model Size Used in the Neural Machine Translation Task
[Table]
B Model Size Used in the Neural Language Modeling Task
[Table]