NAS-Driven Hardware Accelerator Exploration for Edge AI and Quantization Effects on the Pareto Space

究極のロボットの脳を構築しようとしていると想像してください。ただし、厳格なルールがあります。それは、スマートウォッチやドローンのような、非常に小さくバッテリー駆動のガジェットの中に収まらなければならないというルールです。これが「エッジAI」の世界です。そこでは、人工知能は巨大で電力を大量に消費するクラウドサーバーではなく、小さなデバイスの中で生きています。これを実現するために、エンジニアは難しいパズルを解かなければなりません。顔や音を認識できるほど賢い脳でありながら、数分でバッテリーを使い果たしてしまうことなく動作するほど小さく、かつ高速である必要があります。

完璧な脳のデザインを見つけるために、科学者たちは「ニューラル・アーキテクチャ探索（NAS）」と呼ばれる手法を用います。これは、自動化された超高速のシェフのようなものです。彼らは、最高の味（最高精度）を実現しつつ、材料（計算量）を最小限に抑えるための、何百万もの異なるレシピ（ニューラルネットワークのデザイン）を試行錯誤します。通常、これらのシェフは「フル精度」の材料を使って料理をします。これは、正確で重い計りを使用することに似ています。しかし、小さなガジェット向けには、「量子化（クオンタイゼーション）」へと切り替えなければなりません。これは、より小さな、軽い計りを使用すること（数値を表現するために使用するビット数を減らすこと）に相当します。ここで大きな疑問が生じます。もし計りの種類を変えてしまったら、キッチンで完璧に見えたレシピは、現場（実際のデバイス）でも同じ味になるのでしょうか？本論文はこのまさにその問いを掘り下げ、測定システムの変更が最適なロボットの脳の探索にどのように影響するか、そしてその脳を特殊な再構成可能チップにどのように適合させるかを探求しています。

---

小さな脳への3段階の旅

この論文の著者たちは、AIをハードウェアに適合させる問題を解決するために、巧妙な3ステップのパイプラインを提案しています。彼らはこれを「NAS-then-quantized（NASの後に量子化を行う）」アプローチと呼んでいます。これは、コンピュータ上で車の設計を行い、次にその車がデコボコの未舗装路でどのように走行するかをテストし、最後にその地形に合わせてエンジンを構築するプロセスに似ています。

ステージ1：ハードウェアに依存しないフロントエンド
まず、チームは「サロゲート（代理）」モデルを使用します。これは、実際に設計を構築することなく、その設計がどれほど優れているかを推測するスマートな予測モデルです。彼らは、15,625種類もの異なるニューラルネットワークのデザイン（NAS-Bench-201というデータセットから）を含む膨大なライブラリを用いて、この予測モデルを訓練します。この段階は「ハードウェア非依存」です。つまり、特定のチップについては考慮せず、精度と必要な演算数（FLOPs）に基づいて最適なデザインを探します。彼らは2つの探索戦略を使用します。一つは単純な「ランダム探索」（ランダムにデザインを選択すること）、もう一つは「多目的進化アルゴリズム」（自然選択を模倣して、より優れたデザインを進化させること）です。このステージを終えると、上位10個の候補リストが得られます。

ステージ2：量子化の架け橋（クオンタイゼーション・ブリッジ）
次に、現実との照らし合わせが行われます。チームはこれら10個のトップ候補を取り出し、事後学習量子化（PTQ）という手法を用いて、それらを縮小します。具体的には、通常の32ビットではなく、INT4（1つの数値をわずか4ビットで表現）へと切り替えます。これは、高解像度の写真を古い携帯電話に収まるように圧縮することに似ています。彼らはこれを行うために「Brevitas」というツールを使用します。

ここからが興味深いところです。モデルを圧縮すると、第一段階で素晴らしく見えたデザインが、失敗したり悪化したりすることがあります。「量子化の架け橋」はフィルターとして機能します。圧縮にどれだけ耐えられるかに基づいて、デザインの順位を再決定するのです。もしデザインが崩壊してしまえば、それは脱落となります。もし生き残るデザインが一つもなければ、システムはステージ1に戻ってやり直します。彼らがINT4を選択したのは、非常にアグレッシブな手法ではあるものの、わずかな追加訓練（わずか2エポック）を行うことで、モデルが精度を大幅に回復できることを見出したからです。精度は9.49%という低値から83.53%まで跳ね上がり、元のモデルの情報が大部分保持されていることが証明されました。

ステージ3：ハードウェアを意識したバックエンド
最後に、生き残ったデザインには「家」が必要です。チームは、これらの圧縮された脳をCGRA（Coarse-Grained Reconfigurable Array：粗粒度再構成可能配列）と呼ばれる特殊なタイプのハードウェアにマッピングします。CGRAは、あらゆる形に合わせてブロックを並べ替えられるレゴボードのようなものだと考えてください。彼らは、ハードウェアの構成（プロセッシングユニットの数、メモリの深さ、データパスなど）を変更しながら、異なる配置を探索するために進化アルゴリズムを使用します。彼らは物理的なチップを構築するのではなく、「フリー・アナリティカル・オラクル（自由解析的神託）」、つまり数学的なシミュレーターを使用して、各構成がどれほど高速で効率的かを予測します。そして、遅延を最小限に抑え、ハードウェアのリソース使用を最大化する構成を選択します。

大きな驚き：古い地図を捨ててはいけない

この論文における最もエキサイティングな発見は、単なる3段階のパイプラインではありません。それは、「フル精度」の世界と「圧縮された」世界との関係について彼らが学んだことです。

通常、モデルを圧縮する際、圧縮された（INT4の）世界専用の新しい予測器を訓練して、ゼロから探索を開始する必要があると考えてしまいがちです。著者らはこのアイデアを検証しました。彼らは2つのアプローチを比較しました。

1. 専用のINT4サロゲート： 圧縮されたデータに特化して訓練された予測器。
2. FP32ゼロショット・サロゲート： 元のフル精度データで訓練され、それをそのまま使って圧縮モデルの性能を推測する予測器。

結果は驚くべきものでした。シミュレーションにおいて、FP32ゼロショット・サロゲートは、専用のINT4サロゲートよりも優れた性能を示しました。 探索戦略を実行した際、FP32の予測器は、専用のINT4予測器よりも、より幅広い最適なトレードオフ（パレート空間）をカバーするデザインを見つけ出しました。

なぜでしょうか？ 著者らは、フル精度のデータの方が「クリーン」でノイズが少ないからだと示唆しています。たとえ圧縮された世界が異なっていても、どのデザインが「良く」、どのデザインが「悪い」かというランキングは、驚くほど似通っています。論文では、デザインの正確な順序は変わるものの（約21%の確率で、フル精度で優れていたデザインが圧縮後に劣るようになる）、最適なデザインの全体的な構造は十分に安定しており、古い地図（FP32の予測器）が依然として信頼できるガイドとして機能していると述べています。

パズルの安定性

チームは単に推測したのではなく、その混沌を測定しました。彼らは全15,625個のアーキテクチャを調査し、INT4に切り替えたときに「パレートフロント（絶対的な最良のデザインのリスト）」がどれほど移動するかを確認しました。

• シフト（移動）： トップデザインのリストは完全に再編成されました。元のトップデザインのうち、圧縮後の世界でトップとして生き残ったものは一つもありませんでした。
• フリップ（逆転）： 約21.73%の確率で、デザイン間の関係が逆転しました（フル精度ではデザインAがデザインBより優れていたが、圧縮後はデザインBの方が優位になった）。
• 相関： このような混沌にもかかわらず、全体的なランキングの相関関係は0.6655と中程度に保たれていました。これは、優れたデザインの一般的な「形状」が維持されていることを意味します。

また、より小さくシンプルなモデル（演算数が少ないもの）は、圧縮に対して非常に敏感であり、精度を失いやすいことも分かりました。一方で、より大きく複雑なモデルは驚くほど堅牢でした。これは、その構造が、小さな数値によって導入されるエラーを自然に平滑化しているためと考えられます。

ハードウェアの結果

最後に、彼らは最も優れた生存デザインのうち3つを選び、CGRAハードウェアにマッピングしました。その結果、デザイン自体は異なっていても、それらはすべて全く同じ最適なハードウェア構成、すなわち16行66列のプロセッシングユニットと特定のメモリの深さを持つグリッドへと収束することが分かりました。これは、この種のタスクにおいては、ニューラルネットワークのデザインが同様の構成要素を使用している限り、ハードウェアの「スイートスポット」は非常に一貫していることを示唆しています。

結論

この論文は、小さなデバイス向けにAIモデルを縮小したいとき、毎回車輪を再発明する必要はないことを示唆しています。フル精度で最適なデザインを見つけ、それを圧縮テストでフィルタリングし、それを柔軟なハードウェアにマッピングするという3段階のプロセスを経ることで、精度とサイズの完璧なバランスを効率的に見つけることができます。最も重要なことは、フル精度の探索ツールが圧縮された世界においてもガイドとして信頼できることを示しており、私たちのポケットの中に強力なAIを詰め込むための探求において、時間と労力の節約につながるということです。