NV-Embed: Improved Techniques for Training LLMs as Generalist Embedding Models
Motivation 選んだ理由
これまでは LLM (Decoder only) のエンベディングよりもBERT や T5 などの Encoder-Decoder モデルの方が性能的には良かった
この論文で紹介されている, NV-Embed モデルは MTEB SoTAを記録し、CC-BY-NC で公開されている（おそらく訓練データ由来のライセンス）。
Summary どんなもの？
手法的には３つの改良を行っている
対照学習中にLLMの causal アテンションマスクを除去する
LLM出力（各トークンに対するエンベディング）に対して、潜在変数（学習で決定するパラメタ）をもつアテンション層を提案し、入力テキストに対するエンベディングを生成する
モデルの学習には、2段階のcontrastive instruction-tuning を行う
Contribution 先行研究と比べてどこがすごい？
サマリに書いたこれまでのエンベディング研究の集大成的なアーキテクチャ
一般公開されている訓練データのみを使って、今回の性能を達成している
これまでの LLM エンベディングモデルは、GPT-4 で生成された大量の人工データでFT を行っていた
https://gyazo.com/cbc479b9df355a5141b480b49229baa0
Classification のスコアが特に高い
SimCSE や E5 (mE5) が Encoder-Decoder モデルの手法
E5 は BERT を蒸留した MiniLM を元に訓練
SimCSEは BERT/RoBERTa がベースになっている
voyage, text-embedding-3-large, Cohere がプロプライエタリなモデル
Method 技術や手法のキモはどこ？
対照学習中にLLMの causal アテンションマスクを除去する
https://gyazo.com/deb7f6c93e94b34f303ca337d5b8906c
Decoder only モデルは、過去のトークンのみアテンションできるようにマスクされている
これが、Decoder モデルが Enc-Dec モデルよりもエンベディングの表現力を下げる原因になっていた
エンベディングモデルとしての学習時にこの制限を外すことで、表現力を上げている
アテンションマスクを除いてしまえば、Decoder only とEnc-Dec でモデルはほぼ同じ
（GRITではcausal アテンション と双方向アテンション両方のハイブリッドを採用しているが、単純に制限を外すだけでも良いとわかったのが貢献）
LLM出力に対して、潜在変数をもつアテンション層を提案し、入力テキストに対するエンベディングを生成する
https://gyazo.com/dfa0fc0c3d8a5d5031cb154dfb525459
平均プーリングや<EOS>トークンのエンベディングを使用する場合と比較して、一貫して検索と下流タスクの精度を向上させた
モデルの学習に、2段階のcontrastive instruction-tuning を行う
エンベディングにおける instruction-tuning
指示＋クエリとドキュメントを与えたうえで、クエリのエンベディングを生成する
https://gyazo.com/57957edcbfbc525590f2ce6f094953ab
訓練・評価時ともにinstruction に関する指示トークンはマスクされている
ドキュメント側のエンベディングには指示を与えない
指示の例
https://gyazo.com/55ed7a3a27f2b778b2e8c0134ad22802
検索タスクと非検索タスク（分類・クラスタリングなど）を適切に実行できるエンベディングモデル構築のため
検索タスクと非検索タスクの両立は難しい
例えば、in-batch negative は検索向けの学習が効率的である一方、分類などでは誤った学習をする可能性がある
クエリに対して、同一バッチ内のほかの回答を負例として利用する
１段目では、検索データセットに対して、in-batch 負例 とアノテートされたされたハード負例を利用
検索タスクの方が学習が難しいため、こちらに比重を置いている
２段目で、検索データセットと非検索データセットの組み合わせに対して、指示による対照学習を行う
Experiments どうやって有効だと検証した？
https://gyazo.com/877171e8cd455508ce02424c4eb56352
MTEB のリーダーボードでトップをとりました
https://gyazo.com/e44f7f39eca14f62876b3d1f08ad4c32
二段階のinstruction-tuning の効果
２段目の学習で、検索タスクも少し上がっているが、特にクラスタリングが大きく上がっている
Latent-attension はEOS、Mean、普通のアテンション層追加とくらべて安定して高い性能を出している
セルフアテンションを1層追加した所で大勢には影響なく、ほとんどEOSと変わらない
と書いてあるが、二段目ではちょっとEOSにくらべて上回っている例(クラスタリング、ペア分類、分類、STS)が多いような気がする
bidirect vs. causal
おおむね一貫してbidirect が高い（たまにcausalが上回るところもあるが、差は大きくない）
Discussion 議論はある？
ひとまずの完成形、今後のベースラインができた感がある
一般に入手可能なデータだけで学習していて、かつそれほど学習自体は重くはなさそうなので、手を出しやすいはず
モデル自体も重くないので、動かして使ってみる価値は高そう
一方で、エンベディングのための指示はバリエーションが多くはなく、変わった用途に対応できるか？という所は気になる所