TAG：マルチタスク学習で一緒にトレーニングすべきタスクを知る(2/2)

１．TAG：マルチタスク学習で一緒にトレーニングすべきタスクを知る(2/2)まとめ

・TAGは1つのタスクに関してのみモデルのパラメータを更新してその影響を調べる
・ネットワーク内の各タスクが他のタスクとどのように相互作用するかについての情報を収集
・この情報を使うと特定タスク同士が常に互いに有益なであったり敵対する事がわかる

２．TAGの概要

以下、ai.googleblog.comより「Deciding Which Tasks Should Train Together in Multi-Task Neural Networks」の意訳です。元記事は2021年10月25日、Christopher Fiftyさんによる投稿です。

アイキャッチ画像のクレジットはPhoto by The HK Photo Company on Unsplash

タスク間に親和性を持つグループの構築

この課題を検討するにあたり、私たちはメタ学習(meta learning)からヒントを得ました。メタ学習とは、これまでに見たことのない新しいタスクに素早く適応できるようにニューラルネットワークを学習させる機械学習の研究分野です。

古典的なメタ学習アルゴリズムの1つであるMAML(Model-Agnostic Meta-Learning)は、「タスクのコレクション」に対するモデルのパラメータに勾配更新を適用します。その後、「更新されたパラメータ値で計算された「コレクション内のタスクのサブセット」」の損失を最小化するように元のパラメータセットを更新します。(訳注：この部分、わかりやすく意訳できてませんが、先に読み進んでから戻ってくると理解しやすいです)

この方法により、MAMLは、「現在の重み」に対して損失を最小化するのではなく、「1つ以上のステップを学習した後の重み」に対して損失を最小化する特徴表現を学習するようにモデルを訓練します。その結果、MAMLはモデルのパラメータを、現在ではなく未来に向けて最適化するため、以前に見たことのないタスクに素早く適応する能力を持つようになります。

TAGでは、同様のメカニズムを用いて、マルチタスク・ニューラルネットワークの学習動向を把握しています。具体的には、1つのタスクに関してのみモデルのパラメータを更新し、この変更がマルチタスク・ニューラルネットワークの他のタスクにどのような影響を与えるかを調べ、次にこの更新を元に戻します。このプロセスを他のすべてのタスクについて繰り返し、ネットワーク内の各タスクが他のタスクとどのように相互作用するかについての情報を収集します。その後、ネットワーク内のすべてのタスクに関してモデルの共有パラメータを更新することで、通常通りのトレーニングを行います。

これらの統計情報を収集し、トレーニング中の動向を観察すると、あるタスクは常に互いに有益な関係を示し、あるタスクは互いに敵対していることがわかります。ネットワーク選択アルゴリズムは、このデータを利用して、タスク間の親和性を最大化するようにタスクをグループ化することができます。ただし、推論の際にマルチタスクネットワークをいくつ使用するかは、実行者の選択によります。

TAGの概要
まず、同ネットワークでタスクを一緒に学習し、タスク間の親和性(inter-task affinities)を計算します。第2に、ネットワーク選択アルゴリズムは、タスク間の親和性を最大化するタスクのグループ化を見つけます。第3に、結果として得られたマルチタスクネットワークを学習し、展開します。

結果

実験の結果、TAGは非常に強力なタスクグルーピングを選択できることがわかりました。

CelebAデータセットとTaskonomyデータセットにおいて、TAGはこれまでの最先端技術と競合し、それぞれ32倍、11.5倍の速度で動作しました。Taskonomyデータセットでは、この高速化により、タスクグルーピングを見つけるためのTesla V100 GPUの使用時間が2,008時間減少しました。

結論

TAGは、1回のトレーニングでどのタスクを一緒にトレーニングすべきかを決定する効率的な手法です。この手法では、タスクがトレーニングを通じてどのように相互作用するかに注目します。特に、1つのタスクをトレーニングする際にモデルのパラメータを更新すると、ネットワーク内の他のタスクの損失値にどのような影響を与えるかに注目します。このスコアを最大化するようなタスクのグループを選択すると、モデルの性能と強い相関関係があることがわかりました。

謝辞

本研究に多大な貢献をしてくださったEhsan Amid, Zhe Zhao, Tianhe Yu, Rohan Anil, Chelsea Finnに感謝いたします。また、アニメーションをデザインしてくれたTom Small、協力的で高揚感のある研究環境を提供してくれたGoogle Research全体にも感謝しています。