ModernBERT のブログを読んでみた感想

コミューン株式会社で機械学習エンジニアをしている深澤です。
この記事はCommuneアドベントカレンダー(シリーズ2)の 21 日目の記事です。

直前まで全く違う記事を書いていたのですが、ModernBERT というアツいモデルが発表されて話題になっており、Huggingface の公式からブログも出ているので、早速読んでみた感想をアドベントカレンダーの記事にすることにしました。感想だったり、付属事項を調べてみたので記事として昇華することを目指しています。
翻訳したもの、というよりは要点をかいつまんだものになっています。詳細は以下の記事を御覧ください。

• 記事

https://huggingface.co/blog/modernbert

• 論文

https://arxiv.org/abs/2412.13663

• コード

https://github.com/AnswerDotAI/ModernBERT

ModernBERT のここがすごい

• Token length が 8192 に拡張(！？)
  • もともとは 512 です。10倍以上になってる...
• GLUE スコアもかなり向上
• それでいてbase (139M params) and large (395M params) という(現代では相対的に)破格の小ささ
  • かつては BERT を見たらモデルサイズがでかすぎる！！！と思っていましたが、いまや 1B 以下のモデルとなるとなんて小さいんだ…と思ってしまう

ModernBERT のモチベーション

なぜ decoder-only (GPTなど) のモデルが全盛期を迎えている今、encoder-only の古臭いモデルである BERT の改善に乗り出したのか、というモチベーションを説明してくれています。

Decoder-Only モデルはでかすぎる

パラメータサイズがでかすぎるので何をするにもお金がかかるし、推論の時間もかかるし、ということで取り回しづらいことはみなさんもよくご存知かなと思います。多くは語りませんが、記事の中では車に例えられています。

a frontier model like OpenAI's O1 is like a Ferrari SF-23

In contrast, a BERT model is like a Honda Civic.

扱いやすい車である BERT を使った応用先は枚挙に暇がありません。記事内では特に RAG について言及されており、文脈を考慮したベクトルを取り出すのに最も都合の良いのは BERT 系だというのはだいたい同意できるところではないでしょうか。

ということで、BERT を現代の技術を使って良くしようぜ！というプロジェクトが ModernBERT とのことです。2018 年は Computer Science 界隈にとっては千年前に相当しますので...

BERT was released in 2018 (millennia ago in AI-years!)

ModernBERT の改善ポイント

RoPE の導入

LLama にも使われている RoFormer(2020) で提案されたトークン位置埋め込み。当初の BERT は絶対位置埋め込みだったのですが、これを用いているとのことです。

https://huggingface.co/blog/designing-positional-encoding

位置埋込は絶対的から相対的、にだんだん移り変わってきているなと感じますが、一番最新のトレンドは回転行列でこれを表現することだな〜と最近の論文を読んでいて思います。

なお https://dl.acm.org/doi/10.1145/3626772.3657805 EulerFormer も近いアイデアになっていて、これは更に複素数空間で扱おうというアイデア

Attention の変更

まず用いる Attention は Flash Attention に変更。そのうえで Alternating Attention というものを導入しています。
これは毎回 global attention(全単語との注意を見る) をするのではなく、local attention(周辺の単語しか見ない) を時々入れる、というもの。

人間が本を読むときのことを思えば、毎回すべての文章に注意を向けているかというとそうではなく、全体と局所を組み合わせているだろう、とのこと。
これがかなりインパクトがあったみたいです。

似たアイデアは以下とかかな…

• Focal Attention for Long-Range Interactions in Vision Transformers
  • https://proceedings.neurips.cc/paper_files/paper/2021/hash/fc1a36821b02abbd2503fd949bfc9131-Abstract.html
• LGAFormer: transformer with local and global attention for action detection
  • https://dl.acm.org/doi/10.1007/s11227-024-06138-1

Padding をやめる

個人的に一番気になっているところでした。なんと、あの親しんできた Padding をやめようというのが盛り込まれています。

• encoder-only モデルはバッチ内で文章の長さを揃えて渡す必要がある。そのために padding を行ってきた
• これは良い解決策だが、非効率的
• そこで、padding をやめてバッチサイズ1で全部一つの sequence にまとめることにした
• できるだけ複数の文章をたくさんつなげて、モデルの最大入力長に近づけたバッチを作る

言われてみればそ、そうなんだ、という感じですがいくつか疑問(IMO)が出てきたので書き留めておきます。

• 最大長に近づける、とあるけど最大長になるように最後は PAD トークンを追加するのかな？
  • それとも異なる長さのシーケンスをバッチサイズ1として処理し続けるのか
• 一つのシーケンスにしたとして、ModernBERT はシーケンス全体のベクトルを返すことになるのか？一つ一つのベクトルがほしい場合はどうすれば…？
• …推論時にはそれぞれの文章を別々で渡してベクトルを貰えればいいんだろうけど、学習するときにはいろんな文章をくっつけた sequence の MLM を一気にやるんだろうか

みたいなことを呟いてたら答えてもらえました。ありがたい...！！！

Samples are still processed individually through careful masking.

という一文もあるので、よく読めてなかったです🙏

更にこのあともう少し調べていたのですが、Flash Attention 自体がそもそも packing(このブログで言う unpad?) という操作をするようになっていて、すでによく知られた方法みたいでした。Flash Attention を勉強します...

https://qiita.com/jovyan/items/4490b1371b4d790573b5#flash-attentionとパディングについて

その他の細かい改善

• GeLU 活性化層は gated linear unit activation function (GeGLU) に変更
• 不必要な bias 項は除去
• Normalization Layer を追加

学習について

だんだん力尽きつつあるので、箇条書きでまとめています。

データ

• 古いモデルの多くは、一般的にWikipediaとWikibooksからなる限られたコーパスで訓練しており多様性にかけている
• ModernBERT ではウェブ文書、コード、科学論文など、様々な英語ソースからのデータ(2兆個のトークン)で学習を行った

学習プロセスについて

BERT的な学習方法に則ったうえでいくつかの改善を実施

• Next Sentence Prediction は削除
  • やめた既存研究として一番有名なのは RoBERTa ？
    • https://arxiv.org/abs/1907.11692
• mask 率は 15% から 30% に増加させる
  • 引き上げている研究としてすぐに思い浮かんだのは ELECTRA、効率的な学習ができるようになったので 25% にしたよ、という話。
    • https://openreview.net/pdf?id=r1xMH1BtvB

学習では3段階のステージを実施しています。

• まず、1.7T個のトークンを1024個のシーケンス長で学習する。

• 次にシーケンス長8192, 250Bトークンで学習する。

• 最後に、長い文章と短い文章を混ぜ合わせた 50B トークンで学習する(ProLong の学習方法と同じ)

  • How to Train Long-Context Language Models (Effectively)

• 学習率は最初の2つは一定で行い、最後だけ Warmup-Stable-Decay スケジュールの学習率で計算する。

  • Pythia と同じように checkpoint を公開するので研究に役立ててほしいとのこと
    • https://arxiv.org/abs/2304.01373

学習フェーズがサラッと書かれていましたが、ちょっと理解が曖昧なのでもしかしたらまた後で記述を足しておきます。

その他の学習 TIPS

• バッチサイズのウォームアップ: 最初は小さいバッチサイズから始めて、だんだん大きくする
• weight initialization via tiling for the larger model size
  • Microsoft の Phi と同様のテクニックらしい(調べたけどわからなかった)
  • 小さいモデルの重みを並べて大きなモデルの初期重みとする感じ？
  • 乱数で初期化するよりも良さそう
• torch.randperm でデータをシャッフルする際、データサイズがかなり大きいと id が若いところに偏ってしまう部分があるらしいので気をつける
  • こわい…

読んだ感想

完走した感想です。

• sequence packing: こんなことしていいんだって思って勉強になった
• Alternating Attention もよかった。Global, Local の Attention を混ぜちゃおうは自然な発想だし、これのインパクトはかなりあったみたいなので手元の実験などでも活用してみたい
• これまでの研究で進んでいた要素技術をできるだけ組み込んで一気に性能が向上した、というのはこれ以上ない良い話だなと思いました。
  • 個別の要素それぞれについて、Ablation Study の続報があったらぜひ読みたいな

次は論文とコードを読んでみて色々取り入れたいな、と思います。
日本語の ModernBERT もほしいな... 👀

ここまで記事を読んでいただきありがとうございました！