世界を動かす技術を、日本語で。

BERTは単一のテキスト拡散ステップに過ぎない

概要

  • Google DeepMind が提案した Gemini Diffusion は、従来のGPT型生成モデルとは異なる拡散型テキスト生成モデル
  • 拡散モデル の原理をテキスト生成に応用し、マスク率を変化させながら段階的にノイズ除去を実施
  • BERTRoBERTa のようなマスク言語モデルも拡散的生成モデルへ転用可能
  • 実験的に RoBERTa を拡散的生成モデルとしてファインチューニングし、整合性のあるテキスト生成を実現
  • MLM (Masked Language Modeling)と 拡散モデル の関係性や今後の発展可能性を考察

拡散型言語モデルの基本と発展

  • Gemini Diffusion は、従来の GPT 系の1トークンずつ生成とは異なり、 ノイズからブロック単位でテキスト生成 を行う新方式
  • 拡散モデルは元々画像生成で普及し、 ノイズ加算→段階的除去 という2段階プロセスを持つ
    • 画像の場合: ガウスノイズ を加え、U-Net等で逆方向にノイズ除去
    • テキストの場合: マスク を段階的に増やし、逆に除去して元のテキストを再構築
  • マスク言語モデル(MLM) は、拡散モデルの特殊ケース(一定マスク率のみ)であり、マスク率を可変にすれば 拡散的生成 が可能
  • これにより、BERTやRoBERTaも 段階的生成モデル として機能

Transformerアーキテクチャの歴史と分類

  • Transformer は2017年に Encoder-Decoder 構造で登場
  • 2018年以降、 Encoderのみ(BERT)Decoderのみ(GPT) に分化
    • Encoder系(BERT) :双方向文脈、MLMで訓練、分類や検索タスクに強み
    • Decoder系(GPT) :次トークン予測、生成タスクに強み
  • BERTは即座に分類タスクで活用されたが、生成力はGPT系列に軍配
  • 拡散的生成 の登場で、BERT系も生成タスクに新たな可能性

離散言語拡散モデルの仕組み

  • 前方(ノイズ付加)プロセス :段階的にトークンを<MASK>で置換、最終的に全てマスク
  • 逆方向(ノイズ除去)プロセス :部分的にマスクされたテキストから元のトークンを予測・復元
  • マスク率を段階的に変化 させることで、従来のMLMを拡張
  • 各ステップでの損失和を最適化し、 生成能力 を獲得

RoBERTa拡散モデルの実装例

  • HuggingFace Transformers を用い、RoBERTaをWikiTextデータセットで拡散的にファインチューニング
    • 10段階のマスク率(1.0, 0.9, ..., 0.1)をランダムに適用
    • 先頭16トークンは常にマスクせず、 プロンプト条件付き生成 を実現
  • 主な訓練・生成フロー
    • 訓練時:バッチごとにマスク率を選択、<MASK>で置換
    • 生成時:プロンプト+<MASK>で初期化し、段階的にノイズ除去・トークン復元
  • 実験結果:30分の訓練でも 一貫性あるテキスト生成 が可能
    • GPT-2と比較して若干遅いが、意外なほど高品質

まとめと今後の展望

  • RoBERTaのようなMLMモデル も、マスク率を可変にした拡散的訓練で 生成モデル へ転用可能
  • アーキテクチャ変更なしでも、 段階的ノイズ除去による生成 が実現
  • 今後は AR-DiffusionSkip-Step Diffusion などの新手法、実装最適化でさらなる品質・速度向上に期待
  • MLMと拡散モデル の関係性理解が、今後の言語モデル設計の指針となる

Hackerたちの意見

テキスト拡散モデルが出始めたとき、俺もこの人と同じことを思ったよ(「待って、これってただのMLMじゃん」)って。もっとMaskGITのことを考えてたけどね。「拡散」って呼ぶ理由は、モデルが間違ったトークンを正しいものに置き換えることを学ばなきゃならないってことかな(連続拡散の大きな特徴はノイズ耐性だから)。でも、良い間違ったトークンを考えるのは難しいから、誰もやってないと思う。

electraとの比較が気になるな。

それかdebertaでも、でもやっぱり超面白いね!

俺の知る限り、この関連性は2021年に最初に指摘されたんだよね。https://arxiv.org/abs/2107.03006(5ページ目)。俺たちは、テキスト拡散をやりたくて、意味的に似た言葉に言葉を腐らせる(例えば「速い茶色の狐」→「素早い黒い犬」)ことを考えてたんだけど、マスキングの方がモデルが見つけやすいってことが分かったんだ。歴史的には、さらに遡るとhttps://arxiv.org/abs/1904.09324があって、拡散数学にフレーミングせずに生成的MLMを作ったんだよね。

これも関連がある - https://arxiv.org/pdf/1902.04094

俺にとって、拡散ベースのアプローチは流行りのLLMのトークンを一つずつ処理するアプローチよりも、動物の脳で起こっていることに近い「感じ」がする。自分のことを言うと、俺は前に話した言葉に基づいて一つずつ言葉を生成するわけじゃなくて、頭の中にぼんやりしたアイデアを持って始めて、それを言語に整然とするのが難しいんだ。

よくSlackでメッセージを書いた後、5回も編集することがあるんだ…今はそれをやるたびに拡散モデルみたいに感じる。

自分が認知的に意識しているってことは、これは拡散とは程遠い証拠だね。むしろ、思考のトークンを考えるような、熟考に近い感じかな。もし今のLLMのメタファーを見つけなきゃならないとしたら。

あなたは100%の確率で、言葉を一つずつ順番に発音したり書いたりしてるよね。でも、文を始める前に、言いたいことの要点を頭の中でまとめてる。それは、LLMの潜在空間でも、最初のトークンを出力する前に起こることなんだ。

もしかしたら、これは二つの異なる思考モードかもしれないね。ふわっとした思考がまとまることもあるけど、時には線形に考えを繋げることもできる。脳は両方できるのかも。

解釈可能性の研究によると、自己回帰型LLMも自分が何を言うかを事前に計画しているらしいよ。

だから、私はジェミニ拡散にすごくワクワクしてるんだ。[1] - [1] https://deepmind.google/models/gemini-diffusion/

それは、一度に一つのトークンを扱うアプローチに反しているわけではないよ。

これが接続を説明している私のお気に入りのブログ記事だよ:https://sander.ai/2024/09/02/spectral-autoregression.html 彼らは拡散を「スペクトル自己回帰」の一形態と呼んでいて、最初に低周波数の特徴を予測し、その後に高周波数の特徴を予測する傾向があるんだ。

流行のLLMのトークンを一つずつ生成するアプローチについて。個人的には、前に話した言葉を基に一つずつ言葉を生成するわけじゃないんだ。実際、自己回帰型LLMもそうじゃない。確かに、1回のフォワードパスで1トークンしか得られないけど、潜在空間で何が起こっているかを見ると、次のトークンだけを超えた長期的な計画や推論の明確な兆候があるんだ。だから、拡散と比べて私たちにとって必ずしも似ているわけでもないと思う。私たちも一言ずつ順番に言うけど、全体像を頭に入れてるからね。

画像拡散モデルの魅力の一部は、ランダムなノイズから画像を生成するところだと思う。テキスト拡散モデルは、なんで真っ白な状態(つまり、すべて「マスクされた」トークン)から始まるの?ランダムなトークンから始めないのかな?

でも、みんながそうしてるわけじゃないよ。いろんなアプローチが試されてる。ランダムなトークンから始めるものや、マスクから始めるもの、さらにはランダムなベクトル埋め込みから始めるものもある。

モデルに何をさせたいかによるよね。テキストを補完させたいなら、マスクされていない入力テキストの後に、モデルが埋めるべきマスクされたトークンをいくつか提供することになる。もしかしたら、モデルにコードの一部を単純に編集させたいのかもしれない。その場合、編集すべき部分をマスクして、モデルがそのマスクされたトークンを生成したトークンで逐次埋めていくことになる。テキスト拡散モデルの強力な能力の一つは、コードにあるらしい。自己回帰型LLMは、編集する能力を持っているわけじゃないんだ。他のシステムが編集コマンドとして解釈する指示を生成できるだけ。編集したい部分を文字通りマスク解除できるのは、かなり強力なパラダイムで、たくさんのコーディングタスクを速くするだけでなく、改善する可能性もあると思う。テキスト拡散の要素が、GPT Codexのようなコーディングモデルに組み込まれるんじゃないかと予想してる(もうすでにそうかもしれないけど)。コード編集専用に設計された拡散出力ヘッドをトレーニングする理由はないし、同じモデルがそのヘッドを使う必要があるときに利用できるはずだよ。

こういうシンプルな実験を見るのが大好き。サクッと読めて、原理をちょっと理解できるからね。テキストディフューザーに関しての私のつまずきの一つは、理想的にはトークンを離散的に扱うんじゃなくて、むしろ確率的なフィールドとして扱うべきだってこと。画像ディフューザーは、ピクセルが連続的な値を持つ自然な特性がある。色から色へスムーズに移行できる。でも、トークンはそうはいかない。ここでは完全に置き換えをするだけ。トークンにノイズを加えることはできないし、埋め込み空間で作業する必要がある。でも、どうやって埋め込みを直接トレーニングできるの?いろんなアプローチが試されているのを見つけたけど、どれも画像ベースの拡散プロセスよりずっと複雑なんだ。

BERTが出た頃、みんなテキスト生成を試みてたけど、あんまりうまくいかなかったんだよね。参考までに一つ紹介するけど、これね: https://arxiv.org/abs/1902.04094 これは明示的に拡散と関係してるわけじゃないけど、DeepMindのSavinovたちが、トレーニング時に2ステップやってマスキング確率をランダムにするだけで、かなりうまくいくことを発見したんだ。

インラインコード補完モデルを微調整して、カーソルに近づけられないかずっと考えてたんだけど(無理だけど、楽しそう)、私が知ってる限り、ベースとして使えるオープンな拡散モデルはないし、特に良いものもないんだ。早く実用的なものが出てくるといいな。

この拡散モデルのもう一つの簡単な改善点は、ログ確率がトークンがマスクに変わる確率にも影響を与えるようにすることだと思う。だから、高い信頼度のトークンはプルーニングされる可能性が低くなるべきで、もっと早く収束するはず。バックプロパゲーションがそれを利用できるかどうか気になるな。(私はMLエンジニアじゃないけど)

こういうアプローチにもっとワクワクしてる: https://openreview.net/forum?id=c05qIG1Z2B 彼らは自己回帰型トランスフォーマーベースのテキスト生成と潜在拡散を組み合わせてるんだ。オートエンコーダーとトランスフォーマーは、一緒にトレーニングされることもあるよ。