ディープラーニングを軽量化する「モデル圧縮」３手法
株式会社Laboro.AI リード機械学習エンジニア 内木 賢吾
2020.2.13公開 2022.4.25更新
概 要
ディープラーニングは、機械学習の中でもとくに研究が盛んに行われているテクノロジー分野の一つです。アカデミアでの研究だけでなく、企業がAIを導入・活用するにあたってもディープラーニングを採用したAIモデルの開発が多数行われていますが、とくにエッジAIの分野では、デバイスの能力に制約がある中でいかにAIが持つ可能性を最大限に引き出せるかが、重要な検討ポイントのひとつです。
このコラムでは、モデルと呼ばれる入力に対して何かしらの処理をした結果を出力する処理を軽量化するための技術「モデル圧縮」について、その概要を紹介します。まず、モデルの軽量化が必要とされる背景について触れ、その後、モデル圧縮技術の代表的な３つの手法について説明していきます。
目 次
・高まるエッジAIへの期待
・エッジAIの限界を払拭する「モデル圧縮」
・モデル圧縮手法
１．Pruning（枝刈り）
２．Quantize（量子化）
３．Distillation（蒸留）
・まとめ
高まるエッジAIへの期待
近年、現場のデバイス（エッジデバイス）に機械学習のモデルを実装し、その場で予測などを行うエッジAIが注目を集めています。エッジAIの普及によって、例えば、製造ラインの部品検査業務でのでの異常検知などリアルタイム応答が重要となるシーン、さらには自動運転で予想される通信不安定時の予測処理が可能になるといったことが期待されています。
クラウドを利用した従来の予測処理とエッジAIとの違いを掴むため、それぞれを図に示します。図１にはクラウドを用いる一般的な予測処理の流れを、図２にエッジAIを用いる予測処理の流れを示しています。
なお、それぞれの図では学習済みのモデルを用いることを想定しています。また、通常「エッジ」という用語は「 端末の近く」という意味で用いられ、機械学習において「エッジAI」は、エッジデバイス内で取得したデータをAIで処理することを指すケースが多いです。
エッジAIを用いることで、取得した画像などの取得データを送信せずに必要なデータだけを送信するか（図2. 3-a）、エッジデバイス内で処理を完結してクラウドへのデータ送信が不要になる（図2. 3-b）ため、通信コストの削減に加えて、通信遅延や通信障害の影響を少なくすることができます。
今後、５G通信の普及によって低遅延通信が実現すれば、リアルタイム応答に対する課題はある程度解決されるはずです。しかし、それでも遅延が出る状況はあり得ますし、通信の安定性にはやはり不安が残ります。さらに、クラウドにアップロードできない機密情報を扱うといったシーンも考慮すると、エッジデバイスで処理をしたいというニーズはますます高くなることが予想されます。
このように必要とされる状況が想定されるにも関わらず、エッジAIの普及が遅れている理由はどこにあるのでしょうか。
エッジAIの限界を払拭する「モデル圧縮」
エッジAIの普及を足止めする原因として、エッジデバイスの初期コストを含めた運用コストの課題が考えられますが、乗り越えるべき最大の課題、それはエッジデバイスの処理能力の限界にあります。
IoT端末のようなエッジデバイスに用いられる機器は、クラウドなどのGPUサーバーと比較して処理能力が低く、複雑な処理に向きません。一方で、ビジネス現場で求められる画像分類や物体検出といったタスクに用いられる手法は、層を増やすことで性能が向上することが多いです。しかし、層が増えるということはモデルのパラメータが増えることになり、メモリ使用量や計算量が増加します。その結果、限られたリソースの中で処理しなければならないエッジデバイスは早く処理能力の限界に到達してしまい、クラウド環境よりも性能が落ちてしまいます。
最近ではNVIDIA TensorRTやIntel OpenVINOといったエッジAI向けの開発ツールやサービスが公開されており、NVIDIA Jetson Nanoのような小型で安価なGPU開発キットも購入できるようになりました。こうしたツールを活用することでエッジAIが力を発揮するシーンも出始めています。一方、高いレベルでの成果が求められるシーンでは、やはりエッジAIでは対応できないシーンがあることも事実です。極端な例ではありますが、2019年の12月にGoogle Researchが発表しSoTA（State-of-The-Art）を達成した画像認識モデル「Bit」（参考文献①）は、10億個という巨大なパラメータ数を持っていることで話題になりました。2022年4月に同じくGoogle Researchが発表した言語モデル「PaLM」（参考文献②）は5,400億個ものパラメータを持っており、算数の問題を解いたり、冗談の解説をしたり、多くのタスクで人の平均スコア越える性能を持つことで注目を集めました。
モデルの精度と計算量はトレードオフの関係にあります。いかにモデルの精度を保ったままメモリ使用量を削減し、処理時間を高速化できるかが、エッジAIにおいては非常に重要な検討ポイントになります。
そこで考えられる手法が、今回のテーマである「モデル圧縮」です。モデル圧縮は、計算量を小さくすることで推論時間を短縮したり、メモリ使用量の削減によってより複雑なモデルを搭載できるようにしたり、エッジデバイス上での推論を実現できるメリットを期待できる手法です。さらに、エッジデバイスだけではなく、クラウドを利用する場合も同様のメリットを受けることができます。
では、そのモデル圧縮の代表的な手法を３つ紹介していきましょう。
モデル圧縮手法
モデル圧縮について、今回は代表的な手法であるPruning（枝刈り）、Quantize（量子化）、Distillation（蒸留） の3つの手法ついて簡単に説明します。
１．Pruning（枝刈り）
ディプニューラルネットのモデルは、下図のように各ノードが密に結合しています。 そのノード間の重みが小さい箇所の接続を削除する、または影響の小さいノードを削除することでパラメータ数を削減する手法をPruning（プルーニング：枝刈り）と呼びます。
図3がPruningのイメージです。ノードや重みを削除することでパラメータ数を減少させます。記憶する必要があるパラメータが減少することによって、計算する回数が削減され、メモリ使用量が少なくなります。その結果、モデルの軽量化と処理の高速化が期待できます。
どの接続を削除するのかについては様々な手法が提案されています。「重みの絶対値が小さいものを削除する」手法が一般的で、性能も良いとされています（参考文献③）。また、Pruningによって80%〜90%程度の接続を削除できる場合もあります。なお、一般的にはPruningだけでは精度が落ちてしまうため、Pruning後に再学習を行うことで元の精度を保つことが行われます。
２．Quantize（量子化）
Quantize（クオンタイズ：量子化）とは、重みなどのパラーメータをより小さいビットで表現することで、モデルの軽量化を図る手法です。使用するビットを制限することでネットワークの構造を変えずにメモリ使用量を削減できます。使用する数字を簡単な数字に置き換えるようなイメージを持つとわかりやすいかと思います。図4にQuantizeの例を示します。
TensorFlowやPyTorchなどのディープラーニングのフレームワークでは、32ビット浮動小数点精度（float型）を使用することがほとんどです。近年の研究（参考文献④）では、８ビットの量子化であれば1%程度の性能低下であることが報告されています（図５: １×４の意味は、(重みの１ビット)×(活性化関数４bit)を意味する）。中でも使用するビットを１〜２ビットにすることで乗算をXNORゲートで表現することができるため、ASICやFPGAなどでは計算効率が大きく向上します。ただし、使用するビットを小さくするほど精度が低下する可能性が高いことに注意が必要です。
なお、学習時に重みや活性化関数のパラメータを量子化することで、より計算効率の高いネットワークを構築できることが示された報告もあります（参考文献⑤）。
３．Distillation（蒸留）
Distillation（ディスティレーション：蒸留）とは、大きいモデルやアンサンブルモデルを教師モデルとして、その知識を小さいモデル（生徒モデル）の学習に利用する方法です。これにより、大きいモデルに匹敵する精度を持つ小さいモデルを作ることが期待できます。つまり、一度学習したモデルの知識（予測結果）を別の小さいモデルに継承するということです。図６にDistillationのイメージを示します。
Distillationでは、教師モデルの学習などで使用する正解ラベルをhard target、教師モデルの出力（スコア分布）をsoft targetと呼びます。生徒モデルを教師モデルに近づけるように学習するときに、hard tagetとsoft targetの両方を用いる場合やsoft targetのみを用いる場合があります。
soft targetを用いることがDistillationの大きな特徴です。これにより、hard targetに含まれない情報を含めて学習できるというメリットがあります。図６の例において、教師モデルの出力スコアが自動車：0.4、犬：0.0であることから「この画像は、車に似ている部分があり、犬に似ている部分はない」と捉えることができます。hard targetからは「この画像は、電車である」という情報しか含まれていません。つまり、soft targetから「電車とは、車輪があって、窓があって、…」のような一般的な概念を学習できることが期待できます。このように、学習自体に工夫を加えており、さらに情報量を増やしながらモデルを軽量化することがPruningやQuantizeとの違いです。
他にも、Distillationはモデル圧縮以外の用途にも応用できることが知られています。その一つにData Distillation（データ蒸留）と呼ばれる手法があります（参考文献⑥）。Data Distillationは、正解ラベルのないデータに対する教師モデルの予測結果を正解ラベルとして扱い、生徒モデルの訓練を行う手法です。この手法により、従来では利用できなかった正解ラベルのないデータを有効活用でき、かつ、精度向上に有効であることが示されています。
まとめ
本コラムでは、エッジAIの視点からモデル圧縮技術の必要性を説明し、代表的な手法であるPruning（枝刈り）、Quantize（量子化）、Distillation（蒸留）の3つの手法ついて説明しました。
必ずしも全ての状況においてモデル圧縮技術が必要になるわけではありませんが、モデルを軽量化することによってディープラーニングを活用できる可能性が広がることは間違いありません。精度と計算量のトレードオフを解消するため、モデル圧縮はディープラーニングにおける重要な手法のひとつと言えます。
参考文献
① Big Transfer (BiT): General Visual Representation Learning
② PaLM: Scaling Language Modeling with Pathways
③ Deep Compression: Compressing Deep Neural Networks with Pruning, Trained Quantization and Huffman Coding
④ A Survey of FPGA Based Neural Network Accelerator
⑤ Quantized Neural Networks: Training Neural Networks with Low Precision Weights and Activations
⑥ Data Distillation: Towards Omni-Supervised Learning
コラム執筆者
機械学習エンジニア 内木 賢吾
名古屋大学大学院 工学研究科 博士前期課程修了。在学中に質問応答システムを研究。卒業後、ハードウェアエンジニアとして実験機器の開発やデジタル/アナログ回路設計ならびにハーネス設計を経験。その後、車載向け音声認識システムのフロントエンド処理に関する研究開発を担当。2019年7月よりLaboro.AIに参画。
