[論文レビュー] SLAB: Efficient Transformers with Simplified Linear Attention and Progressive Re-parameterized Batch Normalization

Jialong Guo, Xinghao Chen+2 more|arXiv (Cornell University)|May 19, 2024

Advanced Memory and Neural Computing被引用数 7

ひとこと要約

この論文は、Laye rNormを再パラメータ化されたBatchNorm (RepBN)へ段階的に置換し、単純化された線形アテンション (SLA) を組み合わせて、視覚タスクと言語タスクの精度を維持しつつ待ち時間を低減するより効率的なトランスフォーマーを構築するSLABを提示します。

ABSTRACT

Transformers have become foundational architectures for both natural language and computer vision tasks. However, the high computational cost makes it quite challenging to deploy on resource-constraint devices. This paper investigates the computational bottleneck modules of efficient transformer, i.e., normalization layers and attention modules. LayerNorm is commonly used in transformer architectures but is not computational friendly due to statistic calculation during inference. However, replacing LayerNorm with more efficient BatchNorm in transformer often leads to inferior performance and collapse in training. To address this problem, we propose a novel method named PRepBN to progressively replace LayerNorm with re-parameterized BatchNorm in training. Moreover, we propose a simplified linear attention (SLA) module that is simple yet effective to achieve strong performance. Extensive experiments on image classification as well as object detection demonstrate the effectiveness of our proposed method. For example, our SLAB-Swin obtains $83.6\%$ top-1 accuracy on ImageNet-1K with $16.2$ms latency, which is $2.4$ms less than that of Flatten-Swin with $0.1\%$ higher accuracy. We also evaluated our method for language modeling task and obtain comparable performance and lower latency.Codes are publicly available at https://github.com/xinghaochen/SLAB and https://github.com/mindspore-lab/models/tree/master/research/huawei-noah/SLAB.

研究の動機と目的

• トランスフォーマーの正規化とアテンションモジュールの計算ボトルネックを特定する。
• 推論時に安定したBNへ置換できるProgressive Re-parameterized BatchNorm (PRepBN)を提案する。
• 計算コストを削減しつつ性能を維持するSimplified Linear Attention (SLA)を開発する。
• 画像分類、物体検出、言語モデリングタスクでSLABの効率を実証する。
• PRepBNとSLAがLNベースまたは既存の線形アテンション・トランスフォーマーと比較して遅延を低くし、精度が同等または向上することを示す。

提案手法

• 訓練中にLayerNormをRepBNへ段階的に置換する。LN優位のトランスフォーマーからBNベースへ移行するために減衰制御のガンマを用いる。
• RepBN(X) = BN(X) + η X を導入。ηは学習可能なパラメータで、訓練後の標準BNへの後パラメータ化を実現する。
• RepBNを標準BN形式へ変換する推定式（再パラメータ化の補助定理）を提供する。
• Progressive LN -> RepBN: PRepBN(X) = γ*LN(X) + (1-γ)*RepBN(X) を訓練ステップに従ってγが減衰するように定義する。
• Simplified Linear Attention (SLA)を提案する：Sim_SLA(Qi, Kj) = ReLU(Qi) ReLU(Kj)^T を用い、その後正規化された集約とローカル強化のための深さ方向畳み込みを適用する。
• SLAの線形時間計算量とハードウェア適合性を示し、アテンション分散をアテンションマップで視覚化できることを示す。

実験結果

リサーチクエスチョン

• RQ1 LayerNormを段階的に訓練されたRepBNへ置換することは推論遅延を低減しつつ精度を損なわないか？
• RQ2 提案されたSLAは既存の線形アテンション手法の性能と同等またはそれを上回り、計算コストを低く抑えられるか？
• RQ3 PRepBNとSLAは視覚タスクのDeiT、PVT、Swinの異なるバックボーンや言語モデリングでどのように相互作用するか？
• RQ4 標準ベンチマークでPRepBNとSLAを組み合わせた場合の具体的な精度-遅延のトレードオフはどうか？

主な発見

Method	FLOPs	Latency (ms)	Top-1 Acc (%)
Flatten-DeiT-T	1.1 G	15.2	74.1%
SLAB-DeiT-T	1.1 G	9.6	74.3%
Flatten-DeiT-S	4.4 G	15.5	80.4%
SLAB-DeiT-S	4.4 G	10.4	80.0%
Flatten-PVT-T	2.0 G	10.8	77.8%
SLAB-PVT-T	2.0 G	8.0	76.5%
Flatten-CSwin-T	4.3 G	32.4	83.1%
SLAB-CSwin-T	4.3 G	29.3	82.8%
Flatten-Swin-T	4.5 G	10.9	82.1%
SLAB-Swin-T	4.5 G	8.7	81.8%
Flatten-Swin-S	8.8 G	18.6	83.5%
SLAB-Swin-S	8.7 G	16.2	83.6%

• PRepBNは漸進的なLN遷移を伴い、精度を向上させつつBNベースのトランスフォーマーの推論遅延を低減する。
• SLAB-Swin-TはImageNet-1Kで83.6%のTop-1精度を達成し、待機時間は16.2 msでFlatten-Swinより2.4 ms低減、精度はやや高い。
• SLAはバックボーン全体で待機時間を大幅に削減し、Flattenトランスフォーマーと比較して競争力のある精度を維持する。
• 複数のバックボーンにおいて、SLAB派生はスループットの改善や遅延の低減を伴いつつ精度が同等または向上（例: SLAB-DeiT-TはTop-1 74.3%で9.6 ms、Flatten-DeiT-Tは74.1%で15.2 ms）
• 言語モデリングとLLaMA-350M実験で、PRepBNは推論遅延を低減しスループットを高めつつ類似の困惑度を示す。