Understanding and Improving Transformer From a Multi-Particle Dynamic System Point of View

Yiping Lu; Zhuohan Li; Di He; Zhiqing Sun; Bin Dong; Tao Qin; Liwei Wang; Tie-Yan Liu

Source-linked AI summary

Understanding and Improving Transformer From a Multi-Particle Dynamic System Point of View

Yiping Lu, Zhuohan Li, Di He, Zhiqing Sun, Bin Dong, Tao Qin, Liwei Wang, Tie-Yan Liu

arXiv:1906.02762v1Jun 6, 2019cs.LGcs.CLstat.ML

TL;DR

The paper addresses the elusive design principle of the successful Transformer architecture. It interprets Transformer layers as a numerical ODE solver for a multi-particle convection-diffusion system, then replaces Lie-Trotter with Strang-Marchuk splitting to create Macaron Net. Macaron Net achieves higher accuracy than Transformer across supervised and unsupervised tasks, while the reported limitation concerns the treatment of layer normalization in the reformulation.

Problem
The Transformer performs strongly, but its design principle and alternative design choices remain insufficiently understood.
Method
The paper models Transformer layers as an ODE solver for multi-particle convection-diffusion dynamics and derives Macaron Net by replacing Lie-Trotter with Strang-Marchuk splitting.
Results
Macaron Net achieves higher accuracy than Transformer on all evaluated supervised and unsupervised learning tasks with the same number of parameters.
Takeaways & Limitations
Numerical ODE analysis provides a perspective for understanding Transformer representations and designing architectures beyond the standard Transformer.
Takeaways & Limitations
The reformulation remains valid with layer normalization, but the paper notes that normalization may be removed because it is not essential according to cited recent work.

Abstract

from arXiv · show

The Transformer architecture is widely used in natural language processing. Despite its success, the design principle of the Transformer remains elusive. In this paper, we provide a novel perspective towards understanding the architecture: we show that the Transformer can be mathematically interpreted as a numerical Ordinary Differential Equation (ODE) solver for a convection-diffusion equation in a multi-particle dynamic system. In particular, how words in a sentence are abstracted into contexts by passing through the layers of the Transformer can be interpreted as approximating multiple particles' movement in the space using the Lie-Trotter splitting scheme and the Euler's method. Given this ODE's perspective, the rich literature of numerical analysis can be brought to guide us in designing effective structures beyond the Transformer. As an example, we propose to replace the Lie-Trotter splitting scheme by the Strang-Marchuk splitting scheme, a scheme that is more commonly used and with much lower local truncation errors. The Strang-Marchuk splitting scheme suggests that the self-attention and position-wise feed-forward network (FFN) sub-layers should not be treated equally. Instead, in each layer, two position-wise FFN sub-layers should be used, and the self-attention sub-layer is placed in between. This leads to a brand new architecture. Such an FFN-attention-FFN layer is "Macaron-like", and thus we call the network with this new architecture the Macaron Net. Through extensive experiments, we show that the Macaron Net is superior to the Transformer on both supervised and unsupervised learning tasks. The reproducible codes and pretrained models can be found at https://github.com/zhuohan123/macaron-net

1 Introduction

The paper addresses the limited understanding of the Transformer's design by interpreting it through multi-particle dynamics and numerical ODE solvers. This perspective motivates Macaron Net, which uses Strang-Marchuk splitting and achieves higher accuracy than Transformer models across supervised and unsupervised tasks.

The Transformer uses stacked layers combining self-attention and position-wise FFN sub-layers with residual connections.Self-attention captures dependencies across positions, while the FFN encodes each position into higher-level representations.
The paper interprets Transformer layers as a numerical ODE solver for a first-order convection-diffusion equation in a multi-particle dynamic system.Self-attention corresponds to diffusion, position-wise FFNs correspond to convection, and stacked layers approximate particle movement.
Strang-Marchuk splitting replaces Lie-Trotter splitting with two half-step FFN sub-layers surrounding a full-step self-attention sub-layer.The resulting FFN-attention-FFN structure is called a Macaron layer, and the resulting network is called Macaron Net.
Macaron Net achieves higher accuracy than Transformer on all evaluated supervised and unsupervised learning tasks with the same number of parameters.The authors report this result as consistent with the ODE theory.

2 Background

The background connects residual neural networks with numerical ODE solvers and describes the Transformer’s attention and FFN components. It also motivates a multi-particle formulation because variable-length sequences and heterogeneous sub-layers prevent directly applying a single-particle ODE interpretation.

2.1 Relationship Between Neural Networks and ODE: Euler’s method discretizes an ODE sequentially, and its update is mathematically equivalent to a residual neural-network block.The neural block represents γf(x_l,t_l), while t_l indicates the parameters of the l-th layer.
2.1 Relationship Between Neural Networks and ODE: ODE-inspired numerical solvers have been used to design neural networks, including architectures based on high-precision Runge-Kutta methods.The cited works connect parametric ODE solvers with deep residual networks and optimize ODE parameters through backpropagation.
2.2 Transformer: A Transformer layer applies multi-head self-attention followed by a position-wise FFN, with residual connections and layer normalization for both sub-layers.The layer operates on a sequence of vectors and outputs a sequence with the same shape.
2.2 Transformer: Self-attention queries key-value representations to combine information across positions, while multi-head attention attends jointly across representation subspaces.The attention formulation uses Q, K, and V from the previous layer; the multi-head output concatenates multiple attention heads.
2.2 Transformer: The paper argues that a single-particle ODE formulation does not directly fit NLP because sentences have variable length and Transformer sub-layers differ in their information flow.Self-attention uses information from all positions, unlike the position-wise FFN.

3 Reformulate Transformer Layers as an ODE Solver for Multi-Particle Dynamic System

The paper reformulates Transformer layers as a numerical ODE solver for a multi-particle convection-diffusion system, then uses a more accurate splitting scheme to motivate Macaron layers.

3.1 Multi-Particle ODE and Its Numerical Solver: Multi-particle dynamics model each particle using convection from its own movement and diffusion arising from interactions with other particles.
3.1 Multi-Particle ODE and Its Numerical Solver: Splitting methods solve coupled convection-diffusion dynamics by separating the differential operators, with Lie-Trotter applying them successively.The method first advances one term to an intermediate state, then advances the other term.
3.2 Physical interpretation of the Transformer: Transformer layers match this solver structure: self-attention represents diffusion, while position-wise FFNs represent convection across sequence-position vectors.The correspondence is formalized by mapping MultiHeadAtt to F* and FFN to G*.
3.2 Physical interpretation of the Transformer: The Transformer can therefore be viewed as using Lie-Trotter splitting and Euler’s method with time step γ to solve the multi-particle ODE.
3.3 Improving Transformer Via Strang-Marchuk Splitting Scheme: Strang-Marchuk splitting reduces ordering bias by dividing the G step into two half-steps around F, yielding three neural sub-layers.Its local truncation error is O(γ3), compared with O(γ2) for Lie-Trotter; the neural implementation uses FFN-attention-FFN with half-step residual connections.

4 Experiments

The experiments evaluate Macaron architectures on supervised machine translation and unsupervised GLUE tasks under parameter-matched comparisons. Macaron Net outperforms Transformer baselines across the reported tasks.

Experiment Settings: Macaron layers replace Transformer layers while keeping parameter counts the same across corresponding settings.Each Macaron layer contains two position-wise FFN sub-layers, with adjusted inner dimensionality for fair comparison.
Experiment Results: 1.0 BLEU on IWSLT14 and 1.6 BLEU on WMT14 show Macaron small and base outperforming their Transformer counterparts.Macaron big also exceeds Transformer big by 1.8 BLEU points.
Experiment Results: Macaron Net base achieves the best general GLUE score, outperforming a trained BERT base model on all tasks and the original BERT base on 6 of 8 tasks.The remaining two comparisons achieve close performance.
Experiment Results: Across machine translation and GLUE, Macaron Net performs better than Transformer in practice, consistent with the ODE theory.The experiments keep the number of parameters fixed for the reported architecture comparisons.

5 Conclusion and Future Work

The paper interprets Transformer layers as numerical ODE solving steps and replaces Lie-Trotter splitting with Strang-Marchuk splitting to obtain Macaron Net. Reported improvements in translation and GLUE are consistent with this theory.

Conclusion: Transformer feature abstraction is interpreted as approximating particle movement through Lie-Trotter splitting and Euler’s method.The underlying equation is a convection-diffusion equation in a multi-particle dynamic system.
Conclusion: Replacing Lie-Trotter with Strang-Marchuk splitting produces the Macaron architecture.The paper presents this replacement as the basis for the new network design.
Conclusion and Future Work: Improvements in real applications support the effectiveness of Macaron Net and align with the ODE interpretation.The paper identifies deeper ODE connections and improved Transformer components as future directions.

A Proof of the Theorem

The appendix develops local truncation-error estimates for Lie-Trotter and Strang-Marchuk splitting schemes. It also relates Euler discretization to residual-network updates.

Lie-Trotter Analysis: The Lie-Trotter local truncation error is estimated using derivatives of the component functions and higher-order remainder terms.The displayed analysis includes terms involving F′ and G′ and an O(γ3) remainder.
Strang-Marchuk Analysis: The appendix rewrites the Strang-Marchuk scheme to estimate its local truncation error.The proof combines intermediate estimates for the splitting operators.
Theorem Scope: The theorem is stated for time-independent ODEs but extends to time-dependent ODEs by introducing an auxiliary variable.This extension is presented without loss of generality.

B.1 Machine Translation

The machine-translation experiments use IWSLT14 and WMT14 with standard Transformer settings, parameter-matched Macaron replacements, and BLEU evaluation under stated training protocols.

Dataset: IWSLT14 contains about 153K/7K/7K training, validation, and test sentence pairs, while WMT14 contains 4.5M training pairs.The datasets use joint source-target byte-pair-encoded vocabularies of 10K and 37K tokens, respectively.
Model: WMT14 uses six-layer encoder and decoder stacks with base and big configurations, while IWSLT14 uses a smaller setting.The WMT14 hidden dimensions are 512 for base and 1024 for big, with 8 and 16 attention heads.
Model: Macaron models replace every Transformer layer with a layer containing two position-wise FFN sub-layers.The decoder additionally retains encoder-decoder attention and masking.
Model: Parameter counts remain matched by setting each Macaron FFN inner-layer dimensionality to twice the hidden-state dimensionality.This adjustment is applied across base, big, and small settings.
Optimizer and Training: Training uses Adam, the reference learning-rate schedule, label smoothing of ϵls = 0.1, and fairseq-based PyTorch implementations.WMT14 models use four NVIDIA P40 GPUs and IWSLT14 models use one.
Evaluation: BLEU evaluation uses tokenized case-sensitive scoring for WMT14 En-De and case-insensitive scoring for IWSLT14 De-En.Inference uses beam search with dataset-specific beam sizes and length penalties.

B.2 Unsupervised Pretraining

The unsupervised pretraining setup uses Wikipedia and BookCorpus, matched in scale to the original data, and replaces BERT-base Transformer layers with parameter-matched Macaron layers. Pretraining follows BERT’s masked-language-modeling and next-sentence-prediction objectives.

Dataset: Pretraining uses English Wikipedia and a self-collected BookCorpus, totaling roughly 3.4B words, comparable with the corpus used in the original BERT setup.Documents are segmented, normalized, lower-cased, and tokenized before training.
Model: The Macaron base replaces BERT base’s 12 Transformer layers while using 768-dimensional hidden states, 12 attention heads, and the same parameter count.The inner dimensionality of both FFN sub-layers is halved to preserve the parameter budget.
Objectives: Unsupervised pretraining uses masked language modeling and next sentence prediction, following the two-task setup of the original BERT work.The model predicts randomly masked tokens and whether sentence pairs are consecutive.

B.3 GLUE Dataset

The GLUE evaluation covers eight downstream language-understanding tasks, with task-specific metrics and development-set hyperparameter selection. WNLI is excluded, and test labels remain unavailable with one evaluation-server submission per model.

Tasks: The evaluation includes CoLA, SST-2, MRPC, STS-B, QQP, MNLI, QNLI, and RTE, spanning acceptability, sentiment, paraphrase, similarity, inference, and entailment tasks.The tasks use accuracy, correlation, F1, or Matthews correlation coefficient depending on the benchmark.
Tasks: MNLI evaluates three-way textual entailment and reports accuracy separately on matched in-domain and mismatched cross-domain test sections.The labels are entailment, contradiction, and neutral.
Evaluation scope: WNLI is skipped because few previous works outperform majority-class prediction on this task.WNLI requires resolving a pronoun’s referent in a Winograd-style sentence.
Fine-tuning: Fine-tuning searches batch sizes 16/32, learning rates 5e-3/3e-5, epochs 3/4/5, and random seeds, selecting test models by development-set performance.The procedure follows the cited BERT fine-tuning setup.
Evaluation scope: GLUE test labels are not distributed, and each model receives only a single evaluation-server submission.This constrains direct independent inspection of the test labels and repeated submissions.