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

長谷川拓西田京介加来宗一郎吉田仙
日本電信電話株式会社 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 ステップ分のデータについては閾値の決定にのみ用い，計算グラフは保持せずパラメータの更新は行わないものとした.