1 of 15

Thinking Like Transformers

By Gail Weiss, Yoav Goldberg, and Eran Yahav, [arXiv:2106.06981]

ICML2021論文読み会 2021-08-18

中村道宏 (@__nakamichi__)

1

2 of 15

概要

  • RNNには有限オートマトンという対応する計算モデルがある
  • Transformerに対応する計算モデルはあるだろうか?�→作ってみた!

  • RASPというTransformerを抽象化したプログラミング言語を提案
    • Restricted Access Sequence Processing Language
  • RASPを通じて、Transformerの性質や内部の動作を理解できる
    • ある種のタスクを解くのに必要なネットワークの層数やattention headsの個数が見積もれる
    • attentionの重みがどんな計算アルゴリズムを表しているかがわかる
    • “Thinking Like Transformers”

2

3 of 15

RNNとオートマトン

  • オートマトンを通じてRNNを理解する研究は昔から色々ある
    • 学習済みRNNからオートマトンを抽出する
    • オートマトンを通じてRNNの表現力を解析する
    • etc.

3

4 of 15

Transformerと計算モデル

  • Transfomerに対応する計算モデルは?→作ろう!

4

プログラミング言語

RASP

(Restricted Access Sequence Processing Language)

※REPLで触って遊べます https://github.com/tech-srl/RASP

5 of 15

RASPによるTransformerの理解

普通のプログラミングで解けるタスク

RASPでプログラムを書く

RASPプログラムに対応するアーキテクチャのTransformerを学習させる

(RASPからTransformerへの「コンパイル」)

できたモデルはRASPプログラムと同様の性質やふるまいを持つと期待される

5

6 of 15

プログラミング言語RASP

RASP: Restricted Access Sequence Processing Language

  • データ構造:長さnの配列(sequence)とn x n 行列のみが存在
  • 入力:長さnのトークン列と、そのインデックス配列
    • tokens = “hello”
    • indices = [0, 1, 2, 3, 4]
    • Transformerの入力に相当
  • 出力:長さnの配列
  • 要素ごとに作用する各種演算子(s-ops)
    • s1 = [1, 2, 3], s2 = [4, 5, 6]
    • s1 + s2 = [5, 7, 9], s1 + 2 = [3, 4, 5] (broadcasting)
    • s1 > 2 = [F, F, T] (broadcasting)
    • pairwise_mul(s1, s2) = [4, 10, 18]
    • etc.
    • feedforward layerに相当

6

7 of 15

プログラミング言語RASP

Transformerのattention layerに相当する2つの演算を定義

  • select
    • 引数:2つの配列k, qと、2つの配列を引数にとる述語p
    • 戻り値:選択行列(selection matrix)S
    • numpy風に書くと、 S[i, j] = p(k[i], q[j])
    • attentionの行列   に相当
  • aggregate
    • 引数:選択行列Sと配列v
    • 戻り値:配列res
      • numpy風に書くと、 res[i] = S[i,:].dot(v) / S[i,:].sum()
    • attention       に相当

※forループのような制御構文はRASPには存在しない

7

8 of 15

プログラミング言語RASP

8

9 of 15

RASPで書けるプログラム

  • reverse: 文字列反転
  • histogram: 入力内の文字出現頻度をカウント
  • double histogram: 出現頻度が同じ文字種の数をカウント
  • most freq: 文字を出現頻度順にソート
  • Dick-k: k種類の括弧からなる文字列内で括弧がバランスしているかを判定
    • “([([” → バランスしていない
    • “()[]” → バランスしている
  • etc.

9

10 of 15

RASPプログラムの例:double histogram

10

11 of 15

RASPからTransformerへのコンパイル

与えられたRASPプログラムに対応するTransformerアーキテクチャの構成法

  • s-ops(要素毎の演算) → feedforward layer
  • select & aggregate → attention layer
  • 同じ配列に対して複数のselect&aggregateがある場合は、�attention headが増える

こうして構成したTransfomerは、与えられたタスクを解くのに必要な�attention layerの数とattention headの数の上界を与えると期待される

11

12 of 15

Thinking Like Transfomer

実際、RASPから構成したTransformerで前述のタスクをほぼ100%の精度で解けた

12

13 of 15

Thinking Like Transfomer

逆に、RASPによる構成よりlayerとheadの数を減らすと、精度が下がりがち�→RASPによる構成が与えるlayer数、head数の上界はけっこうタイト

13

14 of 15

Thinking Like Transfomer

RASPから構成したTranformerのattentionは選択行列と一致するか?�→だいたい一致するタスクもあれば、そうでないタスクもある

14

15 of 15

まとめと感想

  • Transformerを抽象化した計算モデルであるプログラミング言語RASPを提案
  • 単純なアルゴリズム的タスクに対しては、RASPプログラムから�Transformerを構成することで、Transformerの性質やふるまいを理解できた

  • RASPからTransformerへのコンパイルはできるが、逆がまだできないので、真に機械学習で解きたいタスク(=通常のプログラムが書けないタスク)に対するTransformerの設計や理解もできない
  • TransformerからRASP(または他の計算モデル)へのデコンパイル(?)もそのうちできるようになることを期待

15