TransformersにUltraBERTのネイティブ統合を提案、二重マスク学習で分類と検索に対応

Transformersには10月14日、UltraBERTのネイティブ統合提案が提出されました。同日、研究チームのDMIS Labはモデルの紹介記事を公開し、約1億5,000万および3億8,000万パラメーターの英語エンコーダーと、そのまま利用できる埋め込みモデルを発表しました。このシリーズはApache 2.0ライセンスで提供され、32Kのコンテキストに対応します。今回の新たな動きはモデル紹介記事の公開とフレームワーク統合の提案で、関連論文は10月10日に投稿されています。紹介記事、論文

UltraBERTの主な変更点は事前学習の目的関数です。著者らの実験では、マスク率を低くすると分類に有利で、高くすると検索に有利でした。そこで同じバッチ内に、マスク率5%と50%の系列を4対1の比率で混在させ、エンコーダーを共有しながら、言語モデル出力ヘッドはそれぞれ独立させています。マスク率の低い分岐には13億パラメーターの教師モデルによる蒸留を適用し、高い分岐では通常のマスク言語モデリングを維持することで、蒸留による検索性能の低下を避けています。学習方法

アーキテクチャでは、局所注意と全域注意を交互に配置し、それぞれ異なるRoPE基底を使用します。そのため、ネイティブ統合提案では複数の回転位置埋め込みを管理し、モデルコードをTransformersの注意機構バックエンドのインターフェースに適合させる必要があります。提案者はeager、SDPA、FlashAttention 2の出力比較も添付しています。記載されたテストケースでは、有効なトークンの隠れ状態が元のリモートコードと一致しました。ただし、これは限られたテストケースであり、あらゆる精度や系列長で検証済みという意味ではありません。統合提案

著者らによると、380MモデルはGLUEで91.0、BEIRで47.7を達成しましたが、いずれも各タスクでファインチューニングした後の結果です。ベースチェックポイントは分類または検索向けの学習が必要で、そのままデプロイできる汎用ベクトルモデルと見なすことはできません。エンジニアリングチームにとって検証する価値があるのは、同じ小型バックボーンで文書分類と検索を支え、複数モデルの保守コストを下げられるかどうかです。モデルカード

確認時点では、ネイティブ統合はまだレビュー中で、既存のモデルカードではtrust_remote_code=Trueの設定が必要です。今後はマージと正式リリースへの収録を追い、自社のコーパスで長文処理時のレイテンシ、メモリー使用量、検索品質を測定する必要があります。モデルは英語データで学習されており、現時点のスコアから日本語での性能を推測することはできません。提案の状況、利用上の制限