アップルがLLMのようにスケーラブルな大規模自己回帰画像モデルを開発
3つの要点
✔️ LLMの画像版として、自己回帰学習型の画像モデルAIMを提案
✔️ 事前学習した画像特徴量の質はモデル規模とデータの質に従い向上し、下流タスクの性能は事前学習性能に従い向上
✔️ 20億枚の画像でAIMの70億パラメータを事前学習しImageNet-1kタスクで精度84%を達成に加え、性能飽和の兆しなし
Scalable Pre-training of Large Autoregressive Image Models
written by Alaaeldin El-Nouby, Michal Klein, Shuangfei Zhai, Miguel Angel Bautista, Alexander Toshev, Vaishaal Shankar, Joshua M Susskind, Armand Joulin
(Submitted on 16 Jan 2024)
Comments: Published on arxiv.
Subjects: Computer Vision and Pattern Recognition (cs.CV)
code：
本記事で使用している画像は論文中のもの、紹介スライドのもの、またはそれを参考に作成したものを使用しております。
はじめに
2012年、畳み込みニューラルネットワーク（CNN）をベースにしたAlexNetがトロント大学から発表され、画像認識コンペで優勝しました。この成功により、層を深くすればするほど画像認識精度が向上するという期待が高まり、一気に深層学習の研究が盛んになりました。
層数を深くすると、モデルパラメータ学習時の勾配消失が原因で性能が劣化する問題がありましたが、2015年、マイクロソフトリサーチが発表したResNetのショートカット接続の提案により、152層まで層数を増大しても、性能向上することが確認されています。
このように画像認識タスクが、モデルの大規模化をけん引してきたわけですが、2017年、自然言語処理タスクを解く言語モデルとしてTransformerが発表されました。
2022年、Transformerをベースにした大規模言語モデルであるChatGPTのプロトタイプが、Open AIから公開されました。ChatGPTは、あらゆる質問に対して非常にもっともらしい回答をすることから大きな注目を集め、汎用人工知能到来を予期させました。さらに、モデルパラメータ数の増大に対する性能限界の兆しはなく、モデルの大規模化競争が開始し、一気に大規模言語モデルの研究が盛んになりました。
歴史を振り返ると、画像認識研究において、深層学習の研究が発展し、次いで、言語処理研究において、生成モデル、特にTransformerの大規模化が進展してきたわけですが、今回紹介する論文は、言語処理研究の成果を画像認識研究へ一般化できないかという話になります。
具体的には、このTransformerをベースに大規模”言語”モデルではなく大規模”画像”モデルをつくる話です。ただし、これまでTransformerが画像認識に応用されてこなかったわけではありません。
Transformerを画像認識に応用した有名な技術として、Vision Transformer（ViT）があります。このモデルは、2021年にGoogleから発表されています。この研究により、事前学習用の画像データが十分にあれば、CNNをTransformerで代替可能なことが示されました。
Vision Transformerは、大規模言語モデルと同じくTransfomerをベースにしているわけですが、その事前学習の方法が大規模言語モデルと異なります。どういうことかというと、大規模言語モデルでは、事前学習として自己回帰学習、つまりは外からラベルは与えずに、ある文章のある単語までを示して、次の単語を予測するという学習をするのですが、Vision Transformerでは画像認識の正解ラベル（犬や猫などの分類クラス）を訓練画像とセットで与えて学習を行っています。
今回紹介する論文では、モデルのベースはVision Transfomerと同じですが、大規模言語モデルと同様に、自己回帰学習することでモデルパラメータ数、訓練データ数の増大と共に画像認識精度を向上可能な技術が開発できないかをアップルが検証した論文になります。
それでは、提案された自己回帰画像モデルAIM（Autoregressive Image Models）とその評価結果を説明していきます。
自己回帰画像モデル (AIM)
事前学習の流れ
提案手法である、自己回帰画像モデル（AIM）による事前学習の流れを図1に示します。
Transformerの自己回帰学習は、与えられた文章について、左から順に今までに示された単語から都度次の単語を予測できるように学習するものです。このTransformerの自己回帰学習を画像に応用するには、画像を単語からなる文章として表現する必要があります。そこで、AIMでは、図1のように入力画像を重なりのないパッチ画像（部分領域）に分割し、そのパッチ画像に順番をつけ、次のパッチ画像当てを学習します。
元の画像を分割した各パッチ画像を線形写像により次元削減し特徴抽出します。パッチ画像の特徴をTransformerに入力し、自己注意処理（文脈を加味したパッチ特徴の更新。後述するプレフィックス因果マスクを用いた自己注意処理）を行います。その後、抽出された画像特徴に基づき、MLPで次のパッチ画像を予測します。
モデルはラスター順（図1の各パッチ画像に割り振られた数字の順。上の行から順に読まれます）に次のパッチ画像を予測します。
図1のMLPの出力は、パッチ画像1を与えられパッチ画像2を予測、パッチ画像1、2を与えられパッチ画像3を予測する関係が示されています。パッチ画像が与えられないと予測しようがないので、パッチ画像の予測結果は2から始まり、9までパッチ画像を与えられたら、次に予測するパッチ画像はないので、9で終わっています。
先ほどのモデルが出力する2番目のパッチ画像から9番目のパッチ画像の予測と正解の誤差関数（自己回帰目的関数）を式1に示します。
xハットはモデルパラメータθを持つAIMのk番目のパッチ画像の予測値であり、ハットなしは正解のパッチ画像の値です。各パッチ画像の予測と正解の画素値ベクトルの誤差の二乗和の平均値を最小化するようにθを学習します。（図1の通り、モデルの出力は2番目から9番目までのパッチだとすると、実質k=2から始まると考えるべきかもしれません。）
プレフィックス因果マスク
次のパッチ画像を予測するという問題上、これまでに示されたパッチ画像から次のパッチ画像を予測するという関係を守るべきで、次のパッチ画像に関する情報を得られてしまうと、カンニングのような形になってしまいます。Transformerの自己注意処理は、文章内の別の単語との重み付き和を計算して、各単語の特徴を更新するものですので、単純にすべての単語との重み付き和を計算すると、このカンニングが起きてしまいます。
そこで、因果マスクは、事前学習時の次のパッチ画像の予測誤差を計算する際、2番目のパッチ画像を予測するときに1番目のパッチ画像のみを利用し、3番目のパッチ画像を予測するときに、1番目と2番目のパッチ画像を用いて、今までに示されたパッチ画像の間のみで自己注意処理をします。
しかし、因果マスクを用いると、下流タスク（本論文では、事前学習時のタスクとは別の新たな画像認識タスク）に適応させた際にうまくいかない問題があります。
事前学習においては、自己回帰タスクですので、未来のパッチを見せてはいけないという制約があります。一方で、下流タスク（画像認識タスク）においては、認識対象の画像に関するすべてのパッチ画像を一度に見ても構いません。むしろ、すべてのパッチ画像の特徴を考慮したうえで、画像認識判断を下すのが自然です。しかし、事前学習時に、単純に因果マスクを用いてしまうと、複数の画像パッチを考慮した特徴がうまく学習できません。したがって、下流タスクも考慮した事前学習が望まれます。
そこで、本論文では、プレフィックス因果マスクを利用します。プレフィックス因果マスクは、一様分布で決めたあるパッチ画像枚数＝プレフィックス長さまでのパッチ画像をフルで利用して自己注意処理を行います。それ以降は、因果マスクと同様に自己注意処理するものです。予測誤差を計算する際には、プレフィックス長さを超えるパッチ画像の予測を対象とします。
プレフィックス因果マスク有の事前学習と下流適応のイメージを図2に示します。
図の丸や線が何を表すか本論文には説明がないので推測になりますが、ニューラルネットワークの模式図に似ています。
図の下から、入力層、中間層、出力層が5次元の3層のニューラルネットワークと考えると、ここで入力層の5次元が、5つのパッチに当たると考えることができます。つまり、横並びの丸は、5つのパッチ画像の並びと同じだと考えられます。5つのパッチ画像の自己注意処理による特徴更新時にどのように他のパッチ画像が絡んでくるかをニューラルネットワーク風に示していると考えられます。
ここで、図2の左の絵がプレフィックス因果マスク有の事前学習の状況です。プレフィックス長さが3のプレフィックス因果マスクが適用されており、左から数えて3つの画像パッチまでは全結合されているようなニューラルネットワークと考えることができます。
つまり、左から1番目、2番目、3番目の画像パッチはそれぞれ1番目から3番目の画像パッチに基づいて特徴が更新されている（1番目、2番目の画像パッチに関しては未来の一部のパッチで特徴を更新しているような状態）と考えることができます。
4番目と5番目の画像パッチは、自分自身とそれ以前の画像パッチで更新される（因果マスクが適用されて、4番目にとっての未来の画像パッチである5番目の画像パッチを用いて4番目の画像パッチの特徴が更新されない状態）イメージを表しているのだと考えられます。
図2の右の絵が下流適応のイメージです。下流適応は、認識対象の画像すべてのパッチが全結合されていると考えるのが自然なタスクであるので、全結合が望ましいですが、事前訓練は、全結合を常に許すと常にカンニング状態で学習することになってしまいます。
ですので、プレフィックス因果マスクを使った場合、下流適応を想定した全結合の状態での特徴更新を一部保ったような因果マスクで事前学習を行っていると考えることができます。
MLP予測ヘッド
Transformerの事前学習では、様々な下流タスクに通用するような汎用的な画像特徴を学習できるのが理想です。事前学習時の目的関数に特化した画像特徴が学習されると、下流タスクに適応しにくくなってしまいます。
そこで、下流タスクへの適応能力を上げるため、MLP（Multi Layer Perceptron）予測ヘッドを最終層に接続して、事前学習します。MLP予測ヘッドは各パッチ画像を独立に処理します。
そして、下流タスクへの適応時は、事前訓練時のMLP予測ヘッドは捨て、残るTransformer部を汎用特徴抽出器として用います。
下流適応
大規模モデルの事前学習は、大幅に計算資源を消費します。微調整も骨が折れます。
そこで、下流タスクの学習時は、事前学習で学習したTransformerの重みは固定します。新たに接続する分類ヘッドだけを学習します。これにより、少数の下流タスクデータセットへの過学習防止と下流適応コストを削減できます。
AIMの評価結果
AIMがモデルパラメータ数、訓練データ数の増大に対してスケーラブルな性能を示すか評価します。
20億枚の未整理の画像でAIMを事前学習し、15の下流タスク（細粒度認識、医療画像、衛星画像、自然環境、インフォグラフィックなどの多様な画像認識(=画像のクラス分類)ベンチマーク）に適応させ、15の下流タスクの平均性能を評価した結果を図3に示す。
図3のグラフは、AIMのモデルパラメータ数に対するスケーラブル性を検証したグラフです。横軸が事前学習におけるバリデーションロス、縦軸が15の下流タスクの平均性能である。AIM-*Bと書かれたグラフの各点は、AIMのモデルパラメータ数を表します。右に行くほど、数値が上昇し、順に6億、10億、30億、70億モデルパラメータ数の場合を示します。
つまり、AIMのモデルパラメータ数を増大させた場合の、事前学習性能と下流タスク性能の関係を表すグラフです。AIMのモデルパラメータ数の増大に従い、事前学習性能が向上し、下流タスクも性能向上していることが確認できます。
ちなみにですが、ImageNet-1kベンチメークにおいては、AIMは70億モデルパラメータにおいて、84.0という精度を示しており、MAEというオートエンコーダを用いた既存手法の精度82.2を上回っています。一方で、DiNOv2という手法は精度86.4で、AIMを上回っています。これに対して、本論文では、DiNOv2の評価結果については、高解像度の訓練画像を用いているので、それだけで1から1.5%は上がると述べています。（AIMでも同様に高解像度画像を学習させれば、DiNOv2よりも性能向上できるというニュアンスなのかは不明です。）
図4のグラフは、AIMのモデルパラメータ数に対するスケーラブル性を検証したグラフです。
図4の横軸はAIMのユニークな訓練画像数、縦軸はモデルパラメータ数が6億のときの、15の画像認識ベンチマークにおけるAIMの平均性能です。つまり、モデルパラメータ数が6億のときのAIMの訓練データ数と下流タスク性能の関係を示しています。AIMの訓練データ数が100万枚、1億枚、20億枚と増加するにつれて、認識精度が向上していくのが確認できます。認識精度が飽和する様子は見られません。
図3、図4より、AIMがモデルパラメータ数、訓練データ数の増大に対してスケーラブルに画像認識精度が向上することが示唆されています。
おわりに
本論文では、画像認識においても、LLMのようにモデルパラメータ数、訓練データ数に従った、事前学習性能、下流タスク性能の向上を実現するため、Vision Transformerをベースに、自己回帰目的関数を用いた事前学習が可能な手法AIM(Autoregressive Image Models)を提案しました。
単純な因果マスクを用いた自己注意処理を行うTransformerで自己回帰目的関数を学習する場合、画像認識対象画像全体の関係性をうまく使えず、下流タスクへの適応がうまくいかないため、プレフィックス因果マスクを用いた自己注意処理を組み合わせています。
また、自己回帰目的関数に過適合しない汎用的な画像特徴を学習するため、Transformer部を汎用画像特徴抽出器、その後段につけたMLP部を次パッチ画像予測器とする構成を採用しています。
さらに、下流タスク適応時は適応コストを削減するため、事前学習時のMLP部を画像分類MLPに付け替え、下流タスクに適応するよう画像分類MLPのみを学習しています。
評価の結果、訓練データ数20億枚、モデルパラメータ数70億個まで、事前学習性能、下流タスク性能が向上し、その性能向上の限界は確認されていません。
もしも、本当に画像モデルを大規模化し続ければ画像認識精度が向上し続けるのであれば、精度競争のためのモデル改善競争は基本的には終焉すると言えそうです。歴史を振り返れば、深層学習を深くすることでスケーラブルな性能向上を得ること自体大きな課題であったのですから、それが解決されたことになると思います。
そうなれば、資源をたくさん投入すれば性能向上してしまうのですから、残る研究の方向性としては、どれだけ少ない資源で画像認識できるか、その技術をどう応用するかだけになりえそうです。
しかし、大規模言語モデルと比較してしまうと、今回の画像版LLMは、どうしても画像認識（ある画像のクラス分類）に限定された能力に見えてしまうのが気になるところです。
大規模言語モデルのすごいところは、文脈内学習で、いろいろな質問に回答できてしまうことだと思いますが、今回の大規模画像モデルは、画像認識限定に見えてしまうし、下流タスク適応のための学習は必要ですので、大規模言語モデルのゼロショット（下流タスクのデータ学習なし）でも、非常に精度よく回答できてしまう期待とずれるところがあります。
画像版のLLMと言い切るにはそのようなゼロショットでいろいろなタスクが解けてしまう驚きが欲しいと感じました。
この記事に関するカテゴリー
