IIET: Efficient Numerical Transformer via Implicit Iterative Euler Method

TL;DR

The paper introduces the Iterative Implicit Euler Transformer (IIET), a Transformer architecture that embeds an iterative implicit Euler solver within each decoder layer to improve accuracy while enabling model compression. To address inference efficiency, it also proposes Iteration Influence-Aware Distillation (IIAD), which prunes nonessential iterations based on per-layer iteration influence and distills the resulting efficient model back to preserve performance. Empirically, IIET delivers consistent accuracy gains over vanilla Transformers and PCformer across model sizes, with an efficient variant (E-IIET) achieving substantial inference speedups while retaining most of the performance. The work positions IIET as a practical, scalable alternative to high-order ODE-based Transformers, offering flexible deployment options through its distillation-based efficiency controls. Overall, the approach advances the integration of implicit numerical methods into autoregressive language modeling, delivering both performance and deployment benefits.

Abstract

High-order numerical methods enhance Transformer performance in tasks like NLP and CV, but introduce a performance-efficiency trade-off due to increased computational overhead. Our analysis reveals that conventional efficiency techniques, such as distillation, can be detrimental to the performance of these models, exemplified by PCformer. To explore more optimizable ODE-based Transformer architectures, we propose the Iterative Implicit Euler Transformer (IIET), which simplifies high-order methods using an iterative implicit Euler approach. This simplification not only leads to superior performance but also facilitates model compression compared to PCformer. To enhance inference efficiency, we introduce Iteration Influence-Aware Distillation (IIAD). Through a flexible threshold, IIAD allows users to effectively balance the performance-efficiency trade-off. On lm-evaluation-harness, IIET boosts average accuracy by 2.65% over vanilla Transformers and 0.8% over PCformer. Its efficient variant, E-IIET, significantly cuts inference overhead by 55% while retaining 99.4% of the original task accuracy. Moreover, the most efficient IIET variant achieves an average performance gain exceeding 1.6% over vanilla Transformer with comparable speed.

Figures (6)

• Figure 1: Architectural comparison: (a) Vanilla Transformer; (b) Linear multistep-enhanced Transformer; (c) PCformer with 2nd-order Runge-Kutta predictor and 1st-order Euler corrector; (d) Our proposed Iterative Implicit Euler Transformer (IIET). The iteration steps $r$ in IIET is configurable, with experimental validation determining $r=3$ as the optimal setting in this work. All blocks follow an identical computational procedure as the $block_n$.
• Figure 2: PPL on the Wikitext test set for 55M and 340M IIET across varying iteration steps $r$. Dashed lines indicate Transformer++ and PCformer performance at corresponding parameter scales. Note that IIET's FLOPs is nearly $r+1$ times of Transformer++.
• Figure 3: Ablation study on iteration steps $r$: (a) Impact on model performance. (b) Corresponding effects on inference speed and VRAM utilization.
• Figure 4: Distribution of Block Influence (BI) for Transformer++ and IIET models with varying iteration steps $r$. Higher BI values indicate lower model redundancy.
• Figure 5: Iteration Influence within each layer of the 340M IIET model. Deeper colors indicate larger hidden state changes after this iteration. The 740M IIET results are presented in Appendix \ref{['sec:740m_iiet_heatmap']} due to space constraints.