これは、フィードバックメモリを用いたシーケンシャル・トランスフォーマーの高位表現へのアクセスに関する論文「PyTorch」のPyTorch実装です。
通常のトランスフォーマーはトークンを並行して処理します。各トランス層は、前の層の出力に注目します。フィードバックトランスは、前のステップのすべてのレイヤーの出力に注目します。そのため、これによって繰り返しが発生し、トークンごとに処理する必要があります。これにより、トレーニングが大幅に遅くなります(シーケンスの長さにもよりますが、約5倍から10倍です)。ただし、Feedback Transformerを予測する場合、メモリベクトルをキャッシュすれば次のトークンを予測できるため、より高速です
。トレーニングをスピードアップするために、この論文では短いシーケンス長から始めて、徐々に長くする方法について説明します。また、事前学習済みの並列変圧器を出発点として使用する方法についても説明します
。オリジナルのフィードバックトランスは、すべてのレイヤーの出力を保持するわけではありません。代わりに、すべてのレイヤーの出力の加重合計が保持されます。これにより、予測中のキャッシュに使用されるメモリが減ります。このファイルの前半はこれを実装しています。
更新されたフィードバックトランスフォーマーは、キーと値の計算に使用される重みをレイヤー間で共有します。その後、各ステップのキーと値を一度だけ計算し、キャッシュに保存します。このファイルの後半はこれを実装しています。パフォーマンスを向上させるために、カスタム PyTorch 関数を実装しました
。Tiny Shakespeareデータセットのフィードバックトランスフォーマーをトレーニングするためのトレーニングコードとノートブックは次のとおりです。