[論文レビュー] Why do Nearest Neighbor Language Models Work?

Frank F. Xu, Uri Alon+1 more|arXiv (Cornell University)|Jan 7, 2023

Topic Modeling被引用数 6

ひとこと要約

この論文は、k-Nearest Neighbor Language Models (kNN-LMs) が、ベース言語モデルと同じ学習データから検索を行っても、なぜ perplexity が向上するのかを調査する。主な要因として、異なるトランスフォーマー層からの表現をアンサンブルすること（55% の向上）、近似最近傍探索を使用すること（2.6% の向上）、およびソフトマックス温度のチューニング（最大3.7% の向上）を特定し、これらが kNN-LM の成功を記憶や過学習を超えて説明している。

ABSTRACT

Language models (LMs) compute the probability of a text by sequentially computing a representation of an already-seen context and using this representation to predict the next word. Currently, most LMs calculate these representations through a neural network consuming the immediate previous context. However recently, retrieval-augmented LMs have shown to improve over standard neural LMs, by accessing information retrieved from a large datastore, in addition to their standard, parametric, next-word prediction. In this paper, we set out to understand why retrieval-augmented language models, and specifically why k-nearest neighbor language models (kNN-LMs) perform better than standard parametric LMs, even when the k-nearest neighbor component retrieves examples from the same training set that the LM was originally trained on. To this end, we perform a careful analysis of the various dimensions over which kNN-LM diverges from standard LMs, and investigate these dimensions one by one. Empirically, we identify three main reasons why kNN-LM performs better than standard LMs: using a different input representation for predicting the next tokens, approximate kNN search, and the importance of softmax temperature for the kNN distribution. Further, we incorporate these insights into the model architecture or the training procedure of the standard parametric LM, improving its results without the need for an explicit retrieval component. The code is available at https://github.com/frankxu2004/knnlm-why.

研究の動機と目的

• ベース言語モデルの学習に使用された同じ学習データから検索を行っても、kNN-LMs がなぜ性能向上を達成するのかを理解すること。
• kNN-LMs に寄与する特定の設計選択を分離・評価すること。
• 性能向上が検索、アーキテクチャ的アンサンブル、あるいは正則化やラベルスムージングなどの他の要因に起因するのかを特定すること。
• kNN-LM の設計要因を体系的にアブレーション可能にする一般化された定式化を構築すること。
• 明示的な検索を伴わない標準的なパrametric LMs において、性能向上を再現可能であることを示し、効率を向上させること。

提案手法

• 非パラメトリック部とパラメトリックLMのソフトマックスおよび埋め込み層を統合するkNN-LMの一般化された定式化を提案する。
• 系統的な変更を通じてアブレーションスタディを実施：類似度計算のための入力表現、近似kNN探索と正確kNN探索の使用、kNN部におけるソフトマックス温度。
• 異なる入力表現（例：アテンション層出力 vs. フィードフォワード層出力）を用いてベースLMの変種を学習・比較し、アンサンブル効果を分離する。
• kNN検索を微分可能行列演算に置き換えることで、推論効率を向上させつつ性能を維持する。
• 代替の学習目的およびラベルタイプ（例：ソフトラベル）をテストし、kNN-LMの成功の他の説明要因を除外する。

実験結果

リサーチクエスチョン

• RQ1 kNN-LM は、ベース言語モデルの学習に使用された同じ学習データから検索を行っても、なぜ perplexity を低減するのか？
• RQ2 性能向上は検索、アーキテクチャ的アンサンブル、それとも近似探索による正則化に起因するのか？
• RQ3 kNN部におけるソフトマックス温度はモデル性能に顕著な影響を及ぼすのか？また、元のkNN-LM設定では最適であるのか？
• RQ4 kNN-LM の性能向上は、入力表現や学習手順を変更することで、検索を伴わない標準的なパラメトリックLMでも再現可能か？
• RQ5 近似最近傍探索は一般化にどのような役割を果たし、なぜ正確探索よりも優れているのか？

主な発見

• トランスフォーマーの異なる層（例：アテンション出力とフィードフォワード出力）からの表現をアンサンブルすることで、kNN-LM の perplexity 向上に 55% の寄与が見られ、これは perplexity に対して 6.5% の相対的低減に相当する。
• 正確な検索ではなく近似最近傍探索を使用することで、perplexity が 2.6% 相対的に向上し、精度の低下による正則化効果が示唆される。
• kNN部におけるソフトマックス温度のチューニングにより、最大 3.7% の相対的 perplexity 向上が得られ、温度が重要ながしばしば見過ごされるハイパーパramータであることが示された。
• kNN-LM の成功は、異なるベースモデルとのアンサンブルとは直交しており、性能向上がモデルアンサンブルによるものではないことを示している。
• 記憶、過学習、スパース化、盗み取り確率、ソフトラベル化といった要因は、制御されたアブレーション実験により除外されたため、性能向上はそれらに起因しない。
• kNN検索を微分可能な行列演算に置き換えることで、6.5% 以上の相対的 perplexity 向上を維持しつつ、検索を伴わない効率的な推論が可能になる。