検索ヘッドに基づく大規模言語モデルの長文脈処理の改善

大規模言語モデル（LargeLanguageModel;LLM）が文脈内学習[1]や推論時スケーリング[2,3]などの能力を発揮するには，長い文脈を処理する能力が不可欠である．こうした長文脈処理能力は，検索ヘッド（RetrievalHead）という機序と強く関連することが明らかになっている[4]．検索ヘッドはLLM内部のマルチヘッド・アテンションから特定され，文脈において関連する単語に注目し，情報を取り出す働きがあると考えられている．逆に検索ヘッドを不活性化すると，質問応答などの下流タスクにおいて性能が低下することが報告されている．検索ヘッドの発見は，LLMの長文脈処理のメカニズムを解明する手がかりとなった．しかし，この機序を活用してLLMの性能向上に成功したという報告はない．このように，LLMの解釈性と高度化の間には溝が存在する．例えば，LLMの内部に知識を蓄える機構が発見されたものの[5,6]，この機構を恣意的に操作するとLLMの汎用性能が損なわれることが示されている[7]．また，LLMの多言語能力に関する機構が特定されたが[8]，これらを制御してもLLMの言語間転移には寄与しなかった[9]．そこで，次の問いが生じる—検索ヘッドに着想を得てLLMの長文脈処理能力を向上させられるか？検索ヘッドを不活性化LLM制御後LLM…アテンションアテンションアテンションヘッド1ヘッド3ヘッド2…直接選好最適化（DPO）長文脈評価ベンチマーク学習対象LLM学習前学習後8K16K32K64K128K号を構築し，LLMの長文脈処理能力の改善を狙う．本稿では，検索ヘッドに特化した最適化を行い，言語モデルの長文脈処理能力を向上させる手法を提案する．具体的には，図1に示すように，通常のモデルと，検索ヘッドを不活性化したモデルからそれぞれ，出力をサンプリングする．通常のモデルの出力を正例，不活性化したモデルの出力を負例とし，直接選好最適化（DirectPreferenceOptimization;DPO）[10]を行う．これにより，検索ヘッドに特化した対照的な学習信号を構築し，その機能を強化することを狙う．検索ヘッドを不活性化することから，本手法をRetMask（Retrieval-headMasking）と呼ぶ．RetMaskをLlama-3.1[11]，Qwen3[12]，Olmo-3[13]に適用し，長文処理能力を評価した．その結果，Llama-3.1とQwen3では有効性が確認されたが，Olmo-3では効果が限定的であった．分析により，この差異はモデル内部の検索スコアの分布に起因し，検索機能が特定のヘッドに集中しているモデルにおいて提案手法の有効性が高まることが明らかになった．これらの知見は，検索ヘッドという機序をより強く裏付けると同時に，そのLLM開発への有用性を示すもので，本研究はLLMの解釈性が高度化に繋がることを初めて示した事例と言えよう．Thisworkispublishedwithoutpeerreviewandislicensedbytheauthor(s)underCCBY4.0(https://creativecommons.org/licenses/by/4.0/).org/licenses/by/4.0/).4実験4.1実験設定モデルLlama-3.1-8B-Instruct[11]，Qwen3-8B[12]，Olmo-3-7B-Think[13]を用いて実験を行う．検索ヘッドの認定に用いる閾値は，予備実験の結果に基づき，Llama-3.1は𝜏=0.1，Qwen3とOlmo-3は𝜏=0.05とした．ベンチマーク評価にはHELMET[15]を用いる．HELMETはLLMの長文脈処理能力の評価に特化したベンチマークであり，LLM開発でよく用いられる合成タスクRULER[16]から，検索拡張生成や長文書に対する質問応答などの実世界タスクまで，幅広くカバーしている．その他，LLM開発において一般的に使用されているベンチマーク（数学や一般教養など）で評価した結果を付録Bに示す．学習データRetMaskは指示文を含む任意のデータセットに適用できる．本稿では，LMSYS-Chat-1M[17]を学習データとして用いる．また，提案手法の効果が学習データセットに依存するものではないことを示すため，WildChat[18]を用いた実験も行い，その結果を付録Cに示す．なお，これらの学習データはベンチマークのタスクと重なりがない．ベースライン次の方法で負例𝑦𝑙を生成してDPOを適用したモデルをベースラインとする．(1)小規模モデル：自身よりもパラメータ数が少ないモデルから𝑦𝑙を生成する1）．(2)勝敗ペア：自身𝜋𝜃から応答を生成し，より品質が低いとされる応答を𝑦𝑙とする2）．(3)非RHマスク：検索ヘッド以外のヘッドを|Hret|個ランダムに不活性化したモデルから𝑦𝑙を生成する．(4)ランダムマスク：ヘッドを|Hret|個ランダムに（検索ヘッドを含んでもよい）不活性化したモデルから𝑦𝑙を生成する．ハイパーパラメータを含む他の実験設定の詳細を付録Dに示す．4.2実験結果Llama-3.1,Qwen3,Olmo-3の実験結果をそれぞれにおいて最高性能を達成した．また，勝敗ペアおよ1）Llama-3.1-8B-Instructに対してはLlama-3.1-3B-Instruct，Qwen3-8BとOlmo-3-7B-Thinkに対してはQwen3-0.7Bを用いた．2）品質の測定にはGemma-3-27B-IT[19]を用いた．Llama-3.1をHELMETで評価した結果．8K，16K，32K，64K，128Kはデータセットの入力長を示す．Llama-3.1-8B-Instruct負例生成方法8K16K32K64K128K学習前56.0354.1452.4251.6546.40小規模モデル56.7755.3253.4852.1847.53勝敗ペア56.5054.4252.4751.6246.05非RHマスク56.4555.5553.1952.1447.19ランダムマスク56.6755.9553.1452.3047.04RetMask58.1456.9253.4853.1548.68たQwen3をHELMETで評価した結果．8K，16K，32K，64K，128Kはデータセットの入力長を示す．Qwen3-8B負例生成方法8K16K32K64K128K学習前53.2050.1649.8945.4444.73小規模モデル52.5249.8148.7146.6745.51勝敗ペア52.8050.1449.7145.9344.49非RHマスク53.0250.2848.6746.7945.48ランダムマスク49.9947.0245.7543.8545.86RetMask53.7750.6150.3446.7945.62び小規模モデルのベースライン手法は，RetMaskの性能に及ばなかった．これにより，RetMaskの有効性は選好最適化そのものではなく，検索ヘッドの不活性化に起因することが示唆される．さらに，非RHマスクやランダムマスクと比較しても，RetMaskの性能が高い．このことから，任意のヘッドを不活性化するよりも，検索ヘッドを対象として不活性化する方が学習効果が高いことが確認された．ものの，RetMaskを用いた学習は概ね全ての入力長において最高性能を達成した．一方，表3においては，RetMaskの効果は限定的であった．モデルによる効果の差を分析するため，モデル内部の検索スコアの分布を調査し，5.1節で説明する．5分析5.1検索スコアの分布4.2節で観察された性能差は，検索スコアがマルチヘッド注意の中でどのように分布しているかに起因すると仮説を立てた．この仮説を検証するために，モデルにおける全アテンションヘッドの検索スコアを精査し，その分布を図3に示す．ヘッドは検索に何かしら関与しており，検索機能をorg/licenses/by/4.0/).たOlmo-3をHELMETで評価した結果．8K，16K，32K，64Kは入力長を示す．Olmo-3-7B-Think負例生成方法8K16K32K64K学習前46.5345.8342.4135.07小規模モデル45.2644.4442.6033.92非RHマスク45.5444.6542.9534.22RetMask47.0745.1942.6835.07Llama-3.1-8B-InstructQwen3-8BOlmo-3-7B-Think4.43.73.12.84.57.06.623.810241152102455.437.867.981.2=00-0.050.05-0.10.1-0.50.5-1.0有すると考える．すると，Llama-3.1では検索機能が約3割のヘッドに集中しているのに対し，Qwen-3では約4割，Olmo-3では9割以上のヘッドに分散している．すなわち，検索ヘッドを不活性化した場合，Llama-3.1とQwen3は検索機能の多くを喪失する．一方，Olmo-3では検索ヘッドを不活性化しても，検索機能をある程度有するヘッドが多数残存するため，検索の機序が維持されると考えられる．この結果は前述の仮説を支持しており，RetMaskの有効性は検索機能の集中度に依存することを示している．また，RetMaskの適用に先立ち検索スコアの分布を調査することで，提案手法の有効性を事前に予測できる可能性が示唆される．5.2RetMaskがモデルに与える影響RetMaskの学習がモデルにどのような影響をもたらしたかを分析する．図4では，学習前の検索スコアが高い上位150個のヘッドについて，学習前後の検索スコアの変化を示している．赤い点線は不活性化の閾値を示しており，点線より左側のヘッドが検索ヘッドとしてデータ合成時に不活性化された．が顕著に上昇している．具体的には，検索スコアの平均値が0.017から0.020に上昇し，17.6%の相対的な向上を示した．Qwen3においても，平均値が0.020から0.021へと5%上昇したが，上昇幅が小さい．この事実は，Qwen3の長文脈処理能力の向上がLlama-3.1ほど顕著でないことと整合している．こRetMaskRetMask0.80.60.40.20.0(a)Llama-3.1-8B-Instruct.RetMaskRetMask0.80.60.40.20.0(b)Qwen3-8B.れは，検索ヘッドを不活性化したモデルの出力を負例とすることで，検索ヘッドの機序に特化した学習が行えたためと考えられる．また，検索スコアの改善は全てのヘッドで均一ではなく，不活性化されたヘッドに集中している．Llama-3.1では，不活性化されたヘッドの平均改善幅が0.051であるのに対し，そうでないヘッドでは変化が軽微（平均+0.001）であった．これは，RetMaskが学習時に対象とした検索ヘッドを選択的に強化していることを示している．6おわりに本稿では，検索ヘッドに着想を得てLLMの長文脈処理能力を向上させられるかを検証した．具体的には，検索ヘッドを選択的に不活性化したモデルから応答を合成し，選好最適化の負例とすることで，長文脈処理の強化に特化した最適化手法を提案した．提案手法はLlama-3.1およびQwen3で長文処理能力を改善したものの，Olmo-3での効果は限定的だった．その要因を分析したところ，提案手法の効果はモデル内部における検索機能の集中度に依存することが明らかになった．今後は，検索ヘッドを不活性化した合成データによる選好最適化が有効である理由をより詳細に解明したい．org/licenses/by/4.0/).