[論文レビュー] Discrete-Valued Neural Communication

Dianbo Liu, Alex Lamb+5 more|arXiv (Cornell University)|Jul 6, 2021

Advanced Graph Neural Networks参考文献 49被引用数 13

ひとこと要約

本稿では、トランスフォーマー、グラフニューラルネットワーク、モジュラーネットワークなどの構造的ディープラーニングモデルにおける部品間通信を、共有で学習可能なコードブックを通じて離散的符号に制約することで、離散的記号を生成する、離散的価値を持つニューラル通信（DVNC）を提案する。通信を離散値に制限することにより、分布外（OOD）タスクにおける体系的一般化性能が向上し、ノイズ耐性が向上し、モデル次元数が低減される。この手法は、アーキテクチャの変更なしに、多様なアーキテクチャで一貫した性能向上を達成する。

ABSTRACT

Deep learning has advanced from fully connected architectures to structured models organized into components, e.g., the transformer composed of positional elements, modular architectures divided into slots, and graph neural nets made up of nodes. In structured models, an interesting question is how to conduct dynamic and possibly sparse communication among the separate components. Here, we explore the hypothesis that restricting the transmitted information among components to discrete representations is a beneficial bottleneck. The motivating intuition is human language in which communication occurs through discrete symbols. Even though individuals have different understandings of what a "cat" is based on their specific experiences, the shared discrete token makes it possible for communication among individuals to be unimpeded by individual differences in internal representation. To discretize the values of concepts dynamically communicated among specialist components, we extend the quantization mechanism from the Vector-Quantized Variational Autoencoder to multi-headed discretization with shared codebooks and use it for discrete-valued neural communication (DVNC). Our experiments show that DVNC substantially improves systematic generalization in a variety of architectures -- transformers, modular architectures, and graph neural networks. We also show that the DVNC is robust to the choice of hyperparameters, making the method very useful in practice. Moreover, we establish a theoretical justification of our discretization process, proving that it has the ability to increase noise robustness and reduce the underlying dimensionality of the model.

研究の動機と目的

• 構造的ニューラルアーキテクチャにおける部品間の動的かつスパースな通信の課題に対処すること。
• 部品間通信に再利用可能な離散的記号を導入することで、分布外（OOD）設定における体系的一般化を向上させること。
• モデル次元数を低減し、ノイズ耐性を向上させる理論的裏付けのある、頑健な部品間通信手法を提供すること。
• 専門的部品間で離散的で共有される表現が、人間の言語や脳機能にインspiredされた共通言語として機能できるかを検討すること。

提案手法

• DVNCは、ベクトル量子化変分オートエンコーダー（VQ-VAE）を拡張し、複数のヘッドにわたる離散化と、部品間で共有されるコードブックを導入する。
• 訓練データ上でK-meansクラスタリングを用いて、サイズ L × m の共有コードブックを学習し、連続的活性化を最も近い離散的コードブックベクトルにマッピングする。
• 勾配の直線的推定（straight-through gradient estimator）を用いて、離散的コードブックを介してタスク損失をバックプロパゲートするが、コードブック自体は再構成損失およびコミットメント損失に基づく勾配降下法で更新される。
• コミットメント損失により、活性化が割り当てられたコードブックベクトルに近づくように保証され、タスク損失は送信部品にのみバックプロパゲートされる。
• 部品間の通信は離散トークンインデックスに置き換えられ、モデルアーキテクチャは保持されるが、記号的通信に類似した制約が課される。
• コードブックは、残りのモデルと同時にエンドツーエンドで訓練され、指数的移動平均の更新を用いないため、直接最適化が可能になる。

実験結果

リサーチクエスチョン

• RQ1 部品間通信を離散化することで、構造的ニューラルネットワークにおける体系的一般化性能が向上するか？
• RQ2 共有で学習可能なコードブックに基づく通信が、分布シフトや分布外一般化に対する耐性を向上させるか？
• RQ3 離散的通信は、特に被覆数の対数の観点から、モデルの感度および基本次元数にどのように影響するか？
• RQ4 DVNCの性能はハイパーパrameterの選択に頑健であるか、特にコードブックサイズと損失重み付けに影響を受けるか？
• RQ5 離散的通信が、専門的部品間で共通言語として機能できるか、モジュラリティと再利用性を向上させるか？

主な発見

• DVNCは、GNN、トランスフォーマー、RIMsなど複数のアーキテクチャにおいて、分布外一般化性能を顕著に向上させ、OODベンチマークで一貫した向上を示した。
• Sort-of-Clevrタスクでは、OOD設定下でベースラインモデル比で平均逆順位（MRR）が15%向上した。
• 2D ShapesおよびThree-bodyタスクでは、非離散化ベースラインと比較して、DVNCモデルがOOD一般化性能で20〜30%高い正答率を示した。
• ハイパーパrameterの選択に対して頑健であり、異なるコードブックサイズや損失重みでも安定した性能を示した。
• 理論的分析により、DVNCはモデルの感度を向上させるとともに、被覆数の対数を低減することが確認され、低い基本次元数を示していることが示された。
• トレーニング時間はスケーラブルであり、GNNは1GPUで1タスクあたり約3時間、RIMs/トランスフォーマーは約12時間で学習が完了し、全実験で合計約800 GPU時間の計算リソースを要した。