[Paper Review] LipsFormer: Introducing Lipschitz Continuity to Vision Transformers

Xianbiao Qi, Jianan Wang+3 more|arXiv (Cornell University)|Apr 19, 2023

Neural Networks and Reservoir Computing4 citations

TL;DR

LipsFormer introduces a theoretically grounded, Lipschitz continuous Vision Transformer architecture by replacing unstable components—LayerNorm with CenterNorm, dot-product attention with scaled cosine similarity attention, residual shortcuts with weighted variants, and Xavier initialization with spectral initialization. This design enables stable training without learning rate warmup, achieving 82.7% top-1 accuracy on ImageNet-1K with Swin-Tiny and 83.5% with CSwin-Tiny in 300 epochs.

ABSTRACT

We present a Lipschitz continuous Transformer, called LipsFormer, to pursue training stability both theoretically and empirically for Transformer-based models. In contrast to previous practical tricks that address training instability by learning rate warmup, layer normalization, attention formulation, and weight initialization, we show that Lipschitz continuity is a more essential property to ensure training stability. In LipsFormer, we replace unstable Transformer component modules with Lipschitz continuous counterparts: CenterNorm instead of LayerNorm, spectral initialization instead of Xavier initialization, scaled cosine similarity attention instead of dot-product attention, and weighted residual shortcut. We prove that these introduced modules are Lipschitz continuous and derive an upper bound on the Lipschitz constant of LipsFormer. Our experiments show that LipsFormer allows stable training of deep Transformer architectures without the need of careful learning rate tuning such as warmup, yielding a faster convergence and better generalization. As a result, on the ImageNet 1K dataset, LipsFormer-Swin-Tiny based on Swin Transformer training for 300 epochs can obtain 82.7\% without any learning rate warmup. Moreover, LipsFormer-CSwin-Tiny, based on CSwin, training for 300 epochs achieves a top-1 accuracy of 83.5\% with 4.7G FLOPs and 24M parameters. The code will be released at \url{https://github.com/IDEA-Research/LipsFormer}.

Motivation & Objective

• Address the persistent challenge of training instability in Vision Transformers, especially at initialization.
• Identify that Lipschitz continuity is a more fundamental property than common training tricks like learning rate warmup or normalization.
• Develop a holistic framework to enforce Lipschitz continuity across key components: normalization, attention, residual connections, and weight initialization.
• Provide theoretical upper bounds on the Lipschitz constant to guide stable network design.
• Demonstrate that LipsFormer enables stable, end-to-end training without hyperparameter tuning for learning rate warmup.

Proposed method

• Replace LayerNorm with CenterNorm, a Lipschitz continuous normalization layer that centers features without scaling.
• Substitute dot-product self-attention with scaled cosine similarity attention, proven to be Lipschitz continuous under proper scaling.
• Introduce a weighted residual shortcut with a learnable, zero-initialized gate to stabilize residual updates.
• Apply spectral-based weight initialization instead of Xavier, ensuring bounded spectral norms at initialization.
• Derive a theoretical upper bound on the Lipschitz constant of the full LipsFormer architecture using composition rules of Lipschitz functions.
• Construct LipsFormer-Swin and LipsFormer-CSwin by replacing standard components in Swin and CSwin Transformers with their Lipschitz-continuous counterparts.

Figure 1: Comparison of a LipsFormer block with ConvNeXt, Transformer and Swin-Transformer Blocks. We use different colors to mark our Lipschitz improvements.

Experimental results

Research questions

• RQ1 Can enforcing Lipschitz continuity across key Transformer components lead to more stable training without learning rate warmup?
• RQ2 How does the theoretical Lipschitz constant of a Transformer network relate to its practical trainability and generalization?
• RQ3 What are the specific architectural components in standard Transformers that violate Lipschitz continuity and contribute to training instability?
• RQ4 Can a holistic Lipschitz design outperform incremental improvements like DeepNorm or ReZero in training stability and convergence?
• RQ5 Does a Lipschitz-continuous Transformer maintain or improve performance on standard benchmarks like ImageNet-1K without hyperparameter tuning?

Key findings

• LipsFormer achieves 82.7% top-1 accuracy on ImageNet-1K using Swin-Tiny with 300 epochs, without any learning rate warmup.
• LipsFormer-CSwin-Tiny reaches 83.5% top-1 accuracy with only 4.7G FLOPs and 24M parameters, also without learning rate warmup.
• The theoretical upper bound on the Lipschitz constant of LipsFormer is derived and shown to be finite at initialization, providing a principled stability guarantee.
• Experiments confirm that the network remains stable and converges faster when trained without warmup, validating the assumption of Lipschitz continuity.
• The proposed components—CenterNorm, scaled cosine attention, weighted residual shortcut, and spectral initialization—can be used as drop-in replacements in existing Vision Transformer architectures.
• LipsFormer outperforms standard Swin and CSwin baselines in training stability and generalization, especially in deep architectures, without requiring careful hyperparameter tuning.

Figure 2: Sigmoid, Tanh, ReLU and GELU activation function.