スペシャリストたちが協調する仕組み：Shared Global Workspace

Transformer 2021年08月16日

3つの要点
✔️ 認知科学のGlobal Workspaceから生まれたアイデアをTransformerやRIMsに導入した
✔️ 競争機構とブロードキャスト機構と持つShared Workspaceを提案する
✔️ Shared Workspaceによる全スペシャリスト間の情報共有の有効性を難易度の高い実験で示した

Coordination Among Neural Modules Through a Shared Global Workspace
written by Chengyue Gong, Dilin Wang, Meng Li, Vikas Chandra, Qiang Liu
(Submitted on 1 Mar 2021)
Comments: Published on arxiv.
Subjects: Machine Learning (cs.LG); Artificial Intelligence (cs.AI); Machine Learning (stat.ML)

code：

本記事で使用している画像は論文中のもの、またはそれを参考に作成したものを使用しております。

はじめに

異なる情報を明確に分離するような構造化したモデルが深層学習のトレンドになります。実は1980年代のAI研究は知恵を生み出せるアーキテクチャを如何にデザインするかにフォーカスしていました。中でも異なる役割のモジュールをうまく連携すれば非常に複雑なシステムを作り出せるのではないかという有力なアイデアがありました。

モジュール間の因果関係やお互いに及ぼす影響が重要になるケースではモジュール全体の同期が求められ、問題として残っていました。著者は認知科学のGlobal Workspace理論からヒントを得て、全てのモジュールが同時にアクセスできる共有表現をモジュール構造に取り入れることを提案します。この共有表現はいずれのスペシャリストモジュールによる影響を受け、全てのモジュールに情報をブロードキャストする機能を持ちます。このようなモジュール間の情報を協調する共有表現の構造はTransformersやRIMsに似ているため、これらのアーキテクチャをベースに拡張を行います。

TransformersやRIMsに用いられるAttention機構はポジションに対してペアワイズな相互作用を行います。つまり、全てのポジションに対して２つポジションずつAttentionによる計算を行います。このペアワイズな相互作用は全体のポジション間の情報をうまくシェアすることができないことを論文で指摘した上、モデルの全ての部分（モジュール）がお互いに情報シェアできる仕組みが必要と主張します。

提案手法Shared Workspaceを一言で説明すると、Shared Workspaceに各スペシャリストモジュールが入力に最も関連する場合のみ書き込みを行い、Shared Workspaceから自動的に全てのスペシャリストモジュールに情報を、ブロードキャストするものになります。ここでわかりやすさのために、以下の説明においてスペシャリストモジュールを単にスペシャリストと略します。また、誤解を避けるために提案手法にまつわる専門用語（例えば：Shared Workspace）は訳せずにそのまま用います。

Shared Workspaceによるモジュール構造間の同期

認知科学のGlobal workspace architectureを再現するように、著者たちはShared Working Memoryを通じてスペシャリストたちが疎なコミュニケーションを行うようなアーキテクチャを設計しました。具体的には、TransformerとRIMs（AttentionとSlot-basedのモジュール構造のアーキテクチャ）に対して、Shared Workspaceを加えた上、モジュールに書き込みする権限を競争させる仕組みに拡張しました。このShared Workspace構造によってスペシャリスト同士の同期や調整をうまく行えることに期待しています。

TransformersもRIMも、スペシャリスト間で情報を共有するためにペアワイズ方式のSelf-Attentionメカニズムを利用しているに対して、提案手法は限られた容量のShared Workspaceを通じてスペシャリスト間の情報共有を促進します。各計算ステージにおいて、異なるスペシャリストがShared Workspaceへの書き込み権限を競うステップ１と、Shared Workspaceの内容をすべてのスペシャリストに同時にブロードキャストされるステップ２から構成されています。

具体例をFigure 2に示しています。Transformer(b)とUniversal Transformer(d)にShared Workspace層を追加し、モジュール構造のRIMs(a)とTIMs(c)にモジュール同士のコミュニケーション層をShared Workspace層に置き換える形で簡単な拡張を行います。また、Shared Workspaceに書き込む操作とShared Workspaceの情報をブロードキャストする操作をAttention機構で実現します。* RIMsとTIMsの詳細についてこちらの記事を参考していただければ幸いです。ではShared Workspaceの詳細を3ステップに分けて説明をしていきます。

入力情報からスペシャリストの表現を取得する

ステップ１では、入力情報から各スペシャリストの表現を取得します。このステップではRIMsとTransformersに含まれるスペシャリストの異なる入力に合わせた表現を整えます。

ここで、スペシャリストが各計算ステージに次に紹介する２つのステップの準備を行います。1. 各スペシャリストがShared Workspaceへの書き込みを競います。2. 各スペシャリストがShared Workspaceからブロードキャストされた情報を受け取ります。

Shared Workspaceに情報を書き込む

ステップ２ではスペシャリストたちが、受け取った新しい情報に対して自身を更新すべき際に、Shared Workspaceへの書き込み権限を競います。つまり、入力の情報に関連するスペシャリストがKey-Value Attentionで算出した関連度スコアを高くなるように学習します。なんとなく入力の情報に対して反応し、頻繁に更新するスペシャリストは真に重要な入力が受け取った際にその他のスペシャリストに勝てず、更新できなくなってしまう危険があります。こうした競争させる仕組みによって、異なる入力に対して担当するスペシャリストが反応する分業化ができることに期待できます。

具体的には、Key-Value AttentionのKey, ValueはStep1で得られたスペシャリストの表現を線形変換したもの、QueryはShared Workspaceの表現を変換したものが該当します。つまり、Shared Workspaceが関連したスペシャリストにアクセスする仕組みになります。そしてAttentionのSoftmax計算によって得られるスコアを順に並び、最も関連するk個のスペシャリストがShared Workspaceに書き込みます(式1)。このtop-kの操作は、スタンダートのsoft-attention(すべてのスペシャリスト)とhard-attention(top-1のスペシャリスト)のバランスを取ったものと考えられます。

　(1)

Shared Workspaceから情報のブロードキャスト

ステップ３では、Shared Workspaceの情報を全てのスペシャリストにブロードキャストします。ここでもAttention機構で算出した関連度スコアを元に、スペシャリストの更新度合いを決めます。ただし、ステップ２とは逆にスペシャリストからQueryを生成し、Shared WorkspaceのKey, Valueを生成してSoft-Attentionを計算します。全てのスペシャリストが得られたShared Workspaceの情報を元に各自の潜在表現を更新します（式2）。ここで更新するのは、RIMsの場合はLSTMやGRUで、Transformerの場合はFFN順伝搬レイヤーとなります。

　（2）

ただし、hはLSTMやGRUの潜在表現、Sは関連度スコア、vはShared Workspaceの表現を線形変換したものになります。

Shared Workspaceの一貫性と計算量

一貫性については、ステップ毎にShared Workspaceは更新を行いますが、エピソードが終了するタイミングでリセットします。つまり、RIMsの場合は入力系列が終了する（ゲームが終了する）タイミング、Transformerは最終レイヤーまでの伝搬が終えるタイミングまでに、同一のShared Workspace表現を共有します。

TransformerやRIMsはAttention機構を用いた二つごとのスペシャリスト間の関連度を計算するためO(n^2)の計算量になるに対して、提案するShared Workspace構造はスペシャリストの個数nに対してO(n)の計算量になります。実際には、スペリャシルトの個数nは決まっており、計算量の観点ではShared Workspaceを用いた構造は非常に優れており、大規模な実験に適していると考えられます。（参考までに、関連文献によると人間のワーキングメモリは10以下で、非常に小さいと思われます。）

実験

実験では二つのことを示します。(a) 幅広いベンチマークを用いてShared Workspaceが精度を向上させることを確認し、提案手法の実用性と汎用性を示します。(b) ペアワイズの相互作用を用いずに精度向上を実現できることを確認し、Shared Workspaceがスペシャリスト間の一貫性を維持できることを示します。詳細な実験設定は論文の付録に詳しく示されており、参考にしてください。

画像入力に関する理解を問うタスク

スペシャリストが下流タスクに有用な情報のみを限られたShared Workspaceに書き込むようになると考えられます。このアイデアの検証は、複数の視覚情報を処理するタスクと次のベースラインを用いて行います。

TR（Transformers）：層ごとのパラメータが共有されるTransformer
STR（Sparse Transformers）：スパースなAttention Matrixを用いたTransformer
TR+HC（High Capacity Transformers）：層ごとのパラメータが異なるTransformers
TR+SSW（Transformers with Shared Workspace with soft-competition）：soft-attentionを用いたShared Workspaceを取り入れたTransformers
TR+HSW（Transformers with Shared Workspace with top-k competition）：top-k attentionを用いたShared Workspaceを取り入れたTransformers

・Detecting Equilateral Triangles：正三角形検出タスク

64*64の画像に現れる点群が正三角形になっているかどうかを二値分類するタスクで、ベースラインTRをVision Transformer（ViT）とします。画像を4*4のパッチに分けて画像の系列を入力します。特定の情報に注目すれば解けるタスクなので、キャパシティが制限されたShared Workspaceが重要な情報にのみ注目でき、精度がよくなる仮説が立てられます。結果（Figure 3）はこの仮説を支えており、TR+HSWがBaselineのTRより分散が小さく精度が良くなりました。

・CATER：物体トラッキングタスク

Caterは動画が与えられて目標オブジェクトが動画終了時に6*6のグリッドのどのセルにあるかを当てるタスクになります。最終フレームに目標オブジェクトが隠れていない場合は簡単に解けるが、障害物によって隠れてしまう場合も存在するため、見えない物体でも長期的かつ正確にトラッキングする推理能力が求められます。実際は1秒間に6フレームをサンプリングした画像の系列をTransformerに入力する形になります。Table 1より、提案手法TR+HSWとTR+SSWがベースラインよりわずかに良いことがわかります。

*この実験でははっきりした精度向上が見られないが、6*6の36クラス分類でタスク難易度が高い割にベースラインと同等以上の精度は達したと筆者は考えています。

・Sort-of-CLEVER：関係性推理タスク

物体間の関係性が聞かれる質問文と物体を映した画像が与えられる推論タスクになります。質問を正確に答えるには、聞かれた特定の物体に注目する必要があります。75*75の画像に6種類の色と2種類の形状でランダムに配置される物体が存在します。1枚の画像に対して、「赤い物体の形状はなんでしょうか？」という非関係性問題と、「赤い物体に最も近い物体の形状はなんでしょうか？」という関係性問題を10問ずつ用意されています。タスクはViTのようにパッチに分割した画像を入力とし、Sort-of-CLEVERの答えが有限であるため分類タスクとして扱います。

関係性問題と非関係性問題のどちらもShared Workspaceを取り入れた提案手法の方が、収束は速く、精度が良いことが結果（Figure 4）からわかります。従って、このような離散的な情報を含むタスクにおいてはShared Workspaceが従来のTransformersアーキテクチャより優れていると考えられます。

Physical Reasoning task：物理プロセスの推理タスク

物理プロセスの推論タスクで、バウンディングボールの動きを予測するには、各ボールの動きを別々に捉える必要があります。タスクは最初の10フレームを与えて、30フレーム目と45フレーム目の予測ロスによって評価します。LSTM、RIMsとRMCをベースラインとして、RIMs＋Shared Workspaceによる精度の向上を確認します。複数の異なる条件で実験を行ったが、いずれの結果もRIMsにShared Workspaceをつけたことで精度が向上し、ほとんどの条件においてRMCよりも優れることがわかります。

Shared Workspace for Multiagent Starcraft World Modelling.

最後に、Starcraftというマルチエージェントなゲーム環境（SC2 domain）で実験しました。各エージェントが攻撃力、防衛力やHP値などの状態指標だけでなく、複雑なスキルや特徴を有しているため大変難しい環境となっています。ただし、こういった離散的な属性を持ちながらこれらが相互作用するゲームの性質は、モジュール構造RIMs+Shared Workspaceのアーキテクチャに適しており、提案手法の有効性が確かめると考えられます。実験の詳細は付録Gを参考にしてください。

Table 2におけるRIMsの精度がよくない理由として、ペアワイズな相互作用を行うRIMsはそれぞれ2個のスペシャリスト間でコミュニケーションを行うため、3種類以上の情報を同時に考慮したいこのタスクには適していないと考えられます。結果としてShared Workspaceを用いた提案手法がLSTMやRIMsより精度がよく、異なるスペシャリスト間の情報一貫性を保つアイデアが有効であることが示されている形になります。

まとめ

この論文は、認知科学のGlobal Workspace理論からヒントを得て、RIMsやTransformersにShared Workspaceを追加して拡張することで、全てのスペシャリストモジュール間の情報一貫性を保つことを検証しました。多くの実験を通じてペアワイズな相互作用を行うベースラインよりShared Workspaceを用いた全てのモジュール間で協調を行うことが有効であることを示しました。

個人的な見解にはなりますが、提案手法は目立つ実験結果を出したわけでないものの、因果推論の独立メカニズムや認知科学のGlobal Workspaceといった重要な概念を深層学習に取り入れようとしています。これは、いわゆる強いAIを実現するには既存の深層学習アーキテクチャでは難しいため、新たなアーキテクチャを探索することが背後にあるモチベーションではないかと考えています。現在の大規模モデルとビッグデータによるAIが頭打ちになった際に、こういったアーキテクチャの研究が重要になってくるに違いないので、さらに深掘りする価値はあるのではないかと思っています。