K-最近傍言語モデル

これは、「記憶による一般化：最近傍言語モデル」という論文をPyTorchで実装したものです。自己回帰変換モデルの複雑さを改善するために、k-最近傍を使用します

自己回帰言語モデルでは$p(w_t|c_t)$、$w_t$$t$ステップ中のトークンがどこで、$c_t$がコンテキストなのかを推定します。$c_t=(w_1,w_2,...,w_{t-1})$

この論文では、$(f(c_i),w_i)$検索キー付きのキーと値のペアでの k $p(w_t|c_t)$ 最近傍検索の使用方法を改良しました。$f(c_t)$$f(c_t)$$c_t$これがコンテキストの埋め込みです。この論文（およびこの実装）では、変圧器の最終層のフィードフォワード層への入力を次のように使用しています

$f(c_i)$インデックスにはFAISSを使用しています。

実装

したがって、$k$ NN-LM を実行するには、以下を行う必要があります。

• トランスフォーマーモデルのトレーニング
• のインデックスを作成 $(f(c_i),w_i)$
• 以下の NN 検索を使用して $k$knn-ML を評価する $(f(c_i),w_i)$ $f(c_t)$

この実験では小さなデータセットを使用しているため、インデックス用に数百ギガバイトのディスク容量を消費せずに実行できます。

$k$NN-LM の正式な実装はここにあります。