Improving Bi-encoder Document Ranking Models with Two Rankers and Multi-teacher Distillation
2021.09.09
著者: Jaekeol Choi(ソウル大学校, NAVERの人), Euna Jung, Jangwon Suh, Euna Jung, Jangwon Suh (ソウル大学校)
選んだ理由
今のPretrainedモデルの事情を知りたかったことと、よりよい文表現を既存モデルの蒸留で作れそうだったため。
概要
既存の強い文書ランキングモデル2つを合わせて蒸留することで新しいモデルを作るフレームワーク ”TRMD; two rankers and multi-teacher distillation”を提案
構造の違う2つのモデル(bi-encoder, cross-encoder)を同時に教師として学習に使うことでお互いのいいとこ取りができる。
イメージ:全く戦闘スタイルの違う2人の師匠から同時に見様見真似で学んでいく
結果、教師モデル単体よりも高い性能を示した。
two ranker構造とmulti-teacher distillation
知識の蒸留(Knowledge distillation)とは:
ある既存のモデル(教師モデル)と同じ入力を別のモデル(生徒モデル)にした場合に、教師モデルの出力と生徒モデルの出力ができるだけ同じになるように学習していく。
https://gyazo.com/49dad4bdaef7fb0d10085addca679fec
中間層同士を似せていくパターンもある
サイズの大きなモデルの軽量化などに使われたりする
精度向上やロバスト性向上も期待できる(知らなかった)
提案モデル: 2教師、2ranker
https://gyazo.com/0195a59ac61cd0e2730fad2d0243a8a1
異なる2つのモデルで知識蒸留を行い、student encoderはcross-encoder teacherの表現(CLS表現)とbi-encoder teacherの表現(ColBERTのCLS, query/docment表現)の両方を学習する。
教師モデル
モデル構造、出力される分散表現が異なる2モデルを教師にする
https://gyazo.com/b6154db95915fe4972827eb0961cc80e
モデル構造(表2列目)
Cross-encoder: query-document対全体にself-attentionを行う
クエリと文書間の相互関係を取れる(クエリに応じた文表現ができる)
一般にbi-encoderより性能が高め
bi-encoder: queryはqueryで、documentはdocumentでself-attentionを通す。
文表現がクエリによって変わることがないので文表現を事前計算しておける。のでクエリあたりの推論時間がかなり短い
https://gyazo.com/e2f66701352b35fb26ed1b91383c8f24
生徒モデル
学習の流れ
1. 各教師モデルをそれぞれfine-tuneしておく
2. 各教師モデル、生徒モデルにクエリと文書を入力し、文表現を出す
3. ロスの計算
生徒モデルの推定関連度 と 正解の関連度の差
https://gyazo.com/657cc0d7f5134615a181c69b5ce05acd
zは関連度スコアのあつまり
生徒モデルの文表現(cross用) と cross教師の文表現の差
https://gyazo.com/c999a575a751247dcfef8950c2411265
生徒モデルの文表現(bi用) と bi教師の文表現の差
https://gyazo.com/6e6a530b5c13f66a2c5d996f3a984003
この3つのロスの総和を最小化するように学習していく。
▼2人教師を使って、bi-encoderモデルを作る場合 (図左からcross教師, bi生徒, bi教師)
https://gyazo.com/0195a59ac61cd0e2730fad2d0243a8a1
bi教師から教わる時:
生徒モデルと教師モデルの{Query:Document}embedding(2つCLS vectorがある)を比較する
cross教師から教わる時:
生徒モデルのCLS embedding(真ん中グレーの箱)とcross教師のCLS embeddingを比較する
生徒と教師の型が異なる問題(cross教師, bi生徒)
モデルが違うので出力が異なる...→うまく統合する
今回はbi生徒が2つCLS vectorを出力するので、それをresidual functionに通して1つの分散表現にまとめる(片方を線形変換してもう一方に足す, TwinBERTでやってた)
https://gyazo.com/500fbe2b9657e7964e0effa19f6bf577
↓
両教師モデルで関連度スコアを出力(もとの手法(ranker)そのまま利用)
両関連度スコアを足し合わせたスコアを生徒モデルの関連度スコアとする
https://gyazo.com/5a948945148a54c04fcdbe9463bcbd45
生徒モデルの関連度スコアとデータセットの正解を比較(式3)
実験
データセット: 文書ランキング系(Robust04, Clueweb09b)
評価指標: P@20, nDCG@5, nDCG@20
比較モデル:
教師モデル単体(monoBERT, TwinBERT, ColBERT)
提案モデル(TRMD-〇〇)
提案モデルの蒸留なし(TR-〇〇; Two Rankers)
結果:教師モデルに勝った。
https://gyazo.com/6443381fb9b9a22739c446010cd6afad
一般にcross-encoder がbi-encoderより強いが...
bi-encoderモデルに2教師で蒸留してもcross-encoderモデル単体には及ばない
cross-encoderが生徒のときよりもbi-enocoderが生徒のときのほうが精度の上がり幅は大きい(自分より強い教師から教わるから, biでは使えなかったクエリ・文書間の関係を教師モデルから得た)
蒸留によって精度が上がることも確認できた。(単なるアンサンブルよりも伸びる)
感想、所感
アンサンブルをレベルアップ+軽量化したような提案で汎用性が高そう
技術としてものすごく目新しいわけではないが、既存の枠組みをうまく利用しているところが評価された?
イントロで「精度の高いcross-enocderは計算時間が~~」といっていたのに実験で計算時間に触れていなかったのはなぜ、(多分短くはなっているはず)