Neural ODE Transformers: Analyzing Internal Dynamics and Adaptive Fine-tuning

TL;DR

The paper introduces DiffEqFormer, a non-autonomous neural ODE transformer that parameterizes all attention and feed-forward weights as time-dependent functions, enabling continuous-depth computation and flexible fine-tuning across architectures. By applying spectral analysis to QK and OV pairs and employing Lyapunov exponents for token-level sensitivity, it reveals dynamic properties that counter clustering observed in weight-sharing models and provides interpretable metrics. Empirically, DiffEqFormer achieves competitive perplexities on WikiText103 and OpenWebText, often surpassing corresponding GPT and Llama baselines, and supports adaptive fine-tuning via solver-step adjustments and LoRA. The work offers a principled continuous-depth perspective with practical benefits for model adaptation, memory efficiency considerations, and future integration of stochastic or advanced ODE solvers.

Abstract

Recent advancements in large language models (LLMs) based on transformer architectures have sparked significant interest in understanding their inner workings. In this paper, we introduce a novel approach to modeling transformer architectures using highly flexible non-autonomous neural ordinary differential equations (ODEs). Our proposed model parameterizes all weights of attention and feed-forward blocks through neural networks, expressing these weights as functions of a continuous layer index. Through spectral analysis of the model's dynamics, we uncover an increase in eigenvalue magnitude that challenges the weight-sharing assumption prevalent in existing theoretical studies. We also leverage the Lyapunov exponent to examine token-level sensitivity, enhancing model interpretability. Our neural ODE transformer demonstrates performance comparable to or better than vanilla transformers across various configurations and datasets, while offering flexible fine-tuning capabilities that can adapt to different architectural constraints.

Figures (29)

• Figure 1: (a) The vector field of DiffEqFormer with attention block and feed-forward block constructed from time-dependent weights. (b) The architecture of time-dependent weights.
• Figure 2: (a-b): Spectral dynamics of QK and OV pairs. (c-d): Trajectory of a sequence consisting of 40 points in 3-dimensional space. (c) Attention-only model with shared weight assumptions as described in geshkovski2023emergence: Clusters emerge. (d) Attention-only model with time-dependent weights of increasing magnitude, inspired by observations from trained DiffEqFormer: No clusters occur. (e) Plot of a function in our simulation that mimics the magnitude of $Q(t), K(t), V(t)$ over time like in trained DiffEqFormer.
• Figure 3: Lyapunov exponent values represent the sensitivity of previous words to the next word. Higher values correspond to more intense highlighting in red. (a) The next word is its. (b) The next word is match.
• Figure 4: Validation Perplexity of DiffEqFormer in comparison with GPT models. (a,b) Results on WikiText103 dataset in two architecture settings. (c, d, e) Results on OpenWebText dataset on three architecture settings.
• Figure 5: Finetune validation perplexity across two settings: (a) OpenWebText$\to$WikiText103; (b) WikiText103$\to$OpenWebText . All models are pretrained with 18 function evaluations (or layers) and finetune in different settings of function evaluations 9, 12, 18, 24 with LoRA and full-rank finetune. We compare these with the baseline of corresponding GPT model pretrained and finetune with 24 layers which is on much expensive computation both pretrain and finetune.