独立メカニズムを捉えるために、Transformerを分解する？

Transformer 2021年08月13日

3つの要点
✔️ 独立メカニズム仮説を取り入れたTransformer
✔️ Attention機構を活用してTransformerを複数モジュールに分解する
✔️ Transformerを用いた幅広いタスクにて有効性を確認した

Transformers with Competitive Ensembles of Independent Mechanisms
written by Alex Lamb, Di He, Anirudh Goyal, Guolin Ke, Chien-Feng Liao, Mirco Ravanelli, Yoshua Bengio
(Submitted on 27 Feb 2021)
Comments: Accepted by ICML 2021.
Subjects: Machine Learning (cs.LG); Artificial Intelligence (cs.AI)

code：

はじめに

大活躍しているZero-shot言語生成モデルGPT-3や分布外な画像生成モデルDALL-Eで姿を見せているTransformerアーキテクチャは、全てのポジション情報を一つ大きな潜在表現で学習します。ただし、これは関係しない情報までを同時に処理することを意味し、世の中に存在する独立構造を捉えることを制限してしまいます。本記事は、この問題を指摘してTransformerを改善したTransformer Independent Mechanisms（TIM）を紹介します。

TIMの重要なコンセプトとしては、「物理現象を独立したモジュールの動きとして考えられるのは、これらの背後に存在する独立したメカニズムによって支配されている」という独立メカニズム仮説があります。実はこの独立メカニズム仮説が因果推論のコミュニティにおいては大前提とされているが、深層学習で扱われる研究がまだ少なく、Independent Causal MechanismやAI-SCHOLARでも紹介させてもらったReccurent Independent Mechanismsが挙げられます。本日は、TIMがどう設計されているのか？どういった実験設定で成果を出しているのか？について順に紹介させてもらいます。

TIMのアーキテクチャ

著者たちが提案するTIMはTransformerを複数の独立したパーツに分解して、１つのポジションに複数のMechanismモジュールが配置されるようなアーキテクチャと捉えられます。

ポジション毎に３つのMechanismがある極めてシンプルな例をFigture2に示しています。ポジションの軸（系列モデルの時間軸に相当します）とMechanismモジュールの軸と２つの軸に沿って、それぞれAttention計算で情報をシェアした後、２つの軸に沿ってFFN順伝搬ネットワークで潜在変数を更新するアーキテクチャとなります。また、提案手法TIMはスタンダードのTransformerレイヤーと単純に置き換えられて、Transformerを用いたその他の手法にも簡単に応用できます。では、TIMの詳細についてアルゴリズムの4つステップに従って説明していきます。

Mechanisms間の競争

まずは、各Mechanismsモジュールの専門性（１つのモジュールが１種類の処理を行う性質）を高めるには、独自のパラメータを持ちながらAttention計算のみで情報共有を行います。その上でさらに強い帰納的バイアスを導入したいため、先行研究RIMと同様なAttentionによる関連度スコアを算出する仕組みを導入し、競争を誘発します。

具体的には、アルゴリズムのステップ1の式が示したように、各モジュールの表現(h)を１つの値に線形変換(GroupLinear)した後、Softmax計算でスコアを算出します。そのスコアを元に、Mechanismsモジュールが情報にどれだけアクセスでき、更新されるかを重み付けします。この重みは、後ほどのポジション軸に沿ったAttentionを行って潜在変数を更新する際に用います。

各Mechanismモジュールが取得したい情報を確保するには、他のモジュールの関連度スコアを抑えなければいけなく、Mechanismsモジュールの専門性を高める効果が期待されます。

ポジション軸の情報共有

次のステップ２では各Mechanismモジュールはポジション軸に沿ってAttention計算を行って、線形変換(GroupLinear)をします。ステップ1でも出現したGroupLinearは、一般的な(Linear)線形変換に対して、分割したモジュール(Group)間だけで線形変換を行う層を指します。最後にステップ１で算出したスコアを重みづけて、潜在変数hを更新して均一化します。ここで注意して欲しいのは、ポジション軸の情報共有のみで構成するアーキテクチャは、単に複数個の独立したTransformerを組み合わせたものとして捉えられます。

Mechanismモジュール軸の情報共有

各mechanismsモジュールが独立して情報を処理することをやってきたが、モジュール間の最小限な情報共有も必要と考えられます。ステップ３では、32ユニットの２headsのMulti-head Attentionのみを用いて、ほんの少しの情報だけをMechanismモジュール軸に沿ってAttention計算を行って共有します。

順伝搬で潜在変数を更新

ステップ４では、Mechanism軸とポジション軸と２つの軸でそれぞれFFN順伝搬ネットワークで線形変換を行い、潜在変数hの更新を行います。

実験

著者たちは独立したメカニズムを含むと考えられるデータセットに対して、TIMの有効性を確かめるために２つの問いに答えます。

1. TIMが合理的で有意義な専門性を持つMechanismモジュールの学習ができるのか？こちらについては、トイデータと現実の大規模な音声認識と言語処理タスクで検証します。

2. 独立したメカニズムを持つモデルを活用して、定量的な精度の向上に繋げられるか？こちらについては、音声増強やBERTのMLM等のタスクで検証します。

Transformerは幅広い領域で活用されているかつTransformerをTIMに置き換えるだけで実験できるので、読者の研究テーマに独立メカニズムが存在するかどうか、TIM適用できないかについて考えてみるのがいいかもしれません。

Image Transformer

ここでは、GPT-2のアーキテクチャにTIM取り入れた画像生成モデルを用いて、明確に異なる２つのメカニズムを有する自作タスクを解かせます。

具体的にはMNISTの数字を左に、ランダムに選んだCIFARの画像を右に存在する合成されたデータセットを用いて、励起した１つMechanismモジュールを可視化することで、片方に対して専門化できるかを評価します。

Figture 3右より、TIMはこの合成データセットにおいて異なるMechanismモジュールがそれぞれ両辺に専門化できたことが言えます。興味深いことに、学習の初期においては色の明るさに専門化したモジュールが学習が進むにつれて、左右の異なるデータセットに専門化できるようになってきます。

また、CIFAR-10データセットにおいて物体と背景（Figture 3左）に対しても専門化することができました。

音声増強

音声増強は、実世界のノイズを含んだ音声データの質を高めるタスクです。信号処理技術に基づいた従来のアプローチは、言語的な音でない部分を検出して削除することで実現します。近年Tranformerを用いた手法が従来の手法を超え、有効性を示しています。

ここで言語的意味のある音とそうでない音を明確に分解することを、異なるメカニズムによって生み出されたデータを処理することだと考えれば、TIMは適していることがわかります。

良質な音声データにノイズ音声を加えたDNSデータセットで、音質の良さの評価指標PESQを用いた実験結果をTable 3に示しています。提案手法TIMは、最新の手法PoCoNetの1/8のパラメータだけで、SOTAを達成しています。

またFigure 5の可視化から、Mechanismモジュールの独立性はレイヤーが深くなるにつれて明確に現れて、モジュールの専門化が進んでいることが読み取れます。

さらにVoiceBankデータセットの実験結果（Table 2）から、TIMに競争機構を取り入れることで入力に含まれる無関係情報の変化に影響されず、精度の向上につながることがわかります。

BERTにTIMを取り入れて改良する実験

自然言語処理（NLP）タスクにおいて、Masked Language Modelで事前学習することで多様な下流タスクで有効性を示したBERTを、最初の２層と最後の層を除いて全ての層（9/12の層）をTIMに置き換えました。また、2つのMechanismモジュールを用いてBERTと同様なパラメータを用いて、BERTとの比較実験をしました。

BERTと同様に事前学習を行い、複数のデータセットでFine-tuningをしてそれぞれの精度を比較しました。結果（Table 1）から、簡単な変化を加えただけのTIMが精度を向上させ、TIMがより良いNLPアーキテクチャであることが言えると著者は主張しています。ただし、この結果はわずか良くなった程度であり、はっきりしたことが言えなく、もっと良い結果が求められると筆者が感じております。

CATER：隠れ物体の追跡タスク

CATERは動画終了時に追跡したい物体が、6x6グリッドのどのセルにあるのかを推論するタスクです。また、追跡したいボールがコップの下に隠れてしまうケースも存在しており、そういった時は最後のフレームに物体が映っていないため、物体ごとの移動を追う能力が必要となります。

与えられた動画を１秒間に６枚の画像をサンプリングした後、ResNetで特徴量を抽出します。そして、系列データから予測するタスクとしてLSTMやTransformerとベースラインとして比較実験を行いました。多めのMechanismモジュール（n_s = 8）を用いることで大幅の精度向上を実現できたことがわかります。

まとめ

この論文ではTransformerに対して、全ての情報を一つ大きな潜在変数で扱うことのデメリットを指摘した上で、独立メカニズムの概念を取り入れたTIMを提案しました。そして、画像、音声、言語処理や推論といった幅広いタスクでTIMの有効性を示しました。

Transformerが非常に多くの領域で活躍しているこそ、少しの精度向上で大きなインパクトを生み出す可能性を秘めています。特に「世界が独立したメカニズムによって動かされている」という仮説が本当に正しいのであれば、TIMがより多くのシチュエーションと領域に貢献するに違いないという思いを込めて本論文を紹介させてもらいました。