[論文レビュー] Simplified TinyBERT: Knowledge Distillation for Document Retrieval

Xuanang Chen, Ben He+3 more|arXiv (Cornell University)|Sep 16, 2020

Topic Modeling参考文献 21被引用数 6

ひとこと要約

この論文では、2段階のトレーニングを1段階に統合し、損失関数にハードラベルを組み込むことで文書検索性能を向上させる知識蒸留フレームワーク、Simplified TinyBERTを提案する。MS MARCOおよびTREC 2019 DLベンチマークにおいて、BERT-Baseより15倍高速な推論を実現しながら、両者を著しく上回る性能を達成している。

ABSTRACT

Despite the effectiveness of utilizing the BERT model for document ranking, the high computational cost of such approaches limits their uses. To this end, this paper first empirically investigates the effectiveness of two knowledge distillation models on the document ranking task. In addition, on top of the recently proposed TinyBERT model, two simplifications are proposed. Evaluations on two different and widely-used benchmarks demonstrate that Simplified TinyBERT with the proposed simplifications not only boosts TinyBERT, but also significantly outperforms BERT-Base when providing 15$ imes$ speedup.

研究の動機と目的

• BERTの高い推論コストが展開を制限する文書ランク付けというタスクにおいて、知識蒸留の有効性を調査すること。
• トレーニングプロセスの簡素化とランク付けタスクにおける蒸留効果の向上を目的としたTinyBERTモデルの改善。
• ランク付け性能でBERT-Baseを上回りながら、高い効率性を維持する小型で高速な学生モデルの実現。

提案手法

• 2段階のTinyBERTのタスク固有の蒸留段階を、統合損失関数を用いて1段階のトレーニングステップに統合した。
• 注意、隠れ状態、埋め込み、ソフトラベルの蒸留損失を統合した統一損失関数を提案：$\mathcal{L} = \mathcal{L}_{attn} + \mathcal{L}_{hidn} + \mathcal{L}_{emb} + \mathcal{L}_{soft}$。
• 関連と非関連のドキュメントの区別を向上させるために、ハードラベルを損失関数に組み込んだ。
• 一般蒸留段階を別途必要としない1段階の蒸留プロセスを採用。学生モデルはBERT-Baseの最初のk層で初期化された。
• 固定サイズの学生モデルに対して、温度$T \in \{1,5,10\}$と知識蒸留重み$\alpha \in \{0.2,0.5,0.7\}$のハイパーパrameterサーチを実施した。
• MRR@10やNDCG@10といった標準的なランク付け指標を用いて、MS MARCOおよびTREC 2019 DLトラックベンチマークでモデルを評価した。

実験結果

リサーチクエスチョン

• RQ1 推論コストが高く制限が厳しい文書検索というタスクに、標準的な知識蒸留およびTinyBERTを効果的に適用できるか？
• RQ2 TinyBERTのタスク固有の蒸留の2段階を1段階に統合することで、性能とトレーニング効率が向上するか？
• RQ3 蒸留損失にハードラベルを追加することで、関連と非関連のドキュメントを区別する能力が向上するか？
• RQ4 Simplified TinyBERTの簡素化されたトレーニングプロセスは、15倍の推論高速化を達成しながら、BERT-Baseを上回るランク付け性能を実現するか？
• RQ5 提案された簡素化は、さまざまなモデルサイズおよび再ランク付け深度において、性能を維持または向上させるか？

主な発見

• L6_H768設定のSimplified TinyBERTは、15倍の高速化下でMS MARCO開発セットにおいて、BERT-Baseおよび標準TinyBERTを著しく上回る性能を示した。
• TREC 2019 DLテストセットでは、小規模なL3_H384学生モデルを用いた場合、Simplified TinyBERTはNDCG@10でBERT-Baseを上回ったが、TinyBERTは同じ条件下でBERT-Baseより著しく劣っていた。
• 全再ランク付け深度（トップ10からトップ100）において、Simplified TinyBERTはBERT-BaseおよびTinyBERTを一貫して上回り、MS MARCO開発セットの深度100においてMRR@10で3.5%の向上を達成した。
• 1段階トレーニングによる簡素化により、標準TinyBERTと比較して2段階トレーニング時間が42–52%短縮された。L6_H768では14.37時間（標準TinyBERTは29.95時間）であった。
• アブレーションスタディの結果、1段階トレーニングとハードラベル統合の両方が、独立的かつ相乗的に性能向上に寄与していることが確認された。
• L3_H384バージョンは15倍の高速化を達成しながら、文書検索分野における蒸留モデルの中で最先端の性能を達成しており、BERT-Baseのランク付け効果を上回っている。