Not All Memories are Created Equal: Learning to Forget by Expiringを読んだメモ

概要

　Transformerに与える系列に対して、記憶が有効である期限を事前に計算し、不要な記憶は捨てていくような機構を導入する。

提案手法

　前提として、毎回のタイムステップで、表現ベクトル $h _ t$ がやってくる。これをTransformerで処理して、過去の情報を考慮しつつOutputを出力するような状況を考える。

　普通にやると $h _ 1, \dots, h _ t$ をTransformerに全部入力する。Transformerだと基本的に $O(N ^ 2)$ かかるので、系列が長くなると計算時間が大変なことになる。

　提案手法では以下のようにする。

f:id:tokumini:20211125124933p:plain

　あるタイムステップで表現ベクトル $h _ t$ がやってきたときに、まずそれの期限(Span) $e _ t$ を計算する。 $h _ t$ はその期限が切れるまでの間だけTransformerで考慮されることにする。

　 $t$ は現在時刻であり、添字がかぶるので、過去のステップを $i$ とする。

　 $e _ i$ の計算方法は

$e _ i = L \sigma (w h _ i + b)$

ようするに1層の線形層に、活性化関数としてシグモイド関数をかけて、0からLの範囲になるように引き伸ばす。Lは最大の期限として設定するハイパーパラメータ。

　現在のステップ $t$ における、ステップ $i$ データの寿命を計算する

$r _ {ti} = e _ i - (t - i)$

　これが0より大きいものだけ取り出して、Transformerに入力するということをしたい。しかしそういうことを完全にやってしまうと、0以下のものに対して勾配が流れなくなってしまうので、学習中は以下のようにやや緩和して行う。

$m _ {ti} = \max (0, \min(1, 1 + r _ {ti} / R))$

　ここで $R$ はハイパーパラメータである。この式が何をしたいかというと

f:id:tokumini:20211125213830p:plain

らしい。

　メモリを削減していきたいので、学習のときは $e _ i$ にL1ノルムで制約をかける。 $\alpha$ をハイパーパラメータとして、

$L _ {total} = L _ {task} + \alpha \sum _ i e _ i / T$

という形で損失を付け加える。

学習の工夫

　計算効率とかのために少し手を加えるらしい。興味が持てなかったのでパス。

実験

　過去に一つ重要な情報があり、それをずっと保持しなければいけないようなトイプロブレム的な強化学習や系列のコピーでは全体的に良い。(もちろん提案手法の有効性を主張するためのタスク設定なので当然なのだけど、やや恣意的だなと感じるところはある)

　言語モデリングや物体衝突に関する認識的なやや実践よりのタスクだと、Transformerを用いて過去の情報を持ち越してくるような手法として、Compressive TransformerとAdaptive-Spanというものと比較している。精度とかメモリ効率とか計算効率がそれなりに良いという主張ではあるらしい。

f:id:tokumini:20211126053922p:plain

　ちゃんと有効そうなところへ長い期限が割り当てられていそうという分析は面白かった。

f:id:tokumini:20211126054306p:plain

所感

　不要な記憶は捨てたいよねというモチベーション自体はわかるが、保持すると決めた表現ベクトルが加工されないまま残ることはちょっと違和感があるかもしれない。覚えておくべき事柄も、だんだん洗練された形式で記憶を保持していくということはできないのかな。その点は比較されているCompressive Transformerの方が良さそうな印象はある。まぁでもそれを言っていくとPerceiverになっていくのかな……。