[論文レビュー] Neural Execution Engines: Learning to Execute Subroutines

Yujun Yan, Kevin Swersky+3 more|arXiv (Cornell University)|Jun 15, 2020

Topic Modeling参考文献 6被引用数 7

ひとこと要約

この論文では、強力な一般化を達成するTransformerベースのアーキテクチャであるニューラル実行エンジン（NEE）を提案する。NEEは、ソートや最短経路といったアルゴリズム的サブルーチンを学習し、訓練分布をはるかに超える入力に対しても高い精度を達成する。学習された条件付きマスキング機構と2進数表現を導入することで、訓練時とは異なる長さのシーケンスや未観測の値域を含む入力に対しても、ほぼ完璧な精度を達成する。

ABSTRACT

A significant effort has been made to train neural networks that replicate algorithmic reasoning, but they often fail to learn the abstract concepts underlying these algorithms. This is evidenced by their inability to generalize to data distributions that are outside of their restricted training sets, namely larger inputs and unseen data. We study these generalization issues at the level of numerical subroutines that comprise common algorithms like sorting, shortest paths, and minimum spanning trees. First, we observe that transformer-based sequence-to-sequence models can learn subroutines like sorting a list of numbers, but their performance rapidly degrades as the length of lists grows beyond those found in the training set. We demonstrate that this is due to attention weights that lose fidelity with longer sequences, particularly when the input numbers are numerically similar. To address the issue, we propose a learned conditional masking mechanism, which enables the model to strongly generalize far outside of its training range with near-perfect accuracy on a variety of algorithms. Second, to generalize to unseen data, we show that encoding numbers with a binary representation leads to embeddings with rich structure once trained on downstream tasks like addition or multiplication. This allows the embedding to handle missing data by faithfully interpolating numbers not seen during training.

研究の動機と目的

• アルゴリズム的推論におけるニューラルネットワークの強力な一般化の欠如、特に訓練分布の範囲を超える入力に対しての一般化の欠如を解決する。
• 数値的に類似した値を含む長大なシーケンスにおいて、Transformerの注目メカニズムの劣化が生じ、誤予測と誤差の蓄積を引き起こす問題を克服する。
• 2進表現を用いることで、構造的で補間可能な数値埋め込みを学習し、未観測のデータ分布への一般化を可能にする。
• 1つの訓練済みサブルーチン（例：比較）を、再訓練なしにDijkstra法やPrim法など多様なアルゴリズムに再利用可能なモジュラーなフレームワークを設計する。
• 異なるグラフタイプや入力分布に対して、分布シフト下でもパフォーマンスの低下が最小限に抑えられることを示す、耐久性を確保する。

提案手法

• 各ステップで値とポインタを予測するエンコーダ・デコーダ型のsequence-to-sequence Transformerを用い、動的状態追跡を可能にする。
• 予測されたポインタに基づいてエンコーダーマスクを更新する学習された条件付きマスキング機構を導入し、長大なシーケンスにおける注目メカニズムの忠実性を向上させる。
• 数値を2進形式で表現することで、階層的で指数的スケーラビリティを持つ埋め込みを実現し、未観測の値への補間を可能にする。
• 加法や乗法といった下流タスクで訓練することで、構造的で解釈可能な数値表現を学習する。
• アルゴリズムの論理の学習をデータ分布から分離し、完全なアルゴリズムではなくサブルーチンで訓練することで、モジュラーな再利用を可能にする。
• 同じ学習済みサブルーチン（例：比較）を、Dijkstra法やPrim法など複数のアルゴリズムに同じNeeを適用することで、再訓練なしに転送可能であることを示す。

実験結果

リサーチクエスチョン

• RQ1 Transformerベースのモデルは、訓練分布をはるかに超える範囲で、ソートや最短経路といったサブルーチンを強力に一般化して実行できるか？
• RQ2 標準的なTransformerはなぜ、長大なシーケンスや数値的に類似した入力に対して一般化できないのか？ また、注目メカニズムの忠実性をどのように維持できるか？
• RQ3 2進数表現によって、構造的で補間可能な埋め込みを実現でき、未観測の数値への一般化を可能にするか？
• RQ4 1つの学習済みサブルーチン（例：比較）を、再訓練なしに複数の異なるアルゴリズムにどの程度再利用できるか？
• RQ5 入力データの分布シフト（例えば、異なるグラフタイプや入力範囲）に対して、ニューラル実行エンジンはどの程度耐性を示すか？

主な発見

• Dijkstra法の最短経路とPrim法の最小全域木の両方において、テストグラフサイズが100までに達しても100.00%の精度を達成した。訓練はサイズ8のグラフで行った。
• 訓練時とは異なるグラフタイプのテストグラフに対しても、グラフアルゴリズムのパフォーマンスはほぼ完璧（最短経路99.91%、MST 93.00%）を維持し、分布シフトに対する耐久性を示した。
• 2進数表現を用いることで、構造的で解釈可能な埋め込みを学習し、補間が可能となった。例えば、訓練時に除外された65%の数値が忠実に補間された。
• 学習された条件付きマスキング機構により、強力な一般化が実現された。長さ8の短いシーケンスで訓練したモデルが、長さ100のシーケンスに対しても正確に一般化した。
• 1つのサブルーチン（例：比較）で訓練したNEEは、再訓練なしにDijkstra法やPrim法など複数のアルゴリズムに直接適用可能であり、高いモularityと転送性を示した。
• 2進埋め込みを用いたモデルは、パラメータ効率性とスケーラビリティにおいてワンホット出力よりも優れており、すべてのタスクおよび入力範囲で高いパフォーマンスを維持した。