# 高速・高精度検索のための勾配推定手法を用いた 擬似スパースベクトルの学習法

長谷川拓西田京介加来宗一郎吉田仙
日本電信電話株式会社 NTTメディアインテリジェンス研究所
taku.hasegawa.ps@hco.ntt.co.jp

## 1 はじめに

深層学習技術の発展により, ニューラルネット ワーク（NN）を用いた検索モデルの有効性について 数多く報告されている [1]. 近年では，BERT [2]を はじめとする事前学習モデルの利用により従来技術 を大きく上回る精度で検索が可能となってきた [3].

しかし，大規模検索においては，[3] の様にクエリ と文書を併せて NNに入力する方式では，検索対象 の文書数分の計算が検索のたびに必要となるため,検索速度の観点から実用的でない. そこで，クエリ と文書を独立に NN で密べクトルにエンコードし て，最大内積探索（MIPS）により検索する方式が近年では盛んに研究されている [4]. MIPS は高速な検索を可能にするが，MIPS のインデックスを構築す るための計算時間が膨大となる課題が残る.

そこで我々は，クエリと文書を独立にスパースベ クトルにエンコードし，転置インデックスを利用可能にすることでインデックスの構築および検索の両方を高速に行うことを目指す. 本研究では，長谷川 らにより提案された $\alpha$-SVS [5] の拡張に取り組む. $\alpha$-SVS ではベクトルの各次元に対し，ミニバッチに おける絶対値の大きさの上位 $100 \alpha \%$ 值が小 さな要素を 0 として得た擬似的なスパースベクト ルを用いて損失を計算するが，離散的な forward 計算（top- $\alpha$ ）の勾配が不安定となる問題があった. 本研究では, 離散 forward 関数の偏微分の線形補間に より擬似的な勾配を推定する手法を導入すること でこの問題を解決する．大規模コーパスである MS MARCO を用いた数値実験により，従来の $\alpha$-SVS に 比べ性能が向上したことを示す.

## 2 定義

本研究で取り組む情報検索と関連する用語につい て定義する.定義 1 (擬似スパースベクトル). $k$ 番目のベクトルを $v^{k}=\left[v_{1}^{k}, \ldots, v_{D}^{k}\right]^{\top} \in \mathbb{R}^{D}$ とする. $v^{k}$ の擬似スパース ベクトルの $i$ 次元目を以下で定義する.

$
f\left(v_{i}^{k} \mid \theta_{i}\right)= \begin{cases}v_{i}^{k} & v_{i}^{k}>\theta_{i} \\ 0 & \text { otherwise }\end{cases}
$

ここで $\theta_{i}$ は $i$ 次元目の間値である.

定義 2 (擬似スパースベクトルを用いた転置イン デックス). 擬似スパースベクトルを用いた転置辞書 を表現べクトルの各次元を潜在単語とみなし，次元番号の集合 $\{1,2, \ldots, D\}$ から $X^{p} \times \mathbb{R}_{>0}$ の部分集合の 族への写像として定義する。ここで，インデックス $i$ の像は文書とその文書とクエリの関連度スコアの 部分集合 $\left.\{\left(x^{\mathrm{p}_{k}}, v_{i}^{k}\right) \mid k=1, \ldots, K, v_{i}^{k} \neq 0\right.\}$ である.

## $3 \alpha$-SVS

高次元のスパースベクトルを作ることにより，特徵べクトルを擬似的に潜在単語とみなして転置イン デックスを作成する手法 [6] がこれまで提案されて きたが，単純にL1ノルムを最小化する方法では， ベクトルが部分空間に偏る問題があった.

長谷川らはこの問題を解決するために top- $\alpha$ 学習 とこれを導入した検索モデル $\alpha$-SVS を提案してい る [5]. 図 1 に $\alpha$-SVS のアーキテクチャを示す.

## 3.1 文脈埋め込み層

文脈埋め达み層の出力は $H=\operatorname{BERT}(X)$ として得 られる.ここで, 入力系列 $X$ は，

$
X=\left([\mathrm{CLS}], x_{1}, x_{2}, \ldots,[\mathrm{SEP}]\right)
$

とし, $\{x\}$ はクエリ $x^{q}$ あるいは文書 $x^{p_{k}}$ のトークナ イズされたサブワード系列である． BERT はクエリ と文書に対して共有される。

図 $1 \alpha$-SVS のアーキテクチャ. 左図：訓練時のモデル。 $\alpha$-擬似スパースベクトルはクエリ，正例，負例のペアワイズで 学習される. 右図 : 検索時のモデル。 $\alpha$-擬似スパースベクトルで作成した転置辞書を用いてクエリと文書の内積で関連度スコアが算出される。

## 3.2 ベクトル出力層

ベクトル出力層は BERT の出力 $H$ の [CLS]トーク ンに対応する $h_{0}$ を入力とし，全結合層 2 層と ReLU で高次元空間に写像した $v \in \mathbb{R}^{D}$ が出力となる. さ らに，出力べクトルの L2ノルムが 1 となるように 正規化を行い，超球面上で最適化を行う。

$
v^{\prime}=\operatorname{ReLU}\left(W_{2} \operatorname{ReLU}\left(W_{1} h_{0}+b_{1}\right)+b_{2}\right), \quad v=\frac{v^{\prime}}{\left.\|v^{\prime}\right.\|_{2}}
$

$W_{2} \in \mathbb{R}^{D \times D_{\text {hid }}}, W_{1} \in \mathbb{R}^{D_{\text {hid }} \times 768}, b_{2} \in \mathbb{R}^{D}$ および $b_{1} \in$

## $3.3 \alpha$-擬似スパースベクトル

$\alpha$-SVS は得られたべクトル集合に対し，各次元の 値の絶対値の大きさ上位 $100 \alpha \%$ で閾値 $\theta_{i}$ （式 1）を 決め，值が小さな要素を 0 とみなし $\alpha$-擬似スパース ベクトルを得る. 数式による定義は付録 4 に記載し た.この $\alpha$-擬似スパースベクトルを用いて損失関数 を計算して学習する方法を top- $\alpha$ 学習と呼ぶ.

また，検索時と学習時の上位 $100 \alpha \%$ の閾値はそ れぞれ異なるべクトル集合から計算され，クエリと 文書でもそれぞれ独立で閾値を計算する。訓練時の $\alpha$ は文書，クエリそれぞれ $\alpha_{\mathrm{p}}, \alpha_{\mathrm{q}}$ と表し, 検索時は それぞれ $\alpha_{\mathrm{p}}^{\prime}, \alpha_{\mathrm{q}}^{\prime}$ と表すこととする.

訓練時の $\alpha$-擬似スパースベクトル 各ミニバッチ での与えられたクエリと文書のベクトル集合を $V_{i}^{\mathrm{q}}$ と $V_{i}^{\mathrm{p}}$ とする。 $\alpha$-擬似スパースベクトルは $\alpha_{\mathrm{p}}, \alpha_{\mathrm{q}}$ を 用いて以下のように計算される.

$
f_{\mathrm{q}}\left(v_{i}^{\mathrm{q}}\right)= \begin{cases}v_{i}^{\mathrm{q}} & v_{i}^{\mathrm{q}}>\operatorname{top}\left(V_{i}^{\mathrm{q}} \mid \alpha_{\mathrm{q}}\right) \\ 0 & \text { otherwise }\end{cases}
$

$
f_{\mathrm{p}}\left(v_{i}^{\mathrm{p}}\right)= \begin{cases}v_{i}^{\mathrm{p}} & v_{i}^{\mathrm{p}}>\operatorname{top}\left(V_{i}^{\mathrm{p}} \mid \alpha_{\mathrm{p}}\right) \\ 0 & \text { otherwise }\end{cases}
$

ここで， $\alpha_{\mathrm{p}}=\alpha_{\mathrm{q}}=1$ の場合，密ベクトルによる学習 と等価となる.

検索時の $\alpha$-擬似スパースベクトル 検索時にお いては検索精度を保証できるような上位 $\alpha_{\mathrm{q}}^{\prime}$ を計算 するのに十分なクエリ数が常に得られるとは限ら ない. ゆえに， $\alpha$-SVS では，クエリについては訓練 データのクエリ集合を用いて $\operatorname{top}\left(V_{i}^{\mathrm{q}} \mid \alpha_{\mathrm{q}}^{\prime}\right)$ を計算し,文書については検索対象となる全ての文書集合から $\operatorname{top}\left(V_{i}^{\mathrm{p}} \mid \alpha_{\mathrm{p}^{\prime}}^{\prime}\right)$ を計算する.

## 3.4 損失関数

$\alpha$-SVS ではペアワイズランキング学習を行う．訓練データとして $\mathscr{D}=\left.\{\left(x_{i}^{\mathrm{q}}, x_{i}^{\mathrm{p}_{1}}, x_{i}^{\mathrm{p}_{2}}\right) \mid i=1, \ldots, N\right.\}$ を 与える. ここで $x_{i}^{\mathrm{p}_{1}}$ はクエリの正例文書, $x_{i}^{\mathrm{p}_{2}}$ は負例文書である. $i$ 番目の訓練サンプルに対するヒンジ ロス $L_{i}$ は $\alpha$-擬似スパースベクトルで計算した内積 を用いて以下で定義される。

$
L_{i}=\max \left.\{0, \epsilon-\left(f_{\mathrm{q}}\left(v_{i}^{\mathrm{q}}\right) \cdot f_{\mathrm{p}}\left(v_{i}^{\mathrm{p}_{1}}\right)-f_{\mathrm{q}}\left(v_{i}^{\mathrm{q}}\right) \cdot f_{\mathrm{p}}\left(v_{i}^{\mathrm{p}_{2}}\right)\right)\right.\}
$

ここで $\epsilon$ はヒンジロスのマージンを決定するハイ パーパラメータである.

## 3.5 単一モデルでの 2 段階検索

$\alpha$-擬似スパースベクトルは大規模検索で有用であ るが，文書数が膨大な場合においては非常に小さい $\alpha$ を設定する必要がある．この場合 $v$ で得られる関
連度スコアに対しての近似精度は落ちてしまうた め， $\alpha$-SVS では検索時に同じモデルを用いて 2 段階 に検索を行う。 1 段階目は $\alpha$-擬似スパースベクトル の内積で計算した関連度が上位 $k$ 件の文書を選ぶ. 2 段回目では $\alpha$-擬似スパースベクトルを作成する 前の元のべクトル $v$ を用いて内積で関連度スコアを 計算し, 最終的な関連文書のランキングを得る. 従来では 2 段階検索と呼ばれる場合 2 つの検索モデル を用いて行われるが， $\alpha$-SVS では単一のモデルで 2 段階検索を行うため，2 段回目の検索時ではクエリ や文書をモデルに入力する必要がなく高速に計算可能である.

## 4 提案手法

前節で述べた $\alpha$-SVS の学習をより精度良くする ために，本論文では擬似勾配推定手法を導入する。擬似勾配推定手法はモデルの軽量化のためにモデル の低ビット化の際によく用いられる手法であり [7],離散的な forward 計算が含まれるネットワークの学習を効果的に行うために用いられる。

## 4.1 top- $\alpha$ 学習中における勾配

top- $\alpha$ 学習の forward 計算のうち, 式 1 で定義され ている関数 $f$ は偏微分を行うと

$
\frac{\partial f}{\partial v_{i}^{k}}\left(v_{i}^{k} \mid \theta_{i}\right)= \begin{cases}1 & v_{i}^{k}>\theta_{i} \\ 0 & \text { otherwise }\end{cases}
$

となる. ただし $\theta_{i}=t o p\left(V_{i} \mid \alpha\right)$ である.

top- $\alpha$ 学習では, ミニバッチ内の他の要素次第で 同じ文書べクトルであっても top- $\alpha$ に含まれる場合 と含まれない場合が存在しうる. 式 7 で誤差逆伝搬法を用いて学習を行った場合、top- $\alpha$ に含まれな かった次元は誤差が伝搬されず $V_{i}$ の選ばれ方に大 きく影響を受ける. モデルの初期値にも依存するた め，この不確定性は学習過程に影響を与えうる.

## 4.2 top- $\alpha$ 学習中の勾配推定

本論文では偏微分した関数を線形補完すること により，ミニバッチで選ばれる $V_{i}$ の影響を軽減す る手法を提案する. 式 7 で示した top- $\alpha$ 学習での backward で用いられる関数を，以下の式で置き換え ることにより勾配推定手法を導入する。

$
\frac{\partial g}{\partial v_{i}^{k}}\left(v_{i}^{k} \mid \theta_{i}, \theta_{i}^{\prime}\right)= \begin{cases}1 & v_{i}^{k}>\theta_{i} \\ \frac{1}{\theta_{i}-\theta_{i}^{\prime}} v_{i}^{k}-\frac{\theta_{i}^{\prime}}{\theta_{i}-\theta_{i}^{\prime}} & \theta_{i}^{\prime}<v_{i}^{k} \leq \theta_{i} \\ 0 & v_{i}^{k} \leq \theta_{i}^{\prime}\end{cases}
$

図 2 勾配推定の backward 関数

ここで， $\theta_{i}^{\prime}$ は $\theta_{i}^{\prime} \leq \theta_{i}$ を満たす. 図 2 に従来手法お よび提案手法で用いられる関数を示す.

ここで注意すべきは $\theta_{i}=t o p\left(V_{i} \mid \alpha\right)$ であり,$V_{i}$ が 与えられてから決定されるため, 毎回異なる値をと る. $\theta_{i}^{\prime}$ も同様に決定される必要があるため, 本研究 では下記の 2 通りの方法で $\theta_{i}^{\prime}$ を決定した。
1. $V_{i}$ の最大値を用いた閾値.

$
\theta_{i}^{\prime}=2 \operatorname{top}\left(V_{i} \mid \alpha\right)-\max V_{i}
$

2. top- $k \alpha$ を用いた閾値.

$
\theta_{i}^{\prime}=\operatorname{top}\left(V_{i} \mid k \alpha\right)
$

以後，本論文では 1 を用いたモデルを $\alpha-\mathrm{SVS}_{\text {max }}$ ， 2 を用いたモデルを $\alpha-\mathrm{SVS}_{\mathrm{k} \alpha}$ とする.

## 5 評価実験

## 5.1 実験設定

MS MARCO 2.1 [9] の Passage and Document Retrieval タスクにより実験を行った。データの詳細お よび学習設定については付録 A. 2 に記載した.

比較対象として，クエリから答えとなる文書 を事前学習モデル T5 [10]を用いて予測し転置辞書によりキーワード検索を行う手法である docTTTTTquery [8] およびパッセージおよび文書レ ベルの単語重みをニューラルネットを用いて計算し，これらを利用しキーワード検索を行う手法 DeepCT [11,12]を用いた。

## 5.2 勾配推定手法の効果に関する評価

$\alpha$-SVS および勾配推定手法を導入した $\alpha$-SVS のラ ンキング精度および検索速度について評価し比較 した.

表 1 に実験結果を示す. MSMARCO のリーダー ボードでの評価指標として採用されているMRR @
表 1 MSMARCO パッセージ検索タスクにおける検索精度の結果. 比較手法の結果は文献 [8] から引用.

 &  &  &  \\

図 32 種類の勾配推定手法の 2 つの閾値とその差の推移.

10 では提案手法の $\alpha-\mathrm{SVS}_{\max }$ が最も良い結果となっ た. ただし docTTTTTquery や DeepCT と比べ Recall は低く，R@1000についてはBM25にも劣っている。 このことから $\alpha$-SVS は関連文書を大量の文書から 上位にランキングする能力には長けているが，一部検索漏れを起こしうるということがわかった．特に この問題は勾配推定手法を用いたとしても大きくは 改善されず，さらに勾配推定手法を導入した $\alpha$-SVS は従来の手法にくらべて僅かながら recall が下がっ ている箇所が見られる.よって勾配推定手法を導入した場合，上位のランキング精度は向上するが， Recallへの貢献はみられず，劣化する可能性もある ことがわかった. この問題への取り組みは今後の課題である.

一方で, docTTTTTquery で得られた結果と $\alpha$-SVS で得られた結果を半分ずつ合わせた結果である docTTTTTquery $+\alpha$-SVS の R @ 100 および $\mathrm{R} @ 1000$ の 結果は docTTTTTquery と比較しても最も良い結果 であった. この結果から， $\alpha$-SVS は他の手法では検索漏れを起こしうるパッセージを検索することが 可能であることがわかり， $\alpha$-SVS の有用性が示さ れた。

## 5.32 つの勾配推定手法の比較

2つの勾配推定の手法を比較するために学習中の $\theta$ および $\theta^{\prime}$ の推移を調べた. 図 3 に結果を示す.こ の結果から $\alpha-\mathrm{SVS}_{\text {max }}$ は $\alpha-\mathrm{SVS}_{2 \alpha}$ に比べ各次元の値 の上位 $100 \alpha \%$ の間値が高くなりやすい傾向があり，学習が進むにつれ線形補完される区間が大きくなり やすい傾向があることがわかった. これはあるべク トルは特定の次元に大きい值を持ちやすいことを示 しており， $\alpha$-SVS の出力は L2ノルムが 1 であるこ とから， $\alpha-\mathrm{SVS}_{\max }$ の方がよりスパース性の高いべ クトルが得られやすい傾向にあることがわかった.

## 6 おわりに

本研究では大規模検索モデル $\alpha$-SVS の学習をより 促進するための勾配推定手法を提案し，MS MARCO 2.1 の passage ranking タスクにおいて従来手法を上回る性能を達成した。

従来の $\alpha$-SVS ではランダムに選ばれるバッチ内で の値の上位で閾値決まるため，学習の安定性に欠け るという問題があった. 本研究では，勾配推定手法 を導入することによりバッチによる閾値の摇らぎを 考慮してより効率の良い学習を実現した。

## 参考文献

[1] Bhaskar Mitra and Nick Craswell. An introduction to neural information retrieval. Foundations and Trends in Information Retrieval, Vol. 13, No. 1, pp. 1-126, 2018.

[2] Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: pre-training of deep bidirectional transformers for language understanding. In NAACL, pp. 4171-4186, 2019.

[3] Rodrigo Nogueira and Kyunghyun Cho. Passage re-ranking with BERT. 2019. arXiv:1901.04085.

[4] Vladimir Karpukhin, Barlas Oguz, Sewon Min, Patrick S. H. Lewis, Ledell Wu, Sergey Edunov, Danqi Chen, and Wen-tau Yih. Dense passage retrieval for open-domain question answering. In EMNLP, pp. 6769-6781, 2020.

[5] 長谷川拓, 西田京介, 加来宗一郎, 富田準二. 高速な情報検索に向けた文脈考慮型スパース文書ベクトルの獲得. 人工知能学会全国大会論文集, Vol. 2020, pp. 4Rin145-4Rin145, 2020.

[6] Hamed Zamani, Mostafa Dehghani, W. Bruce Croft, Erik G. Learned-Miller, and Jaap Kamps. From neural re-ranking to neural ranking: Learning a sparse representation for inverted indexing. In CIKM, pp. 497-506, 2018.

[7] Zechun Liu, Baoyuan Wu, Wenhan Luo, Xin Yang, Wei Liu, and Kwang-Ting Cheng. Bi-real net: Enhancing the performance of 1-bit cnns with improved representational capability and advanced training algorithm. In Proceedings of the European Conference on Computer Vision (ECCV), September 2018.

[8] Rodrigo Nogueira and Jimmy Lin. From doc2query to docttttquery. https: //cs. uwaterloo.ca/ jimmylin/publications/ Nogueira_Lin_2019_docTTTTTquery-v2.pdf, 2019.

[9] Payal Bajaj, Daniel Campos, Nick Craswell, Li Deng, Jianfeng Gao, Xiaodong Liu, Rangan Majumder, Andrew McNamara, Bhaskar Mitra, Tri Nguyen, Mir Rosenberg, Xia Song, Alina Stoica, Saurabh Tiwary, and Tong Wang. MS MARCO: A human generated machine reading comprehension dataset. 2018. arXiv:1611.09268v3.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. CoRR, Vol. abs/1910.10683, , 2019.

[11] Zhuyun Dai and Jamie Callan. Context-aware sentence/passage term importance estimation for first stage retrieval. CoRR, Vol. abs/1910.10687, , 2019.

[12] Zhuyun Dai and Jamie Callan. An evaluation of weakly-supervised deepct in the trec 2019 deep learning track. In TREC, 2019.

[13] Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In ICLR, 2015.

## A 付録

## A. 1 定義

定義 3 ( 2 段階検索タスク). クエリ $x^{\mathrm{q}}$ と $K$ 個の文書 $X^{\mathrm{P}}=\left.\{x^{\mathrm{p}_{k}} \mid k=1, \ldots, K\right.\}$ が入力として与えられ、初 めにモデルは $X^{\mathrm{p}}$ からクエリと関連する $T_{1}$ 個の文書 の順序集合を返す. 次にモデルは先の $T_{1}$ 件の文書 をリランキングし， $T_{2}$ 個のクエリに関連する文書の 順序集合を返す。ただし，順序集合とはクエリと文書の関連度によって順序づけられた集合である。

定義 4 ( $\alpha$-擬似スパースベクトル). ベクトル集合の $i$ 番目の要素の集合 $V_{i}$ を $V_{i}=\left.\{v_{i}^{k} \mid k=1, \ldots,\right.\}$ とす る. $\alpha$-擬似スパースベクトル $v^{k}$ の $i$ 次元目は

$
f^{\prime}\left(v_{i}^{k}, V_{i} \mid \alpha\right)= \begin{cases}v_{i}^{k} & v_{i}^{k}>\operatorname{top}\left(V_{i} \mid \alpha\right) \\ 0 & (\text { otherwise })\end{cases}
$

ここで，top $(\cdot \mid \alpha)$ は与えられた集合の上位 $100 \alpha \%$ の 閾値を返す関数である.

## A. 2 実験設定の詳細

## データセット.MS MARCO 2.1 [9] の Passage and Document Retrieval タスクにより実験を行った.本タスクでは ReRanking と Full Ranking が Passage 単位，Document 単位それぞれで存在しており，本実験では Passage Rankingを用いて評価を行った.

ReRanking タスクではあらかじめ BM25 を用いて絞 り込まれた上位 1000 件の passage が与えられている が，Full Ranking では約 880 万件の検索対象文書か ら検索を行うことが求められる. ランキング指標と しては Mean Reciprocal Rank (MRR) Craswell09a が 用いられている。

学習用設定. Train Triples Small セットを用い て，バッチサイズ 100，エポック数 1 とし， 4 枚の NVIDIA Tesla V100 GPU により学習した. 最適化に は Adam [13] を用い， $\beta_{1}=0.9, \beta_{2}=0.999, \epsilon=10^{-8}$ とした. モデルパラメータは 0 で初期化したバ イアスを除き， $N(0,0.02)$ で初期化した. 学習率は $1 \times 10^{-6}$ とし, 学習の半分までは 0 から線形で増加 させ，最終ステップで $9 \times 10^{-7}$ になるように線形に 減衰させた. 勾配は最大ノルム 1 でクリッピングし た. BERT は base モデル（768 次元）を利用し， 2 層 の出力層の中間層の次元数は 1000 , 最終層（スパー スベクトル v）の次元数 $D$ は 30000 とした. hinge loss のマージン $\epsilon$ は 1.0 とした. トークナイズには表 2 MS MARCO 2.1 の passage ranking のデータの詳細.

BERT Wordpiece トークナイザ1）（語彙数 $30 \mathrm{~K}$ ）を用 いた.

top $\alpha$ 学習のためのハイパーパラメータとして $\alpha_{\mathrm{q}}=0.1 \%$ とした. また，バッチサイズ 100 とした 場合，バッチ内で上位 $\alpha \%$ を計算すると 閾值がバッ チ内のデータに強く依存しすぎるため，計算を安定 させるために現在のステップを含め過去 20 ステッ プ分の訓練データのベクトルを保存しておき，これ らを合わせて用いることで top $\alpha$ の閾値を安定させ た. ここで，現在のステップを除く 19 ステップ分 のデータについては閾値の決定にのみ用い，計算グ ラフは保持せずパラメータの更新は行わないものと した.

 長谷川拓西田京介加来宗一郎吉田仙 日本電信電話株式会社 NTTメディアインテリジェンス研究所 taku.hasegawa.ps@hco.ntt.co.jp  