[论文解读] LipsFormer: Introducing Lipschitz Continuity to Vision Transformers

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

Neural Networks and Reservoir Computing被引用 4

一句话总结

LipsFormer 通过用理论上有保障的 Lipschitz 连续性组件替代不稳定的组件，提出了一种理论基础扎实的 Lipschitz 连续 Vision Transformer 架构：将 LayerNorm 替换为 CenterNorm，将点积注意力替换为缩放余弦相似度注意力，将残差捷径替换为加权变体，并将 Xavier 初始化替换为谱初始化。该设计实现了无需学习率预热的稳定训练，在 300 个周期内于 ImageNet-1K 上达到 Swin-Tiny 的 82.7% 顶级-1 准确率，以及 CSwin-Tiny 的 83.5% 准确率。

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}.

研究动机与目标

• 解决 Vision Transformer 在初始化阶段持续存在的训练不稳定性挑战。
• 识别出 Lipschitz 连续性是比学习率预热或归一化等常见训练技巧更根本的性质。
• 构建一个整体框架，以在关键组件（归一化、注意力、残差连接和权重初始化）中强制实现 Lipschitz 连续性。
• 推导 Lipschitz 常数的理论上限，以指导稳定网络的设计。
• 证明 LipsFormer 可实现无需学习率预热超参数调优的稳定端到端训练。

提出的方法

• 用 CenterNorm 替代 LayerNorm，CenterNorm 是一种不进行缩放的 Lipschitz 连续归一化层，仅对特征进行中心化。
• 将点积自注意力替换为缩放余弦相似度注意力，该方法在适当缩放下已被证明具有 Lipschitz 连续性。
• 引入一种带有可学习、零初始化门的加权残差捷径，以稳定残差更新。
• 采用基于谱的权重初始化，替代 Xavier 初始化，以确保初始化时谱范数有界。
• 利用 Lipschitz 函数的复合规则，推导出完整 LipsFormer 架构的 Lipschitz 常数的理论上限。
• 通过将 Swin 和 CSwin Transformer 中的标准组件替换为其 Lipschitz 连续对应物，构建 LipsFormer-Swin 和 LipsFormer-CSwin。

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

实验结果

研究问题

• RQ1 在关键 Transformer 组件中强制实现 Lipschitz 连续性是否能实现无需学习率预热的更稳定训练？
• RQ2 Transformer 网络的理论 Lipschitz 常数与其实际可训练性和泛化能力之间有何关系？
• RQ3 标准 Vision Transformer 中哪些具体架构组件违反了 Lipschitz 连续性并导致训练不稳定性？
• RQ4 整体 Lipschitz 设计是否能超越 DeepNorm 或 ReZero 等渐进式改进，在训练稳定性和收敛性方面表现更优？
• RQ5 Lipschitz 连续的 Transformer 是否能在无需超参数调优的情况下，保持或提升在 ImageNet-1K 等标准基准上的性能？

主要发现

• LipsFormer 在使用 Swin-Tiny 且训练 300 个周期的情况下，无需任何学习率预热，即可在 ImageNet-1K 上实现 82.7% 的顶级-1 准确率。
• LipsFormer-CSwin-Tiny 在仅 4.7G FLOPs 和 24M 参数下，也无需学习率预热，达到了 83.5% 的顶级-1 准确率。
• 推导出 LipsFormer 的 Lipschitz 常数的理论上限，并证明其在初始化时为有限值，提供了原则性的稳定性保障。
• 实验结果证实，当网络在无预热条件下训练时，其保持稳定且收敛更快，验证了 Lipschitz 连续性的假设。
• 所提出的组件——CenterNorm、缩放余弦注意力、加权残差捷径和谱初始化——可作为即插即用的替代方案，直接应用于现有 Vision Transformer 架构中。
• LipsFormer 在训练稳定性和泛化能力方面优于标准 Swin 和 CSwin 基线模型，尤其在深层架构中表现更优，且无需精细的超参数调优。

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