知っておくべきこと：マスク言語モデルは驚くほど優れたゼロショット学習者である

この投稿へようこそ！「TIL」として、これは意図的に小さくまとめられたブログ記事であり、重要な詳細のみを含んでいます。さらに詳しく知りたい場合は、技術レポートをご覧くださいか、HuggingFace でモデルを遊んでみてください！

TL;DR

従来（もちろん例外はありますが）、BERT などのエンコーダーモデルは、コアとなるエンコーダーモデルの上にタスク固有のヘッドを搭載して使用されてきました。機能的には、これは事前学習時に使用されるマスク言語モデリングヘッド（MLM ヘッド）に蓄積された言語モデリングの良さをすべて捨てて、バックボーンを再利用してさまざまなタスクを実行しようとすることを意味します。

このアプローチは非常にうまく機能しています：これが支配的なパラダイムとなっている理由があります！しかし、もし生成ヘッド自体がゼロショットでもほとんどのタスクを実行できるのであればどうでしょうか？私たちが試したのはまさにこれであり、かなり良い結果が出ました！私たちは「ModernBERT-Large-Instruct」を導入しました。これは ModernBERT-Large をベースに、驚くほどシンプルなメカニズムでインストラクションチューニングされたエンコーダーです。このモデルは、タスク固有のヘッドの代わりに ModernBERT の MLM ヘッドを使用して、分類や多肢選択タスクを実行することができます。以前の手法とは異なり、私たちの方法ではアーキテクチャの変更も複雑なパイプラインも必要とせず、さまざまなタスクで強力な結果を達成しています。

知識 QA タスクにおいて、エンコーダーが通常苦手とする分野で驚くほど高い能力を発揮します：MMLU-Pro リーダーボードでは、Qwen2.5-0.5B や SmolLM2-360M といった 10 億パラメータ未満のすべてのモデルを上回り、Llama3-1B（はるかに多くのトークンでトレーニングされ、かつパラメータ数が 3 倍）にも非常に近い性能を達成しています！

NLU タスクでは、ModernBERT-Instruct を同じデータセットでファインチューニングした場合、従来の分類ヘッドと同等かそれ以上の性能を発揮します。

これらの結果は、非常にシンプルでエキサイティングなトレーニングレシピによって達成されました。今後さらに改善の余地が十分にあります👀👀

私もぜひ試してみたいです！

このモデルは HuggingFace で ModernBERT-Large-Instruct として利用可能です。カスタムアテンションマスクやそれに類するものを必要としないため、ゼロショットパイプラインの設定と使用は非常に簡単です：

ModernBERT-Large-Instruct の使い方を見るにはこちらをクリック -->

import torch

from transformers import AutoTokenizer, AutoModelForMaskedLM

モデルとトークナイザーの読み込み

model_name = "answerdotai/ModernBERT-Large-Instruct"

tokenizer = AutoTokenizer.from_pretrained(model_name)

device = 'cuda' if torch.cuda.is_available() else 'cpu'

if device == 'cuda':

model = AutoModelForMaskedLM.from_pretrained(model_name, attn_implementation="flash_attention_2")

else:

model = AutoModelForMaskedLM.from_pretrained(model_name)

model.to(device)

分類または多肢選択のための入力を整形します。これは MMLU からランダムに選ばれた例です。

text = """あなたは質問と選択肢を与えられます。正しい答えを選択してください。

QUESTION: (G, .) が群であり、すべての a, b ∈ G に対して (ab)^-1 = a^-1b^-1 が成り立つとき、G は次のどれか

CHOICES:

• A: 可換半群 (commutative semi group)
• B: アベル群 (abelian group)
• C: 非アベル群 (non-abelian group)
• D: これらのいずれでもない

ANSWER: [unused0] [MASK]"""

予測を取得する

inputs = tokenizer(text, return_tensors="pt").to(device)

outputs = model(**inputs)

mask_idx = (inputs.input_ids == tokenizer.mask_token_id).nonzero()[0, 1]

pred_id = outputs.logits[0, mask_idx].argmax()

answer = tokenizer.decode(pred_id)

print(f"Predicted answer: {answer}") # 出力結果：B

さらに詳しく知りたい場合は、モデルのファインチューニング方法に関する例が掲載されている当社のミニクックブック GitHub リポジトリをご覧ください！

-->

導入

エンコーダーモデルは従来、タスク固有のヘッドを備えたすべてのタスクで最も高い性能を発揮してきました。必ずしも問題ではないものの、これは少し無駄に感じられます：MLM ヘッド（その元の事前学習時のヘッド）が完全に捨て去られてしまうからです。実際にはこのアプローチは機能しますが、何か見落としているような気もします。さらに、これによりゼロショット能力に対して大きな制限が生じます：タスク固有のヘッドが通常必要とされるため、これを回避して良好なゼロショット性能を得るためにさまざまな工夫が必要となってきました。

MLM エンコーダーの下流利用における簡略化された不完全な歴史

エンコーダーモデルを用いたゼロショット分類は長年にわたり様々なアプローチが試されてきた活発な研究分野です。最も一般的な手法は、テキストの包含関係（entailment）を流用するもので、MNLI などのタスクで学習させた後、与えられたラベルが入力テキストによって論理的に導かれるか（entailed か）を予測します。tasksource/ModernBERT-large-nli といった大規模な TaskSource データセット上で訓練された非常に強力なモデルも存在します。

これは生成型 BERT をマルチタスク学習者として探求した最初の研究ではありません：プロンプト手法や、パターン利用トレーニング（PET）法によるサンプル効率的な学習、あるいはモデルを自己回帰的にする試みなど、いくつかの先行研究があります。また、UniMC のように、意味的に中立な語彙化子（verbalizers）（例えば「A」「B」などの意味のある単語ではないもの）を用いてタスクを多肢選択形式に変換し、カスタムアテンションマスクを採用する手法など、 ours と非常に類似したアプローチも存在します。

しかし、これらのすべての方法には欠点があります。一部の手法は脆く（特に語彙化子の違いに敏感）、あるいは有望ではあるがまだ完全ではない性能しか達成できません。一方、他の手法は非常に良い結果を達成する一方で、大幅な複雑さを付加しています。その一方で、デコーダーの世界（あるいはお望みなら LLM トピア）では、インストラクションチューニングが極めて急速に進化しており、大規模で恐ろしいとされる大規模言語モデル（LLM）は、インストラクション学習のおかげで生成型分類、特にゼロショットにおいて非常に優れた能力を発揮するようになっています。

しかし、これにも欠点があります：小規模な大規模言語モデル（LLM）は、エンコーダーに通常より劣っており、微調整を施せば大型モデルと同等の性能を発揮することさえあります。さらに、 autoregressive 型の LLM を実行する計算コストは、小規模なモデルであっても、単一の順方向パスでタスクを実行するエンコーダーと比較して一般的に大幅に大きくなります。

ModernBERT-Large-Instruct

私たちのアプローチは、おそらく、あるいはたまたま、両方の利益を得られるかもしれないことを示すものです：もし MLM が、パイプラインやアーキテクチャの複雑さを追加することなく、単一の順方向パスで生成方式によりタスク（ゼロショットのものさえも！）を処理でき、さらに容易に微調整してドメイン内での性能を向上させることができるならどうでしょうか？

これがここで示す可能性です！私たちは非常にシンプルなトレーニングレシピを使用します。ModernBERT の MLM ヘッドを用いた FLAN スタイルの指示微調整です。カスタムアテンションマスクは使用せず、複雑なプロンプトエンジニアリングも、重厚なデータ前処理パイプラインもありません：単に FLAN から単一トークンで回答可能なタスクのみをフィルタリングし、下流評価に使用するデータセットの一部の例を除外するだけです。

仕組みについて

全体プロセスの高レベルな概要

私たちの核心的な洞察は二重です：ModernBERT は単一のヘッドで、ゼロショットまたは完全微調整のいずれにおいても、ほとんどの自然言語理解（NLU）タスクを実行でき、この挙動は極めてシンプルなトレーニングレシピによって解放され、非常に強い可能性を示唆しています。

仕組みは非常にシンプルです:

すべてのタスクは、モデルが単一のトークンで回答できる形式に整形されます。この単一トークンは入力における最終トークンでもあります。これは常にアンカートークン（[unused0]）で前置され、次のトークンが単一のトークンによる回答であることをモデルに伝えます。

モデルには質問、短い指示、および候補となる選択肢のリストが与えられます。すべての選択肢は、単一トークンの語彙化子（verbalizer）で前置されます。これは、モデルがこのラベルを割り当てた場合に予測するトークンです。

その後、モデルは回答として最も可能性の高いトークンを予測し、スコアが最も高い候補となる語彙化子が回答として選択されます。

このアプローチにはいくつかの利点があります:

• 学習または推論のためにアーキテクチャの変更は不要です。
• マスク言語モデリング（Masked Language Modeling）をサポートするあらゆるモデルで、そのまま試すことができます。
• 実験を開始するために必要なデータの前処理は非常に少なくて済みます。
• 同様に、プロンプトエンジニアリングを大幅に削減します: タスクを実行するには、非常に短いテンプレートとすべてのラベルの説明を書くだけで十分です。

トレーニングの詳細

前述の通り、トレーニングレシピは意図的にシンプルに保たれています。これは主にスコープクリープ（scope creep）を防ぐためです。より優れた処理パイプラインや、より現代的な指示セットを使用することで多くの潜在的な改善点を探求できますが、これらすべてを単一トークンタスクに変換するには複雑なプロセスが必要となります。

データ：単一トークンの回答のみを保持するために、20M サンプルにダウンサンプリングされ、フィルタリングされた FLAN-2022 データセット。非常にシンプルなフィルタリングプロセスは、潜在的な回答をトークン化し、回答が 1 つ以上のトークンを含む例をすべて除外することです。評価データセットからの例も、過学習を防ぐために除外されました。

目的：私たちは、答えを含む語彙表（verbalizer）であるべき単一のマスクされたトークンを予測する Answer Token Prediction (ATP) 目的を使用します。最終的なトレーニング目的は、80% の ATP と 20% のダミー MLM 例の混合であり、ここでマスクされたトークンには無意味なラベルが与えられます（後述）。

ベースモデル：LightOn 社やその他のパートナーと共に最近紹介した ModernBERT-Large (395M パラメータ) です。これは、他の代替案よりもはるかに能力の高いベースモデルであることが証明されました。

ダミー例

モデルのトレーニング時に、Answer Token Prediction は壊滅的な忘却（catastrophic forgetting）を引き起こし、モデルが特定のトークンの予測のみを学習して全体的な推論能力を失う可能性があると考えました。これを防ぐために、トレーニング目的の混合を導入しました。具体的には、20% の例に通常の MLM 目的（テキスト内の 30% のトークンをランダムにマスクし、モデルがそれらすべてを同時に予測する）を割り当て、残りの 80% に Answer Token Prediction 目的を採用します。

ただし、私たちはこれを誤って実装してしまい、結果的にこれらのサンプルを空の例として扱い、「ダミー MLM 例」と呼ぶことになりました。問題はラベリングにあり、[MASK] トークンに適切なラベルが割り当てられるべきところを、すべて [MASK] をそのラベルとして与えてしまったのです。これにより、テキスト内に複数の [MASK] トークンが存在する場合、モデルはすぐにそれらすべてに対して [MASK] を予測するよう学習し、これらの例における損失は瞬く間にほぼゼロまで低下しました。

ふむ、単純なミスで、簡単に修正できるはずです、对吧？その通りです。ただし、私たちは予期せぬ現象を観察しました：3 つの事前トレーニング設定（100% ATP、80%ATP/20%MLM、80%ATP/20%dummy）を評価したところ、ダミー例バリアントが最も優れたパフォーマンスを示し、その差は明確でした！この現象について十分に深く探求して説明できる段階には至っていませんが、私の個人的な理論では、これはドロップアウト（dropout）に似た正則化の一種として機能していると考えられます。

性能

ゼロショット結果

ゼロショットの結果は非常に励みになり、ある意味で驚くべきものです！

最良のものとの競合（2B 未満のモデル向け MMLU-Pro リーダーボード）

知識ベースの多肢選択問題 (MMLU および MMLU-Pro): ModernBERT-Large-Instruct は MMLU で 43.06% の精度を記録し、同サイズのモデルである SmoLLM2-360M(35.8%) を上回り、Llama3-1B(45.83%) に迫る結果となりました。MMLU-Pro では、そのパフォーマンスがリーダーボード上で非常に高い位置を占めるものであり、自身のサイズクラスをはるかに超える活躍を見せ、より大きな大規模言語モデル (LLM) と競合するレベルです。

分類タスク: 平均的には、この手法はすべての既存のゼロショット手法を上回ります。ただし、データセットごとの詳細で見ると必ずしもそうではありません。この方法には強い可能性があり全体的に非常に良い結果を出していますが、一部のデータセットでは性能が低下し、他のデータセットでは過剰な評価を示すケースもあります。これは、今後の手法開発において大きなポテンシャルがあることを示唆しています。

ファインチューニング後の結果

MLM ヘッドだけで十分である

トピック分類、テキストの含意 (MNLI)、感情分析など多様なタスクにおいて、各タスクに対して ModernBERT-Large-Instruct をファインチューニングすることは、従来の分類ヘッドベースのアプローチと同等のパフォーマンスを発揮することが確認されました。特定のデータセットでは、むしろそれらを上回る結果も得られています。実際、この方法こそが最後のギャップを埋め、DeBERTaV3 よりも優れた分類器として ModernBERT を実現するための鍵であると私は考えています。

ここで注意すべき点は、これらのタスクの一部のトレーニングセットが、私たちの事前学習ミックスの中に比較的小さな割合ではあるが存在していることです。しかし、複数のエポックにわたってファインチューニングを行うことで両方の手法が明確に「ドメイン内」領域に入るため、この効果はむしろ最小限であると予想されます。

現代性が重要である

恥知らずな自己剽窃だが適切なミーム

最後に、この可能性がすべての事前学習された MLM エンコーダーに本質的に備わっているのか、それとも ModernBERT に固有のものなのかを知りたかった。この問いに答えるため、RoBERTa-Large などの古いモデルや、現代的なアーキテクチャを持つものの小規模で多様性に欠けるデータでトレーニングされたモデルに対して同じアプローチを適用したところ、パフォーマンスは大幅に低下しました。

モデル

MMLU

ModernBERT-Large-Instruct

43.06

GTE-en-MLM-Large

36.69

RoBERTa-Large

33.11

これは、ModernBERT-Large-Instruct と非常に類似したアーキテクチャ（効率化の調整を除く）を採用している GTE-en-MLM-Large との間には大きなパフォーマンスの差があることから、MLM エンコーダーにおける強力な生成型下流タスクのパフォーマンスは、十分に大規模で多様なデータミックスでトレーニングされていることに大きく依存していることを示唆しています。RoBERTa-Large から GTE-en-MLM-Large への相対的に小さなパフォーマンス向上は、より優れたアーキテクチャの採用が役割を果たす一方で、その効果はトレーニングデータのそれほど顕著ではないことを示唆しているようです。

今後の展望

これらの結果は有望ですが、まだ非常に初期段階です。これらが本当に示しているのは、MLM ヘッドが多目的ヘッドとしての可能性を示すことだけであり、その限界を押し広げるにはほど遠いものです。その他にも以下のような課題があります:

より良く、多様なテンプレート化の探求

トレーニングメカニズムおよびダミー例の影響に関するより詳細な分析

より優れた構成を持つ最新の指示データセットでのテスト

ファーストショット学習能力の調査

より大規模なモデルサイズへのスケーリング

…その他にも多くの課題があります!

これらすべてが、将来の研究に向けた非常に有望な方向性として私たちに映ります。実際、これらのいくつかについては、すでに優れた人々が取り組んでいるという話を聞いています…

究極的に、ここで提示した極めてシンプルなアプローチの結果は、エンコーダーモデルにとって新たな可能性を開くものだと私たちは信じています。ModernBERT-Large-Instruct モデルは HuggingFace で利用可能です。