[論文レビュー] LiteTransformerSearch: Training-free Neural Architecture Search for Efficient Language Models

Mojan Javaheripi, Gustavo Henrique de Rosa+7 more|arXiv (Cornell University)|Mar 4, 2022

Topic Modeling被引用数 5

ひとこと要約

LiteTransformerSearchは、推論誤差を予測するための高相関Proxyとしてデコーダーのパラメータ数を用いることで、トレーニング不要の神経ネットワーク構造探索（NAS）手法を提案する。これにより、ターゲットデバイス上で直接、高速かつハードウェアに最適化されたモデル探索が可能となり、ベースラインモデルと比較して最大2.5倍の高速化と2.0倍のメモリ使用量削減を達成しながら、14のタスクで同等またはそれ以上の性能を発揮する。

ABSTRACT

The Transformer architecture is ubiquitously used as the building block of large-scale autoregressive language models. However, finding architectures with the optimal trade-off between task performance (perplexity) and hardware constraints like peak memory utilization and latency is non-trivial. This is exacerbated by the proliferation of various hardware. We leverage the somewhat surprising empirical observation that the number of decoder parameters in autoregressive Transformers has a high rank correlation with task performance, irrespective of the architecture topology. This observation organically induces a simple Neural Architecture Search (NAS) algorithm that uses decoder parameters as a proxy for perplexity without need for any model training. The search phase of our training-free algorithm, dubbed Lightweight Transformer Search (LTS), can be run directly on target devices since it does not require GPUs. Using on-target-device measurements, LTS extracts the Pareto-frontier of perplexity versus any hardware performance cost. We evaluate LTS on diverse devices from ARM CPUs to NVIDIA GPUs and two popular autoregressive Transformer backbones: GPT-2 and Transformer-XL. Results show that the perplexity of 16-layer GPT-2 and Transformer-XL can be achieved with up to 1.5x, 2.5x faster runtime and 1.2x, 2.0x lower peak memory utilization. When evaluated in zero and one-shot settings, LTS Pareto-frontier models achieve higher average accuracy compared to the 350M parameter OPT across 14 tasks, with up to 1.6x lower latency. LTS extracts the Pareto-frontier in under 3 hours while running on a commodity laptop. We effectively remove the carbon footprint of hundreds of GPU hours of training during search, offering a strong simple baseline for future NAS methods in autoregressive language modeling.

研究の動機と目的

• エッジやリアルタイムアプリケーション向けに、推論誤差、レイテンシ、メモリ使用量のバランスを取った効率的なTransformerアーキテクチャを見つける課題に対処すること。
• NAS中に高価なトレーニングを必要としないように、モデル性能のトレーニング不要Proxyを特定すること。
• タスクパフォーマンスとハードウェア制約の両方を最適化する、デバイス内でのマルチ目的NASを可能にすること。
• 数百時間にわたるGPUトレーニングを回避することで、NASの炭素足跡を削減すること。
• 将来の自動逐次言語モデリングにおけるNASのシンプルでスケーラブルなベースラインを提供すること。

提案手法

• 多様なTransformerアーキテクチャにおいて、デコーダーのパラメータ数と推論誤差の間には高い相関関係があるという経験的事実を活用し、モデルランク付けのゼロコストProxyとしての有効性を示している。
• 進化計算アルゴリズムを用いてNASを実行し、パラメータ数Proxyに加え、レイテンシとピークメモリ使用量の実際のターゲットデバイス測定値を統合している。
• GPUやスーパーネットトレーニングを一切不要とせず、一般のハードウェア（例：ラップトップ）上で全探索パイプラインを実行している。
• 推論誤差、レイテンシ、メモリ使用量の3つを同時に最適化するモデルのパラトリア・フロンティアを構築している。
• 各Transformer層が異なるハイパーパrameter（例：ヘッド数、d_model、d_ff）を有する非均質な探索空間をサポートしている。
• 探索中に実際のハードウェア測定値を用いることで、シミュレーションの不正確さを回避し、ターゲットデバイスの正確性を保証している。

実験結果

リサーチクエスチョン

• RQ1 デコーダーのパラメータ数は、自動逐次Transformerにおける推論誤差を予測する信頼できるトレーニング不要Proxyとして機能するか？
• RQ2 パラメータ数と推論誤差の相関関係は、非均質的および均質的なTransformerアーキテクチャの両方でどの程度成立するか？
• RQ3 トレーニング不要のNAS手法は、探索コストと炭素足跡の削減を図りながらも、競争力のあるパフォーマンスを達成できるか？
• RQ4 デバイス内でのハードウェアに最適化されたNASは、実世界の展開に適したパラトリア・オプティマルモデルを効果的に特定できるか？
• RQ5 提案手法は、多様なハードウェアプラットフォームにおいて、レイテンシ、メモリ、精度の点で既存のベースラインを上回ることができるか？

主な発見

• 均質的および非均質的探索空間の両方において、デコーダーのパラメータ数と推論誤差の順位相関係数が一貫して高く（r > 0.9）あり、Proxyの信頼性が裏付けられた。
• LiteTransformerSearchは、GPUやトレーニングを一切使用せず、一般のラップトップ上で16層のGPT-2およびTransformer-XLにおいて、3時間未満で推論誤差、レイテンシ、メモリ使用量の3次元パラトリア・フロンティアを完全に抽出した。
• 16層のGPT-2およびTransformer-XLにおいて、同じ推論誤差を達成しながら、最大2.5倍の高速実行時間と2.0倍のピークメモリ使用量削減を達成した。
• ゼロショットおよびワンショット設定において、LTSモデルは14のタスクで350MパラメータのOPTモデルを上回り、平均的な精度が高く、最大1.6倍の低レイテンシを達成した。
• 数百時間にわたるGPUトレーニングを排除することで、NASの炭素足跡を大幅に削減し、トレーニングベースの手法に対する持続可能な代替手段を提供した。
• この手法は、外部計算リソースを必要とせず、ターゲットデバイスにネイティブにデプロイ可能であり、リアルタイムかつハードウェアに最適化されたアーキテクチャ探索を可能にした。