ODE-ViT: Plug & Play Attention Layer from the Generalization of the ViT as an Ordinary Differential Equation

TL;DR

ODE-ViT reframes Vision Transformers as continuous-depth models by treating the attention mechanism as an autonomous ODE, enabling stable, well-posed dynamics and improved interpretability with fewer parameters. The authors enforce local Lipschitz continuity and use center normalization, while a Lie–Trotter based decomposition splits the attention block into learnable subflows; they also introduce a plug-and-play teacher–student framework where a pretrained discrete ViT guides the ODE trajectory via MSE and JaSMin losses. Empirically, ODE-ViT achieves competitive CIFAR-10/100 performance with up to an order of magnitude fewer parameters and demonstrates robustness to step count; the teacher–student setup yields substantial gains over the base ODE-ViT. The work provides a theoretical and practical foundation for continuous-depth ViTs, opening avenues for higher‑order solvers, adaptive steps, and multi-modal extensions.

Abstract

In recent years, increasingly large models have achieved outstanding performance across CV tasks. However, these models demand substantial computational resources and storage, and their growing complexity limits our understanding of how they make decisions. Most of these architectures rely on the attention mechanism within Transformer-based designs. Building upon the connection between residual neural networks and ordinary differential equations (ODEs), we introduce ODE-ViT, a Vision Transformer reformulated as an ODE system that satisfies the conditions for well-posed and stable dynamics. Experiments on CIFAR-10 and CIFAR-100 demonstrate that ODE-ViT achieves stable, interpretable, and competitive performance with up to one order of magnitude fewer parameters, surpassing prior ODE-based Transformer approaches in classification tasks. We further propose a plug-and-play teacher-student framework in which a discrete ViT guides the continuous trajectory of ODE-ViT by treating the intermediate representations of the teacher as solutions of the ODE. This strategy improves performance by more than 10% compared to training a free ODE-ViT from scratch.

Figures (10)

• Figure 1: Exemplification of the hypothesis that the MSE does not need to converge to the same point as the teacher, but only within the contraction region. In this region, the classification head predicts the correct class identically to the teacher. Both CLS(T) refer to the last representation of both attention blocks.
• Figure 2: Overview of the plug-and-p lay Teacher--Student framework. The input image is processed through both Token Representation modules and passed through their respective attention layers. The orange trajectory represents the ODE--ViT architecture with teacher supervision, while the blue pathway corresponds to the Teacher ViT. The forced trajectory is optimized by minimizing \ref{['eq:min error']} from the [CLS] tokens. The classification head ($\pi$) is initialized from the teacher and kept frozen during training; in our experiments, only one case benefited from unfreezing it. It is also exemplified that the proportion of states required to reach each ViT hidden state varies across layers.
• Figure 3: Analysis plot with the distances between the [CLS] token of the ViT and the [CLS] token of the ODE--ViT for all samples in the test set. The green color means that the ODE--ViT and the ViT correctly classify the sample, while the red means that both models badly classify it. Finally, the yellow means that only the teacher correctly classify the sample.
• Figure 4: Distribution of Lyapunov exponents for CIFAR-10 (left) and CIFAR-100 (right) classes. The color and horizontal position of each bubble represent the classification accuracy of the ODE-ViT for the corresponding class, while the vertical axis shows the mean Lyapunov exponent of its dynamics. In both datasets, a mild correlation emerges between the Lyapunov exponent and the classification accuracy, indicating that classes with more stable dynamics (lower Lyapunov exponents) tend to exhibit higher accuracy.
• Figure 5: Attention maps of the [CLS] token from the ODE--ViT at its final state. Images in the left are taken from the ImageNet dataset, while the images in the right are taken from CIFAR-100. The model used to extract these representation was trained on CIFAR-100 using the ODE--ViT Base after trained with the teacher-student framework. The results demonstrate that the model has learned a generalizable representation.

Theorems & Definitions (1)

• Proposition 1: Approximation Errorconf/nips/SanderAP22