\documentclass[japanese]{jnlp_3.8}
\usepackage{jnlpbbl_2.3}
\usepackage[dvipdfmx]{graphicx}
\usepackage{amsmath}
\usepackage{amssymb}
\usepackage[hyphens]{url}
\usepackage{array}
\usepackage{hangcaption_jnlp}
\allowdisplaybreaks


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%% author's definition
% custom packages
\usepackage{bm}
\usepackage{bbm}
\usepackage{booktabs}
%%%%\usepackage{subcaption}
\usepackage{multirow}
%%%%\usepackage{makecell}
%%%%\usepackage{color}
\usepackage{cleveref}
% custom commands
\crefname{equation}{式}{式}
\crefname{figure}{図}{図}
\crefname{table}{表}{表}
\crefname{section}{}{}
\creflabelformat{section}{#2#1章#3}
\crefname{subsection}{}{}
\creflabelformat{subsection}{#2#1節#3}
\crefname{subsubsection}{}{}
\creflabelformat{subsubsection}{#2#1項#3}
%%%%\usepackage{fancyvrb}
\crefname{subfigure}{図}{図}
%%%%\renewcommand\thesubfigure{(\alph{subfigure})}
%%%%\captionsetup[subfigure]{labelformat=simple}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

\Volume{32}
\Number{4}
\Month{December}
\Year{2025}
%%%%\Issuetitle{「　　　　」特集号}

\received{2025}{4}{1}
\revised{2025}{8}{1}
\accepted{2025}{9}{20}

\setcounter{page}{1272}

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\jsubject{一般論文}
\jtitle{$k$NN 言語モデルの低頻度語に対する振る舞いの分析}
\jauthor{西田　悠人\affiref{NAIST} \and 森下　　睦\affiref{Future} \and 出口　祥之\affiref{NAIST} \and 上垣外英剛\affiref{NAIST} \and 渡辺　太郎\affiref{NAIST}}
\jabstract{%
検索拡張言語モデルの一つである $k$NN 言語モデルは，任意のテキストデータから構築された大規模なデータストアに推論時に直接アクセスすることで，文脈を適切に把握し，言語らしさを高精度に予測可能であることが報告されている．$k$NN 言語モデルは明示的な記憶であるデータストアの活用によって低頻度現象の予測性能を改善できるという仮説が提唱されてきたが，既存の検証では主に低頻度の文脈の検索性能に焦点が当てられ，低頻度語の予測性能については十分に検討されてこなかった．本研究では，$k$NN 言語モデルが低頻度語に対してどのように振る舞うかを，予測確率や検索精度，データストア内におけるトークンの幾何的な分布，直積量子化による近似誤差の観点から詳細に分析した．実験の結果，$k$NN 言語モデルは低頻度語の予測性能の改善に寄与せず，むしろ高頻度語の予測を改善することを示した．
\renewcommand{\thefootnote}{}
\footnote[0]{本稿は，2025 Annual Conference of the North American Chapter of the Association for Computational Linguisticsで発表予定のもの \cite{nishida-etal-2025-long}の翻訳版である．
なお，本研究の基となった研究を言語処理学会第30回年次大会にて発表している \cite{nlp_2024_nishida}が，本稿には\cref{sec:background}「背景」および，\cref{sec:context_frequency_analysis}と\cref{sec:quantization_analysis}，\cref{sec:in_domain}の分析が追加されている．}
\renewcommand{\thefootnote}{\arabic{footnote}}
}
\jkeywords{言語モデル，近傍探索，低頻度語}

\etitle{Long-Tail Crisis in Nearest Neighbor Language Models}
\eauthor{Yuto Nishida\affiref{NAIST} \and Makoto Morishita\affiref{Future} \and Hiroyuki Deguchi\affiref{NAIST} \and Hidetaka Kamigaito\affiref{NAIST} \and Taro Watanabe\affiref{NAIST}}
\eabstract{%
The $k$-nearest-neighbor language model ($k$NN-LM), one of the retrieval-augmented language models, improves the perplexity for given text by directly accessing a large datastore built from any text data during inference. A widely held hypothesis for the success of $k$NN-LM is that its explicit memory, i.e., the datastore, enhances predictions for long-tail phenomena. However, prior works have primarily shown its ability to retrieve long-tail contexts, leaving the model's performance remain underexplored in estimating the probabilities of long-tail target tokens during inference. In this paper, we investigate the behavior of $k$NN-LM on low-frequency tokens, examining prediction probability, retrieval accuracy, token distribution in the datastore, and approximation error of the product quantization. Our experimental results reveal that $k$NN-LM does not improve prediction performance for low-frequency tokens.
}
\ekeywords{Language Model, Nearest Neighbor Search, Long-tail Phenomena}

\headauthor{西田，森下，出口，上垣外，渡辺}
\headtitle{$k$NN言語モデルの低頻度語に対する振る舞いの分析}

\affilabel{NAIST}{奈良先端科学技術大学院大学}{Nara Institute of Science and Technology}
\affilabel{Future}{フューチャー株式会社}{Future Corporation}



%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\begin{document}
\maketitle


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\section{はじめに}

検索拡張言語モデル \cite{grave2017unbounded,guu-etal-2018-generating,Khandelwal2020Generalization,He2020Learning,pmlr-v162-borgeaud22a}は，ベースとなる言語モデルと外部のデータストアから検索された近傍事例を組み合わせて予測確率を計算する，言語モデリングの新たなパラダイムであり，言語らしさを高精度に予測可能であることが報告されている．
このパラダイムにおける代表的なモデルの一つが$k$NN言語モデル \cite{Khandelwal2020Generalization}であり，ベースとなる訓練済み言語モデルの予測確率と，データストア内の近傍事例から計算された$k$NN確率を各時刻で補間する．
データストアは，任意のテキストデータのすべてのトークンに対するベースモデルの中間表現をキー，キーに対応するトークンをバリューとしてキャッシュすることで構築される．
また，各時刻の$k$NN確率は当該時刻のベースモデルの中間表現（文脈化埋め込み）とデータストア内のキーとの距離に基づいて検索された近傍トークンを用いて計算される．
したがって，$k$NN言語モデルは，推論時にデータストアを通じて直接テキストデータにアクセスすることで，与えられたテスト入力における各トークンの確率を高精度に予測することが可能である．

検索拡張言語モデルが言語らしさを高精度に予測可能な要因として，言語モデルのパラメトリックな記憶とは異なる外部の記憶であるデータストアの明示的な利用が，特に低頻度トークンに対する予測性能を向上させるという仮説が提唱されている \cite{grave2017unbounded,li20ga_interspeech}．
\citeA{Khandelwal2020Generalization}は，$k$NN言語モデルの成功の要因も上記の仮説と同様にデータストアへの直接的なアクセスが低頻度の現象に有効なためであると主張し，仮説の妥当性を示唆するいくつかの定性的な事例を報告している．
しかし，$k$NN言語モデルは推論時に予測対象のトークンである目的トークンの近傍となる埋め込みを検索して利用するにもかかわらず，既存の検証では主に低頻度の文脈，すなわち目的トークンの前に現れるトークン系列に焦点が当てられてきた．
$k$NN言語モデルが低頻度の文脈に対して予測性能を改善するのであれば，低頻度の目的トークンに対しても同様の改善が期待されるが，本稿の事前実験（\cref{sec:preliminary_experiments}）において，目的トークンの頻度とその文脈の$n$-gramの頻度は無相関であることが示された．
このため，$k$NN言語モデルが低頻度の目的トークンに対してどのように振る舞うかは依然として明らかでない．

上記の仮説をさらに検証するため，本研究では低頻度の現象のなかでも低頻度の目的トークンに焦点を当て，詳細な定量分析を実施した．
GPT2-XL \cite{radford2019language}をベースモデルとし，WikiText-103を再分割したデータセットを用いた実験によって分析を行い，予測確率や検索精度，データストアにおける文脈化トークン埋め込みの分布，直積量子化の近似誤差の観点から次の4つの知見が得られた．
(1) 低頻度トークンに対する$k$NN確率はベースモデルが割り当てる確率より低く，高頻度トークンに対する$k$NN確率は高いことが観察され（\cref{sec:knn_prob_analysis}），これまでの仮説に反して，$k$NN言語モデルは低頻度トークンの予測性能を改善せず，むしろ高頻度トークンの予測性能を向上させることが示された．
(2) 低頻度の目的トークンは多くの場合に近傍探索でヒットしないことが観察され（\cref{sec:knn_search_analysis}），この低頻度トークンの検索の困難さが$k$NN言語モデルの低頻度トークンに対する予測性能の低下を引き起こす要因の一つとなっていることが示唆された．
(3) データストア内において，低頻度トークンは高頻度トークンよりも疎に分布しており，かつ多くのケースでコンタミネーションが発生していることが明らかとなった（\cref{sec:datastore_distribution_analysis}）．このことが近傍検索をより困難にしていると考えられる．
(4) 低頻度トークンはデータストアの量子化に際してより大きな誤差を生じており（\cref{sec:quantization_analysis}），近似近傍検索\footnote{計算の効率化のために探索空間の削減などを行い，厳密な最近傍ではないがそれに近いデータ点を高速で検索するアルゴリズム．}をさらに困難にしていることが示唆された．


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\section{背景}
\label{sec:background}

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\subsection{検索拡張言語モデル}

検索拡張言語モデルは，言語モデルが推論時に外部のデータストアから関連する情報を検索し，その情報を利用して言語モデリングをはじめとする自然言語処理タスクに適用する手法である．
検索拡張言語モデルにおいて，検索される情報の粒度はテキストチャンク（文書や段落，文）と句，トークンに大別できる \cite{asai2024reliable}．


大局的な情報（テキストチャンク）を扱う検索拡張言語モデルでは，元の入力に関連する文や文書をデータストアから検索し，検索結果を入力に追加して言語モデルで推論を行う入力拡張 \cite{chen-etal-2017-reading,Lewis2020RAG,yao2023react,asai2024selfrag}の方式が広く用いられている．
入力拡張方式において代表的な手法である検索拡張生成 (RAG; \citeB{Lewis2020RAG}) は，事実知識を問う質問応答などの知識集約型タスクでの有効性が報告されている．
その他にも，検索されたテキストの中間表現を言語モデルの中間層における表現空間で統合する方法 \cite{wang-etal-2023-shall,pmlr-v162-borgeaud22a,wang2024instructretro}が提案されている．

局所的な情報（句やトークン）を扱う検索拡張言語モデルでは，言語モデルの推論時に言語モデルの出力するトークン分布と検索結果から計算されたトークン分布を補間することで予測を行う出力補間 \cite{Khandelwal2020Generalization,zhong-etal-2022-training,min-etal-2023-nonparametric}の方式が用いられる．
トークンを検索する$k$NN言語モデル \cite{Khandelwal2020Generalization}や，句を検索するノンパラメトリックマスク言語モデル \cite{min-etal-2023-nonparametric}が提案されている．


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\subsection{$k$NN言語モデル}

$k$NN言語モデル \cite{Khandelwal2020Generalization}は，局所的な情報を検索して利用する検索拡張言語モデルにおいて代表的な手法である．
極端なドメイン適応を要する言語モデリングの設定\footnote{寛容なライセンスの下で公開されたデータによって構築されたOLCコーパスで学習された言語モデルを用いて，制約のあるライセンスの下で公開されたテキストでの言語モデリングを行う設定．}では，RAGに代表される文書レベルの検索を利用する方法よりも，トークンレベルの検索を利用する$k$NN言語モデルの性能が優れていることが示されている \cite{min2024silo}．
また，$k$NN言語モデルが局所的な検索・補間を行うことを利用して，局所的な制御を要求されるスタイル変換タスクへの応用 \cite{trotta-etal-2022-nearest,nawezi-etal-2023-style}や，感情分析タスクにおけるトークンレベルの確信度較正への応用 \cite{hashimoto-etal-2025-efficient}が提案されており，その有効性が報告されている．


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\subsection{言語モデルによる低頻度語の予測}

近年の大規模言語モデル (LLM) は様々な自然言語処理タスクにおいて高い性能を示している．
しかし，LLMは低頻度な知識についての質問応答の性能が低い\cite{kandpal2023large,zong-etal-2025-comparisonqa}ことや，低頻度トークンに対する予測性能が低い\cite{pinto2024fairlanguagemodelparadox}ことが報告されており，低頻度現象を含む入出力に対する性能には依然として問題が残る．
このような低頻度現象に対するLLMの性能低下の問題を理解し対処するために研究が進められてきており，低頻度トークン予測に強い影響を持つサブネットワークの同定\cite{liu2025emergent}や，低頻度トークンの学習を促進する損失関数の提案\cite{su-etal-2024-mile}が行われている．
$k$NN言語モデルはこのような言語モデルの低頻度現象の予測の困難さを克服するための有望な方法であると考えられてきた\cite{Khandelwal2020Generalization}．
本研究では，$k$NN言語モデルの低頻度トークンの予測性能の分析を通じて，言語モデルの低頻度現象におけるトークンレベルの検索の有効性についての示唆を与える．


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\section{問題設定}
\label{sec:problem_statement}

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\subsection{$k$NN言語モデル}

$k$NN言語モデル \cite{Khandelwal2020Generalization}は検索拡張言語モデルの一つであり，\cref{fig:knnlm}に示すように，ベースとなる言語モデルと外部のデータストアから検索された近傍事例を組み合わせて利用して入力テキストの言語らしさを予測する．
事前に任意のテキストからデータストアを事前に構築し，推論時にはデータストアから検索された$k$個の近傍トークンを用いて$k$NN確率を計算し，ベースモデルによる予測確率と補間する．

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%% fig. 1
\begin{figure}[t]
\begin{center}
\includegraphics{32-4ia7f01.eps}
\end{center}
\hangcaption{$k$NN言語モデルの動作．目的トークン \textit{CNA} がデータストアから近傍事例として検索され，その事例を用いて計算された$k$NN確率によって補間することで予測確率が改善する．$k$NN言語モデルは明示的な記憶であるデータストアを用いることで，\textit{CNA} のような低頻度トークンを含む低頻度現象をより正確に予測できるという仮説が提唱されてきた．}
\label{fig:knnlm}
\end{figure}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\paragraph{データストアの構築}
データストアに格納するテキスト（トークン系列）を $\bm{x}=(x_{1},\cdots,x_{|\bm{x}|}) \in \mathcal{V}^*$ とする．ここで，$\mathcal{V}^*$は語彙$\mathcal{V}$のクリーネ閉包である．
データストア $\mathcal{D} \subseteq \mathbb{R}^{D} \times \mathcal{V}$ は，テキスト中のすべてのトークンに対するキーとバリューのペアの集合として次のように構築する:
\begin{equation}
    \label{eq:datastore}
    \mathcal{D} =
    \left \{(\bm{f}_\theta(\bm{x}_{<t}), x_{t}) \right \}_{t=1}^{|\bm{x}|}.
\end{equation}
それぞれのキーは時刻$t$における文脈系列$\bm{x}_{<t}$に対してベースモデル$\theta$を用いて計算される$D$次元の中間表現$\bm{f}_\theta(\bm{x}_{<t})$である ($\bm{f}_\theta \colon \mathcal{V}^* \to \mathbb{R}^{D}$)．
バリューはキーに対応する目的トークン$x_{t}$である．
なお，キー・バリューのペアは，キーが文脈化埋め込みであるため，トークンタイプ\footnote{語彙$\mathcal{V}$の要素であり，トークンの表層の種類．}ごとに格納されるわけではなく，すべてのトークンを対象として格納され，あるトークンタイプに対応するペアの数は，テキスト中における当該トークンタイプの生起頻度と一致する．
データストアは与えられたテキストにおけるすべてのトークンに対する高次元の中間表現を格納するため，非常に大規模となる．
そのため，データストア内のキーは一般的に直積量子化 \cite{jegou2011product}によって圧縮される．


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\paragraph{$k$NN確率の計算}
テスト入力$\bm{x}$が与えられたとき，ベースモデル$\theta$はそれぞれの目的トークン$x_{t}$に対して中間表現$\bm{q}_{t}=\bm{f}_\theta(\bm{x}_{<t})$を計算する．
この表現をクエリとしてデータストア$\mathcal{D}$から上位$k$個の近傍事例$\mathcal{N}(\bm{q}_t; \mathcal{D}) \subset \mathcal{D}$を検索する．
目的トークン$x_{t}$に対する$k$NN確率$p_{\text{$k$NN}}$は，クエリ$\bm{q}_{t}$と各近傍事例のキー$\bm{k}$との距離に基づいて次のように計算される：
\begin{equation}
\label{eq:$k$NN_probability}
    p_{\text{$k$NN}}(x_{t}|\bm{x}_{<t}) \nonumber \propto 
    \sum_{(\bm{k},v) \in \mathcal{N}(\bm{q}_t; \mathcal{D})}{\mathbbm{1}_{x_{t}=v}\exp \left(\frac{-d(\bm{k},\bm{q}_{t})}{\tau}\right )}.
\end{equation}
ここで，$d(\cdot,\cdot)$ は距離関数，$\tau$ はソフトマックス関数関数の温度パラメータである．
なお，温度パラメータ$\tau$は元の$k$NN言語モデル \cite{Khandelwal2020Generalization}では導入されていないが，本研究では，エンコーダ・デコーダ構造に対応する$k$NN機械翻訳 \cite{khandelwal2021nearest}に着想を得て，確率分布の平滑性を制御するために導入した．
$k$NN確率の計算においては，距離が近いトークンほど高い確率が割り当てられ，あるトークンが近傍に含まれない場合にはその$k$NN確率は0となる．
ここで，温度$\tau$は距離の影響度を調整する役割を果たし，$\tau \to \infty$の極限においては，確率は距離に依存せず近傍内でのトークンの相対頻度に一致する．


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\paragraph{補間}
文脈系列$\bm{x}_{<t}$が与えられたとき，目的トークン$x_{t}$の予測確率$p$は，$k$NN確率$p_\text{$k$NN}$とベースモデルによる予測確率（LM確率）$p_\text{LM}$の線形補間により次のように計算される：
\begin{equation}
\label{eq:interpolation}
    p(x_{t}|\bm{x}_{<t}) = \lambda p_\text{$k$NN}(x_{t}|\bm{x}_{<t}) + (1 - \lambda)p_\text{LM}(x_{t}|\bm{x}_{<t}).
\end{equation}
ここで，$\lambda$は線形補間における$k$NN確率の重みを制御するハイパーパラメータ（補間係数）である．
$\lambda > 0$の条件のもとでは，$k$NN確率$p_\text{$k$NN}$がLM確率$p_\text{LM}$より大きいとき，$p > p_\text{LM}$となり予測性能が改善される\footnote{$p = \lambda p_\text{$k$NN} + (1 - \lambda)p_\text{LM} > \lambda p_\text{LM} + (1 - \lambda)p_\text{LM} = p_\text{LM}$ であるため．}．

推論時には，$k$NN言語モデルはデータストアを通じて直接テキストデータにアクセスする．
\citeA{Khandelwal2020Generalization}は，このような明示的な記憶の利用によって，低頻度現象の予測が改善されるという仮説を提唱した．


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\subsection{研究課題}

検索拡張言語モデルは，明示的に外部の記憶を活用することで低頻度トークンに対して特に有効であるという仮説が提唱されてきた \cite{grave2017unbounded,li20ga_interspeech}．
検索拡張言語モデルの一つである$k$NN言語モデルも同様に，データストアから得られる$k$NN確率との補間により，低頻度の目的トークンに対する予測を改善することが期待される．
しかし，既存の検証では低頻度の文脈に対する予測改善は示されている \cite{Khandelwal2020Generalization}ものの，\cref{sec:preliminary_experiments}で示すように，$n$-gramの文脈系列の頻度と目的トークンの頻度には相関がみられなかった．
したがって，低頻度の目的トークンに対する$k$NN言語モデルの振る舞いは依然として明らかではない．

この乖離を解消するため，本研究では，トークンの頻度と$k$NN言語モデルの振る舞いの関係を分析する．


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\paragraph{研究課題1：$k$NN言語モデルは低頻度の目的トークンに対する予測性能を改善するか？}
トークン頻度とLM確率および$k$NN確率の関係を調査し，$k$NN言語モデルは低頻度トークンの予測を悪化させ，むしろ高頻度トークンの予測を改善することを示す（\cref{sec:knn_prob_analysis}）．

上記の結果を踏まえ，低頻度トークンに対する性能低下の要因を$k$NN確率の計算過程に基づいて分析する．


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\paragraph{研究課題2：$k$NN言語モデルは低頻度トークンを適切に検索できるか？}
$k$NN確率は近傍事例を用いて計算されるため，検索の精度に依存する．
トークン頻度と検索の精度の関係を調査し，頻度が低いほど検索にヒットしない割合が増加する傾向が観察された（\cref{sec:knn_search_analysis}）．
すなわち，低頻度トークンの検索の困難さが予測性能を悪化させる要因となっていると考えられる．

さらに，低頻度トークンの検索が困難である要因について調査するためにデータストアの分析を行う．


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\paragraph{研究課題3：低頻度トークンはデータストアにどのように分布するか？}
トークン頻度とデータストアの幾何的な分布の関係を調査し，低頻度トークンはデータストア内で疎に分布し，かつ他のトークンと混交して分布していることを示した（\cref{sec:datastore_distribution_analysis}）．
このようなデータストアの分布が低頻度トークンの検索を困難にしていると考えられる．


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\paragraph{研究課題4：低頻度トークンは量子化されたデータストアにおいて正確に表現されるか？}
$k$NN言語モデルでは，大規模なデータストアを扱うために直積量子化による圧縮を施す．
この量子化の過程における誤差とトークン頻度の関係を調査し，低頻度トークンは量子化誤差が大きくなる傾向にあることを示す（\cref{sec:quantization_analysis}）．
このことが低頻度トークンの近似近傍探索を困難にしていると考えられる．


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\section{実験設定}
\label{sec:experimental_settings}

本節では，本研究における$k$NN言語モデルの分析に使用する共通の実験設定について述べる．


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\subsection{ベースモデル}
\label{sec:base_lm}

本研究では，ベースとなる言語モデルとしてGPT2-XL \cite{radford2019language}を用いた．
このモデルは1.5Bパラメータを持ち，Reddit上の外部リンクから収集された英語Webテキストにより事前学習されている．
$k$NN言語モデルのベースモデルとしてGPT-2を用いることは一般的であり，既存の検証においてもその有効性が示されている \cite{shi-etal-2022-nearest,xu2023nearest,wang-etal-2023-knn}．
なお，GPT-2は言語モデリングタスクの評価データとして一般的に用いられるWikipediaを事前学習データから意図的に除外している．
Wikipediaドメインのテキストデータを訓練データに含む言語モデルをベースモデルとして用いた分析については\cref{sec:in_domain}で行った．


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\subsection{トークン頻度}
\label{sec:token_frequency}

本研究では，トークン頻度と$k$NN言語モデルの性能との関係を多角的に分析するため，2種類の頻度を考慮した．
第一に，$k$NN言語モデルのデータストア構築に用いたテキスト中のトークン頻度（データストア頻度）である．
第二に，ベースモデルの事前学習に使用されたテキストにおけるトークン頻度（事前学習頻度）である．
なお，ベースモデルであるGPT-2の学習データは非公開であるため，その代替として，GPT-2と同様の方針で収集されたOpenWebTextCorpus \cite{Gokaslan2019OpenWeb}のトークン頻度を事前学習頻度として用いた．


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\subsection{データセット}
\label{sec:dataset}

評価にはWikiText-103 \cite{merity2017pointer}を使用した．
オリジナルのWikiText-103における訓練・検証・テストデータは重複する語彙が多いため，低頻度トークンの挙動の分析には適さない．
そこで本研究では，オリジナルの訓練データを語彙の重複が少なくなるように再分割し，新たな訓練・検証・テストデータを構築することで，低頻度トークンが多く含まれるような設定で評価を実施した．
データ再分割の具体的な手順を次に示す．
\begin{enumerate}
  \item ベースモデルのサブワードトークナイザで元の訓練データをトークン化する．
  \item オリジナルの訓練データの文書集合から，データ内の生起頻度が1である$n$-gram ($n=1,2,3,4$) の占める割合が大きい順に文書を抽出する．なお，抽出する文書の数はオリジナルの検証・テストデータの合計トークン数に同程度となるように決定する．
  \item 抽出した文書の集合をランダムに2つに分割し，新たな検証・テストデータとし，抽出されなかった残りの文書集合を新たな訓練データとする．
\end{enumerate}
\cref{tab:dataset_size}に，オリジナルおよび再分割データにおけるトークン数を示す．
また，\cref{fig:test_data_count}の左側と右側はそれぞれ，元のデータと再分割データのテストデータにおけるデータストア頻度（すなわち対応する訓練データにおけるトークン頻度）ごとおよび事前学習頻度ごとのトークン数を示している．
図より，再分割テストデータにはデータストア頻度が$10^3$未満のトークンがオリジナルデータより多く含まれており，かつ事前学習頻度の$10^4$未満のトークンがやや多く含まれていることが分かる．
したがって，本節で作成した再分割データは，低頻度トークンの分析においてオリジナルの分割よりも適しているといえる．

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%% table 1
\begin{table}[b]
\input{07table01.tex}%
\caption{データセットのトークン数}
\label{tab:dataset_size}
\end{table}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%% fig. 2
\begin{figure}[t]
\begin{center}
\includegraphics{32-4ia7f02.eps}
\end{center}
\hangcaption{WikiText-103のオリジナルのテストデータおよび再分割テストデータにおけるトークン出現数と，データストア頻度（左側）および事前学習頻度（右側）との関係．再分割データは，オリジナルデータよりも多くの低頻度トークンを含んでいる．}
\label{fig:test_data_count}
\end{figure}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\subsection{$k$NN言語モデル}
\label{sec:knn_lm}

本研究では，$k$NN言語モデルの実装にknn-transformers \cite{alon2022neuro}を用いた．
データストアの構築および近傍探索にはFaiss \cite{johnson2019billion,douze2024faiss}を使用し，距離関数には二乗ユークリッド距離を用いた．
データストア内のキーは，ベースモデルの最終層におけるフィードフォワードネットワークへの入力部分に対応する1,600次元の中間表現とした．
効率化のため，データストアにはIVFPQ \cite{jegou2011product}による直積量子化を施し，コードブック数を64，コード数を256，セントロイド数を4,096に設定した．
Coarse quantizerにはIndexFlatL2を用いた．
近傍探索時には，32個の近傍クラスタを探索する設定とした．
$k$NN確率の補間係数 $\lambda$ は\citeA{Khandelwal2020Generalization}に倣って0.25とした．
推論および評価は，オリジナルまたは再分割後の訓練データから構築したデータストアを用いてそれぞれのテストデータに対して行った．


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\subsection{事前実験}
\label{sec:preliminary_experiments}


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\subsubsection{テストデータ全体に対するパープレキシティ}

$k$NN言語モデルの詳細な分析を行う前に，低頻度トークンを多く含むデータでの$k$NN言語モデルの有効性を調査するために，ベースモデルと$k$NN言語モデルの性能をオリジナルおよび再分割データを用いて比較する．
評価は，各データの訓練データからデータストアを構築し，それぞれのテストデータでパープレキシティ (PPL) を測定することで行った．
なお，$k$NN言語モデルにおける近傍数$k$およびソフトマックス関数の温度$\tau$は，検証データでPPLが最小となるように$k \in \{2,4,8,16,32,64,128,256,512,1{,}024\}$ および $\tau \in \{0.01, 0.1, 1, 10, 100, 1{,}000\}$ から選択し，オリジナルデータ・再分割データの双方において$k=1{,}024$, $\tau=10$とした．

\cref{tab:ppl}に，各データセットのテストデータにおけるベースモデルおよび$k$NN言語モデルのPPLを示す．
再分割データにおいても$k$NN言語モデルによって性能が改善しており，低頻度トークンを多く含むデータにおいても$k$NN言語モデルが有効であることが示された．
しかし，再分割データにおける改善の幅はオリジナルデータと比較して小さく，低頻度トークンが多く含まれていることによってその有効性が制限されていることが示唆された．

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%% table 2
\begin{table}[t]
\input{07table02.tex}%
\caption{WikiText-103テストデータにおけるPPL（値が小さいほど性能が高いことを示す）}
\label{tab:ppl}
\end{table}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\subsubsection{低頻度文脈に対する予測}
\label{sec:context_frequency_analysis}

$k$NN言語モデルのような検索拡張言語モデルは，低頻度現象に対して特に効果的であるという仮説が提唱されている \cite{grave2017unbounded,merity2017pointer,li20ga_interspeech}．
\citeA{Khandelwal2020Generalization}は，目的トークンの直前に現れる文脈系列が低頻度であるときに$k$NN言語モデルの予測性能が改善することを示唆する定性的な例を示した．
しかし，再分割テストデータを用いて文脈$n$-gram（目的トークンに先行する$n$トークンの系列）の頻度と予測確率の関係を調査したところ，\cref{fig:prob_vs_context_ngram}に示すように，$k$NN言語モデルが低頻度文脈に対して特に予測性能を改善するわけではないことが示唆された．
$n=1,2,3$では，文脈系列の頻度が$10^3$以下の範囲において$k$NN確率はLM確率を下回っており，$n=4$では，文脈系列の頻度が$10^2$以下の範囲において両確率はほとんど差がない．$n=5$では，文脈系列が低頻度の場合においても$k$NN確率はLM確率を0.02程度上回っているが，これは長い文脈$n$-gramの頻度の疎性によって多くの文脈系列が頻度0となってしまうためである．

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%% fig. 3
\begin{figure}[p]
\begin{center}
\includegraphics{32-4ia7f03.eps}
\end{center}
\caption{再分割テストデータにおける文脈$n$-gramの頻度と$k$NN確率およびLM確率の関係．}
\label{fig:prob_vs_context_ngram}
\end{figure}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%% table 3
\begin{table}[p]
\input{07table03.tex}%
\hangcaption{再分割テストデータにおける目的トークンの頻度と文脈$n$-gram ($n=1,\cdots,5$) の頻度のピアソンの相関係数．}
\label{tab:preliminary_corr}
\end{table}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

また，$k$NN言語モデルのデータストアは文脈系列の中間表現に対応するキーと目的トークンに対応するバリューのペアとして構築されるため，低頻度現象は「低頻度の文脈系列」と「低頻度の目的トークン」に大別できる．
したがって，$k$NN言語モデルの低頻度現象に対する挙動を正確に理解するためには，文脈系列だけでなく目的トークンに着目した分析が必要である．
文脈系列に関して観察された関係が目的トークンにも当てはまるかを調べるために，再分割テストデータにおける文脈$n$-gram ($n=1,\cdots,5$) の頻度と目的トークンの頻度との相関を調査し，\cref{tab:preliminary_corr}に示した．
文脈と目的トークンの頻度の間には強い相関はみられず，低頻度の文脈系列で観察された現象が低頻度の目的トークンに対しても同様に観察されるとは限らないことが示された．
そのため，
本研究では$k$NN言語モデルの低頻度目的トークンに対する挙動を分析し，$k$NN言語モデルの性能に寄与する要因を明らかにする．


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\section{分析}
\label{sec:analysis}


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\subsection{$k$NN言語モデルは低頻度トークンの予測確率を改善しない}
\label{sec:knn_prob_analysis}

\citeA{Khandelwal2020Generalization}は，明示的な記憶であるデータストアを利用することで，$k$NN言語モデルが低頻度現象の予測性能を改善するという仮説を提唱し，低頻度の文脈の予測を改善する定性的な例を示した．
この仮説をさらに検証するために，本研究では文脈ではなく$k$NN言語モデルによって実際に検索される目的トークンに着目し，目的トークンの頻度と予測確率の関係を定量的に分析した．

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%% fig. 4
\begin{figure}[b]
\vspace{-0.5\Cvs}
\begin{center}
\includegraphics{32-4ia7f04.eps}
\end{center}
\hangcaption{再分割テストデータにおける，トークン頻度と$k$NN確率およびLM確率の期待値の関係．頻度が低い場合には$k$NN確率がLM確率を下回り，頻度が高い場合にはその逆の傾向が観察された．}
\label{fig:knn_prob}
%%%%\label{fig:knn_prob_datastore}
%%%%\label{fig:knn_prob_openwebtext}
\end{figure}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\subsubsection{目的トークンの頻度と予測確率の関係}

\cref{eq:interpolation}によると，目的トークンに対する$k$NN確率がLM確率を上回るとき，$k$NN言語モデルの予測が改善される．
したがって，低頻度トークンに対して$k$NN確率とLM確率の大小を比較することで，$k$NN言語モデルによる予測性能の改善が生じているかどうかを確認できる．
\cref{fig:knn_prob}は，再分割テストデータにおけるデータストア内のトークン頻度と，$k$NN確率およびベースモデル確率の期待値との関係を示している．
%%%%\cref{fig:knn_prob_datastore}より，
    \cref{fig:knn_prob}(a)より，
データストア頻度が約$10^5$以下のトークンでは$k$NN確率がLM確率よりも低く，特に$10^3$以下ではその差が顕著であることがわかる．
これは，$k$NN確率を補間することで低頻度トークンの予測性能がかえって悪化していることを示している．
一方で，頻度が$10^5$を超えるトークンでは，$k$NN確率がおおむねLM確率以上となっており，高頻度トークンの予測性能は改善していることが観察された．
%%%%また，\cref{fig:knn_prob_openwebtext}より，
    また，\cref{fig:knn_prob}(b)より，
事前学習頻度においても同様に，$10^7$以下の頻度の低い範囲ではおおむね$k$NN確率がLM確率よりも低く，$10^8$以上の頻度の高い範囲ではおおむね$k$NN確率がLM確率以上となる傾向が確認された．


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\subsubsection{$k$NN言語モデルのハイパーパラメータと予測確率の関係}

$k$NN言語モデルの予測確率の振る舞いをより深く理解するために，ハイパーパラメータ（近傍数$k$および温度$\tau$）の影響を調査した．
ベースモデルの語彙をデータストア頻度でソートし，LOW, MED, HIGHの3カテゴリに等分割し，それぞれのカテゴリについて予測確率の平均とハイパーパラメータの関係を可視化した．

\cref{fig:knn_prob_by_hyperparameters}に，再分割テストデータにおけるそれぞれの頻度カテゴリにおける$k$NN確率およびLM確率とハイパーパラメータの関係を示す．なお，近傍数による影響は温度$\tau=1$として観察した．
%%%%\cref{fig:knn_prob_by_k}より，
    \cref{fig:knn_prob_by_hyperparameters}(a)より，
$k$NN確率は近傍数$k$の増加によってほとんど変化せず，LOWカテゴリに属するトークンでは常にLM確率を下回っていることがわかる．
%%%%\cref{fig:knn_prob_by_temp}より，
    \cref{fig:knn_prob_by_hyperparameters}(b)より，
温度$\tau$を大きくすることで，LOWカテゴリにおける$k$NN確率の低さがさらに顕著になることが示された．
これらの結果は，$k$NN言語モデルが低頻度トークンに対して予測確率を悪化させる傾向が，ハイパーパラメータの調整では解消できないことを示している．

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%% fig. 5
\begin{figure}[b]
\begin{center}
\includegraphics{32-4ia7f05.eps}
\end{center}
\caption{頻度カテゴリごとのハイパーパラメータと$k$NN確率およびLM確率の関係．}
\label{fig:knn_prob_by_hyperparameters}
%%%%\label{fig:knn_prob_by_k}
%%%%\label{fig:knn_prob_by_temp}
\end{figure}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\subsubsection{定性分析}

トークン頻度と予測確率の関係をより直観的に理解するため，\cref{tab:example}に再分割テストデータにおける定性的な例を示す\footnote{なお，本節および表中の目的トークンおよび近傍トークンは，実際には\texttt{intoxicated}のように最初に空白文字を含むトークンであるが，読みやすさのために\textit{intoxicated}のように表記している．}．
目的トークン \textit{intoxicated}, \textit{FreeBSD}, \textit{recommendation} はいずれもデータストア内の頻度が$10^3$以下であり，$k$NN確率はLM確率よりも低い．
一方で，頻度が$67,141$のトークン \textit{song} に関しては，$k$NN確率がLM確率よりも高い．

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%% table 4
\begin{table}[t]
\input{07table04.tex}%
\hangcaption{定性的な例 ($k=1024$, $\tau=10$)．目的トークン \textit{intoxicated}, \textit{FreeBSD}, \textit{recommendation} は，データストア頻度が$10^3$以下であり，$k$NN確率がLM確率を下回り，$k$NNヒット率も60\%未満であった．一方で，目的トークン \textit{song} はデータストア頻度が$67,141$であり，$k$NN確率がLM確率より高く，近傍探索によるヒット率も高かった．表の右側には，各目的トークンの再分割テストデータにおけるサンプル文脈と，対応する上位3つの近傍トークンを示している．低頻度トークンにおいては，それぞれの文脈に対して誤ったトークンの近傍事例として検索されているのに対し，高頻度トークンでは正しい目的トークンが近傍に含まれていた．}
\label{tab:example}
\end{table}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

以上より，\textbf{$k$NN言語モデルは低頻度トークンの予測性能を改善しない}こと，および高頻度トークンに対して予測性能を向上させる傾向があることが明らかとなった．


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\subsection{低頻度トークンに対する$k$NN検索は困難である}
\label{sec:knn_search_analysis}

\cref{sec:knn_prob_analysis}で示したように，$k$NN確率はトークンの頻度が低いほど小さくなる．
本節では，その原因として低頻度トークンは近傍探索において検索されにくいという仮説を立ててその検証を行う．
特に，目的トークンが近傍事例に1件も含まれない場合には$k$NN確率は0となり，低頻度トークンの予測確率に大きな負の影響を与えるため，近傍事例集合$\mathcal{N}$に目的トークンが1件以上含まれている割合（$k$NNヒット率）に焦点を当てて分析を行った．

\cref{fig:knn_hit_ratio}に，トークン頻度と$k$NNヒット率の関係を示す．
データストア頻度と事前学習頻度に共通して，近傍数$k$が増加するとトークン頻度に関係なく$k$NNヒット率が向上する傾向が観察された．
%%%%\cref{fig:knn_hit_ratio_datastore}より，
    \cref{fig:knn_hit_ratio}(a)より，
データストア頻度が低いトークンは$k$NNヒット率が低い傾向が観察された．
$k=1{,}024$においては，データストア頻度が$10^5$以上のトークンではヒット率はほぼ100\%である一方で，頻度が10未満の場合ではヒット率は50\%未満（すなわち過半数のケースで$k$NN確率が0）となった．
%%%%また，\cref{fig:knn_hit_ratio_openwebtext}より，
    また，\cref{fig:knn_hit_ratio}(b)より，
事前学習頻度においても同様に，頻度の低いトークンはヒット率が低い傾向が観察された．特に，頻度が約$10^3$以下のトークンでは$k=1{,}024$であってもヒット率は50\%を下回った．
この結果は，\textbf{多くの低頻度トークンにおいて目的トークンが全く検索できていない}ことを示しており，このことが$k$NN確率の低下を引き起こしていると考えられる．

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%% fig. 6
\begin{figure}[t]
\begin{center}
\includegraphics{32-4ia7f06.eps}
\end{center}
\hangcaption{再分割テストデータにおけるトークン頻度と$k$NNヒット率の関係．低頻度トークンにおいては，目的トークンが近傍事例に全く含まれない場合があった．}
\label{fig:knn_hit_ratio}
%%%%\label{fig:knn_hit_ratio_datastore}
%%%%\label{fig:knn_hit_ratio_openwebtext}
\end{figure}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

近傍数$k$を大きくすることで$k$NNヒット率は改善することが観察されたが，近傍事例に占める目的トークン以外のノイズが増加し，目的トークンの相対頻度が低下すると考えられる．
\cref{eq:$k$NN_probability}より，近傍に占める目的トークンの相対頻度が低下すると（特に温度$\tau$が大きい場合において）$k$NN確率は低下する．
このヒット率と近傍の相対頻度のトレードオフは，\cref{sec:knn_prob_analysis}で観察された，近傍数を大きくしても$k$NN確率が改善しない現象の要因となっていると考えられる．

\cref{tab:example}の定性的な例も頻度が低いほどヒット率が低い傾向を裏付けており，\textit{intoxicated}, \textit{FreeBSD}, \textit{recommendation} といった低頻度トークンのヒット率は低く，\textit{song} のような高頻度トークンではヒット率が高い．


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\subsection{低頻度トークンはデータストアで疎かつ混交して分布する}
\label{sec:datastore_distribution_analysis}

\cref{sec:knn_search_analysis}では，低頻度トークンが近傍探索で検索されにくいことを示した．
その理由の一つとして，データストア内のキーおよびトークンの分布に起因する可能性がある．
本節では，\textbf{分布の疎性}と\textbf{最近傍の混交率}の2つの観点から，データストアの低頻度トークンに関する分布構造を分析する．


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\subsubsection{データストアのキー分布の疎性}

特定のトークンタイプに対応するキーが疎に分布している場合，目的トークンに対応するクエリからの距離が大きくなり，近傍検索が困難になると考えられる．
この仮説を検証するため，それぞれのトークンタイプにおいて，当該トークンタイプに対応するキーの集合の重心から各キーへのユークリッド距離を計算し，その変動係数\footnote{標準偏差を平均で除した値で，分布のばらつきを示す．}によって分布の疎性を計測した．

\cref{fig:datastore_cv}に，再分割訓練データで構築したデータストアにおける，トークン頻度と距離の変動係数の関係を示す．
%%%%\cref{fig:datastore_cv_datastore}より，
    \cref{fig:datastore_cv}(a)より，
データストア頻度が低いトークンほど変動係数が高い傾向がみられた．
%%%%また，\cref{fig:datastore_cv_openwebtext}より，
    また，\cref{fig:datastore_cv}(b)より，
事前学習頻度においても同様に，頻度が低いトークンは変動係数が高い傾向がみられた．
なお，$10^8$を超えるような非常に高頻度なトークンはこの傾向から外れているが，このようなトークンではその多義性の影響により例外的な傾向が現れていると考えられる．
以上より，\textbf{低頻度トークンはデータストアで疎に分布している}ことが示され，このことが近傍探索が困難な要因の一つであると考えられる．

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%% fig. 7
\begin{figure}[b]
\begin{center}
\includegraphics{32-4ia7f07.eps}
\end{center}
\hangcaption{データストアにおけるトークン頻度と距離の変動係数の関係：頻度の低いトークンほどデータストア内で疎に分布している傾向がある．}
\label{fig:datastore_cv}
%%%%\label{fig:datastore_cv_datastore}
%%%%\label{fig:datastore_cv_openwebtext}
\end{figure}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\subsubsection{最近傍の混交率}

データストアのなかで，目的とするトークンタイプに対応するデータ点の近傍に異なるトークンタイプのデータ点が存在する場合，近傍探索で正しいトークンを検索することが困難となる．
この仮説を検証するために，各トークンタイプにおいて，最近傍のトークンタイプが一致しない割合（混交率）を算出した．

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%% fig. 8
\begin{figure}[t]
\begin{center}
\includegraphics{32-4ia7f08.eps}
\end{center}
\hangcaption{トークン頻度とトークンタイプの混交率との関係．データストア頻度が$10^2$以下または事前学習頻度$10^3$以下の低頻度の範囲および，データストア頻度が$10^5$以上または事前学習頻度$10^8$以下の高頻度の範囲において，頻度が低いトークンタイプほど混交率が高い傾向にあった．}
\label{fig:datastore_contamination}
%%%%\label{fig:datastore_contamination_datastore}
%%%%\label{fig:datastore_contamination_openwebtext}
\end{figure}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

\cref{fig:datastore_contamination}は，再分割訓練データで構築したデータストアにおける，トークン頻度と最近傍の混交率の関係を示す．
%%%%\cref{fig:datastore_contamination_datastore}より，
    \cref{fig:datastore_contamination}(a)より，
データストア頻度が低いほど混交率が高く，特に頻度10未満のトークンでは，ほとんどの場合で最近傍が異なるトークンタイプであった．
%%%%\cref{fig:datastore_contamination_openwebtext}より，
    \cref{fig:datastore_contamination}(b)より，
事前学習頻度においても同様に，頻度の低いトークンは混交率が高い傾向がみられ，$10^2$未満のトークンでは9割を超える割合で最近傍が異なるトークンタイプであった．

これは，\textbf{低頻度トークンは近傍に混入が多い}ことを示しており，検索困難性の一因である．
%%%%\cref{fig:datastore_contamination_openwebtext}は，
    \cref{fig:datastore_contamination}(b)は，
再分割テストデータにおける事前学習頻度と混交率の関係を示しており，事前学習頻度が低いトークンほど，近傍に他のトークンタイプが混入しやすい傾向が確認できる．


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\subsection{低頻度トークンはデータストアにおける量子化誤差が大きい}
\label{sec:quantization_analysis}

$k$NN言語モデルでは大規模なデータストアを扱うために，データストア内のキーを量子化し，近似近傍探索を行う．
このときの近似誤差が，低頻度トークンの近傍検索が困難であるもう一つの要因となっていると考えられる．
$k$NN言語モデルにおいてキーを圧縮するために用いられる直積量子化 \cite{jegou2011product}は，データストア内のキーの平均誤差を最小化するようにコードブックを学習するため，大部分のトークンを高精度に量子化することができるが，頻度の低いトークンに対しては精度が低下する可能性がある．
そこで本研究では，直積量子化によって量子化されたキーを復元したベクトルと量子化を行う前のキーのユークリッド距離（復元誤差）をトークンタイプごとに平均することで，データストアの量子化における誤差を測定した．

\cref{fig:reconstruction_error}に，再分割訓練データで構築したデータストアにおける，トークン頻度とキーの復元誤差の関係を示す．
%%%%\cref{fig:reconstruction_error_datastore}より，
    \cref{fig:reconstruction_error}(a)より，
データストア頻度が低いトークンほど直積量子化の近似誤差が大きい傾向が観察された．
%%%%\cref{fig:reconstruction_error_openwebtext}より，
    \cref{fig:reconstruction_error}(b)より，
事前学習頻度においても同様に，頻度が低いほど近似誤差が大きい傾向がみられた．
したがって，\textbf{低頻度トークンはデータストア量子化における復元誤差が大きい}ことが示され，このことが近似近傍探索を困難にしていると考えられる．

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%% fig. 9
\begin{figure}[t]
\begin{center}
\includegraphics{32-4ia7f09.eps}
\end{center}
\hangcaption{トークン頻度と直積量子化の再構成誤差との関係．低頻度トークンはデータストア内で近似誤差が大きい傾向にある．}
\label{fig:reconstruction_error}
%%%%\label{fig:reconstruction_error_datastore}        
%%%%\label{fig:reconstruction_error_openwebtext}
\end{figure}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%% fig. 10
\begin{figure}[t]
\begin{center}
\includegraphics{32-4ia7f10.eps}
\end{center}
\hangcaption{再分割テストデータにおける，直積量子化の復元誤差と$k$NN確率およびLM確率の関係．復元誤差が大きい場合には$k$NN確率がLM確率を下回り，予測性能が悪化する．}
\label{fig:reconstruction_error_vs_probability}
\end{figure}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

量子化の復元誤差の大きさが予測確率を毀損していることを確認するために，\cref{fig:reconstruction_error_vs_probability}に再分割テストデータにおける予測確率と復元誤差の関係を示す．
図より，誤差が小さい場合には$k$NN確率がLM確率を上回る一方で，誤差が大きいと$k$NN確率がLM確率を下回る傾向にあることが観察された．
また，トークンタイプごとに$k$NN言語モデルによる予測の改善幅 (8$p_{\text{$k$NN}} - p_{\text{LM}}$) と復元誤差の相関を測定すると，ピアソンの相関係数は$-0.20$，スピアマンの順位相関係数は$-0.16$であった．
これらの負の相関は，\textbf{復元誤差の大きなトークンほど$k$NN言語モデルの予測性能が悪化する}ことを示している．


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\subsection{in-domain設定における$k$NN言語モデルの分析}
\label{sec:in_domain}

\cref{sec:experimental_settings}および\cref{sec:analysis}における実験では，Wikipediaが除外された事前訓練データによって訓練されたGPT2-XLをベースモデルとして用いた．
$k$NN言語モデルのデータストアは事前訓練データに含まれないWikipediaデータを用いて構築し，同ドメインのテストデータによって評価を行ったため，本研究で示した実験はドメイン適応の設定であると見做せる．
本節では，本研究で観察された$k$NN言語モデルによるPPLの改善や低頻度トークンの性能低下がドメイン適応によるものであるかを弁別するために，in-domainの設定，すなわちモデルの訓練データにテストデータのドメインと同一のドメインのデータが含まれている設定での実験を行う．


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\subsubsection{実験設定}

Wikipediaを含むデータで訓練されたベースモデルによる実験を行うために，事前学習済みGPT2-XLをWikiText-103再分割訓練データの全体を用いてファインチューニングし，そのモデルをベースモデルとして用いた．
その他の設定は\cref{sec:experimental_settings}に準じた．


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\subsubsection{再分割テストデータ全体に対するパープレキシティ}

再分割テストデータにおけるin-domain設定のベースモデルおよび$k$NN言語モデルのPPLを\cref{tab:ppl_indomain}に示す．
なお，$k$NN言語モデルにおける近傍数$k$およびソフトマックス関数の温度$\tau$を，検証データでPPLが最小となるようにチューニングし，$k=1{,}024$, $\tau=10$とした．
表より，in-domain設定であっても$k$NN言語モデルはPPLの改善をもたらすことが観察され，\cref{sec:experimental_settings}でみられたPPLの改善がドメイン適応によるものではないことが示された．

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%% table 5
\begin{table}[b]
\input{07table05.tex}%
\caption{in-domain設定における，WikiText-103再分割テストデータのPPL}
\label{tab:ppl_indomain}
\end{table}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%% fig. 11
\begin{figure}[b]
\begin{center}
\includegraphics{32-4ia7f11.eps}
\end{center}
\hangcaption{in-domain設定における，再分割テストデータに対する$k$NN確率およびLM確率の期待値の関係．}
\label{fig:knn_prob_indomain}
\end{figure}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\subsubsection{$k$NN言語モデルによる低頻度トークンの予測}

\cref{sec:analysis}と同様の設定で，再分割テストデータにおけるデータストア内のトークン頻度と$k$NN確率およびLM確率の期待値との関係を調査し，\cref{fig:knn_prob_indomain}に示した．
図より，in-domain設定であっても$k$NN言語モデルによる低頻度トークンの予測の悪化が生じることが示された．


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\section{関連研究}
\label{sec:related_work}

$k$NN言語モデルの強みや限界を理解するために，その性能の鍵となる要因が分析されてきた．
\citeA{drozdov-etal-2022-cant}は，データストアの構築に用いるデータと評価データの間に多数の$n$-gramの重複があることが，$k$NN言語モデルの高い性能にとって重要であることを示した．
この知見は，データストアの構築に用いる訓練データとの重複が少ないWikiText-103の再分割評価データにおいて，$k$NN言語モデルの性能が低下するという本研究の結果と一致している．
\citeA{wang-etal-2023-knn}は，全体のPPLではなく個々のトークンの予測確率に注目し，$k$NN言語モデルが過半数のトークンの予測性能を悪化させることを示し，open-endのテキスト生成タスクにおける性能低下を招くことを明らかにした．
本研究では，この予測性能を悪化させるトークンが主に低頻度のもので構成されていることを明らかにした．
\citeA{xu2023nearest}は，様々なアブレーション実験を通じて，明示的な記憶の利用が$k$NN言語モデルの成功において主要因ではないことを示唆しており，本研究の低頻度トークンに関する分析もこの仮説を支持する結果となった．

また，$k$NN言語モデルおよびその派生である$k$NN機械翻訳 \cite{khandelwal2021nearest}の性能を向上させるため，近傍数$k$や補間係数$\lambda$を動的に決定する手法 \cite{zheng-etal-2021-adaptive,jiang-etal-2021-learning,jiang-etal-2022-towards}や，2種類のデータストアを用いて距離を再計算する手法 \cite{bhardwaj-etal-2023-adaptation}が提案されてきた．
しかし，これらの手法はいずれも本研究で示された低頻度トークンの予測性能への悪影響を明示的には対処していない．


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\section{おわりに}
\label{sec:conclusion}

本研究では，低頻度の目的トークンの予測に焦点を当て，$k$NN言語モデルの振る舞いを詳細に分析した．
その結果，$k$NN言語モデルは低頻度の目的トークンに対する予測性能を改善しないことが明らかとなった．
このことは，低頻度トークンに対する近傍探索の困難さに起因しており，データストアにおけるトークンの分布や直積量子化に伴う誤差といった要因によって近傍探索が困難になっていると考えられる．

本研究により，$k$NN言語モデルは高頻度トークンについての尤度の予想性能を改善するものの，低頻度トークンの予測性能は悪化することがわかった．高頻度トークンの予測性能の改善は，その出現頻度から，コーパス全体の平均PPLの改善には寄与するが，実用的には高頻度以外のトークンに関する予測性能が求められるタスクも存在する．例えば，回答部分の尤度によって答えを決定する多肢選択式の質問応答タスク \cite{hendrycks2021measuring,pmlr-v174-pal22a,robinson2023leveraging}などでは，必ずしも高頻度トークンの予測性能を改善したからといって，タスク性能が改善するとは限らない．
本研究の知見に基づくと，$k$NN言語モデルの改善には低頻度トークンに対して頑健な埋め込みおよび検索手法の設計や誤差の少ないデータストア圧縮技術の開発が必要である．
一つのアプローチとして，目的トークンに逆文書頻度 (IDF) 重みを適用することで，情報量の多い低頻度トークンに対する出力確率の調整を行うことが考えられる．
ただし，低頻度トークンの$k$NNヒット率が低いことを考慮に入れると，IDF重みを$k$近傍に適用しても目的の低頻度トークンの確率を向上させることができない．そのため，この方法は近傍数$k$を極めて大きくする，あるいはデータストアに事前にトークンの逆頻度に比例するようなデータ点を外挿することでヒット率を向上させてから適用する必要があるが，このような手法の確立は今後の課題である．
その他には，トークン頻度に基づいて埋め込みを正規化するZipfian whitening \cite{yokoi2024zipfian}を用いることでデータストアにおける頻度バイアスを軽減できる可能性がある．
ただし，$k$NN言語モデルのデータストアのキーベクトルは語彙の要素ごとの静的埋め込みではなく，与えられたテキストの各トークンの文脈化埋め込みであるため，この手法を素朴に適用すると$\mathbb{R}^{|\mathcal{D}| \times D}$の行列に対するSVDを求めることになり，スケーラビリティに課題が残る．
また，これらの手法の適用によって低頻度トークンを重みづけすることで，低頻度トークンの予測確率の向上に寄与することが期待できるが，高頻度のトークンの予測確率の低下を招き全体としてPPLが低下する副作用が生じる可能性がある．しかし，上述のような区間尤度を用いるタスク設定においては，副作用が生じる場合においても実用上有用であると考えられる．

本研究の知見が，$k$NN言語モデルのより効果的な活用に貢献することを期待する．



%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\acknowledgment

本研究はJST次世代研究者挑戦的研究プログラム JPMJSP2140 の支援を受けたものです．



%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\bibliographystyle{jnlpbbl_1.8}
\bibliography{07refs}



%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\begin{biography}

\bioauthor{西田　悠人}{%
2022年豊橋技術科学大学工学部情報・知能工学課程卒業．
2024年奈良先端科学技術大学院大学先端科学技術研究科情報科学領域修士課程修了．
同年より同博士課程に在籍．
機械翻訳や言語モデルの研究に従事．
}

\bioauthor{森下　　睦}{%
2017年奈良先端科学技術大学院大学情報科学研究科博士前期課程修了．
同年よりNTTコミュニケーション科学基礎研究所研究員．
2022年東北大学大学院情報科学研究科博士後期課程修了．博士（情報科学）．
2024年よりフューチャー株式会社シニアアーキテクト．
主に機械翻訳，言語資源構築に関する研究に従事．
}

\bioauthor{出口　祥之}{%
2019年愛媛大学工学部情報工学科卒業．
2021年愛媛大学大学院理工学研究科電子情報工学専攻修士課程修了．
2024年奈良先端科学技術大学院大学先端科学技術研究科情報科学領域博士課程修了．博士（工学）．
NICTにて研究技術員として勤めた後，2024年よりNTTコミュニケーション科学基礎研究所ポスドク．
自然言語処理，機械翻訳の研究に従事．
}

\bioauthor{上垣外英剛}{%
2017年東京工業大学大学院総合理工学研究科知能システム科学専攻博士課程修了．博士（工学）．同年NTTコミュニケーション科学基礎研究所入所．2018年より東京工業大学科学技術創成研究院未来産業技術研究所助教．2022年より奈良先端科学技術大学院大学先端科学技術研究科情報科学領域准教授．自然言語処理分野の研究に従事．情報処理学会，言語処理学会，ACL各会員．
}

\bioauthor{渡辺　太郎}{%
1994年京都大学工学部情報工学科卒業.
1997年京都大学大学院工学研究科情報工学専攻修士課程修了.
2000年Language and Information Technologies, School of Computer Science, Carnegie Mellon University, Master of Science取得.
2004年京都大学博士（情報学）.
ATRおよびNTT, NICTにて研究員, また, グーグルでソフトウェアエンジニアとして勤めた後, 2020年より奈良先端科学技術大学院大学教授.
自然言語処理や機械学習, 言語モデル, 機械翻訳の研究に従事.
}

\end{biography}




%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\biodate





\end{document}



