BERTは単なる単一テキスト拡散ステップにすぎない
(nathan.rs)- BERT と RoBERTa のようなマスク付き言語モデルも テキスト拡散モデル の概念として解釈できる
- 従来の Auto-regressive(例: GPT)とは異なり、ブロック単位の生成と段階的復元方式の適用可能性が実証されている
- マスキング率をステップごとに調整して訓練すると、RoBERTa でも自然なテキスト生成が可能であることが実験で確認された
- 既存の構造を変えないで訓練目標のみを調整するだけでも、生成能力を確保できる
- GPT-2 と比較すると、RoBERTa ベースの拡散モデルでも一定レベルの一貫したテキスト生成結果を導出する
概要
Google DeepMind が公開した Gemini Diffusion は、従来の GPT 系列とは異なり、テキストを一度にブロック単位で生成する 拡散ベース言語モデル である。この方式は、ランダムノイズを段階的に洗練してテキストを作るアプローチだ。関連論文(Large Language Diffusion Models)を調べた結果、離散型言語拡散は マスク付き言語モデリング(MLM) の一般化という事実を確認した。つまり、このアイデアを使って BERT 系モデルでもテキスト生成が可能かを実験的に検討した。
注: 後続の DiffusionBERT 論文は、同様のアイデアをより厳密に検証している。
Transformer の簡単な歴史
2017年に最初に提案された Transformer は、エンコーダー-デコーダー構造だった。2018年には、エンコーダー(BERT: 双方向、マスクド復元中心)とデコーダー(GPT: 自己回帰、逐次予測中心)が分離され、それぞれに特化したモデル群が登場した。
- エンコーダー専用(BERT 系)
- 文全体の文脈を入力し、特定の部分を
<MASK>で隠して残りから復元 - 文表現、分類などで性能上の強み
- 文全体の文脈を入力し、特定の部分を
- デコーダー専用(GPT 系)
- 与えられた順序で次のトークンを予測
- 生成、要約、翻訳などで頭角を現す
BERT は分類などにすぐ活用されたが、その後 GPT 系列は生成能力が向上し、さらに多様な用途を確保した。
離散的言語拡散モデル
拡散モデルは元々画像生成で普及した。 画像の場合:
- 順方向プロセス: きれいな画像に徐々にガウシアンノイズを追加して、純粋なノイズ状態に到達
- 逆方向プロセス: 深層学習モデルで段階的にデノイズを繰り返し、元データを復元
テキストへの適用では、最もシンプルな方法は マスキングベースのノイズ過程 である。
- 順方向(マスキング)
- t=0 で元のテキスト、ステップが進むほど一部のトークンを
<MASK>にランダム置換 - 最終ステップでは全体が
<MASK>で埋め尽くされる
- t=0 で元のテキスト、ステップが進むほど一部のトークンを
- 逆方向(デノイズ)
- Transformer エンコーダーが、与えられたマスキング設定で元トークンの復元を学習
- 低いマスキング率なら復元しやすく、比率が高まるほど復元が難しくなる
- 高いマスキング率から低い比率の順に繰り返して、全シーケンスを生成可能
このような拡散フレームワークでは、モデルは様々なマスキング率のステップごとにデノイズ損失を合算して学習する。BERT のマスキング復元目標は事実上テキスト拡散の一部である。マスキング率のスケジュールと反復デノイズを組み合わせれば、BERT の目標を自然言語生成手順へ拡張できる。
RoBERTa Diffusion 実験
RoBERTa は 2019 年に発表された BERT のハイパーパラメータとデータ拡張、より単純な学習目的(MLM Only)が特徴的なモデルである。 実験では HuggingFace transformers、datasets ライブラリを使って RoBERTa のデフォルト重み、トークナイザー、Trainer を利用した。 WikiText データセットをベースに、以下のプロセスでファインチューニングを実施した。
- 10段階拡散スケジュール(mask_probs: 1.0~0.1)から1つを各バッチごとにサンプリングしてマスキング
- カスタム
diffusion_collatorでマスキング確率を決定し、各トークンごとに確率的に<MASK>を適用 - プロンプト文脈の維持のため最初の16トークンは常に保持
データマスキング(カスタム collator):
- 各サンプルのトークン束をパディング後、マスキング確率をランダム選択
- 最初の16トークンを除く残りに対して確率的に
<MASK>を適用 - マスキング済みデータと正解ラベルセットを返す
生成(推論):
- 256トークン長のシーケンスを入力し、最初の16トークンはプロンプト、残りは
<MASK> - 各ステップでモデルが予測したトークンをサンプリングして埋め、さらに一定比率を再びリマスキング
- マスキング比率を徐々に減らして繰り返し、最後には全体が復元される
サンプル生成結果:
- プロンプトに続く自然な後続テキストを含み、実際にかなり一貫性のあるテキストが生成される
- 一部の特異点は、WikiText のデータ前処理形式(例:
@-@ハイフン)に起因する
GPT-2 との比較
- GPT-2 はやや高速で一貫性も高いが、RoBERTa Diffusion も期待を上回る性能を示す(さらに漸進的な改善の余地が高い)
- 新しい AR-Diffusion、Skip-Step Diffusion 方式と最適化により品質/速度の向上が可能
結論
- **RoBERTa のように設計されたマスク付き言語モデルでも、マスキング率を調整しながら訓練すれば自然言語生成エンジンになり得る
<MASK>トークンでテキストを段階的に劣化/復元する学習方式だけでも、完全に生成的モデルへ変換できる可能性を確認- 構造そのものを変更せず、訓練目標の変更だけで生成能力を確保できる
- 結果として、BERT 系モデルも本質的には1つのテキスト拡散モデルと一致する
1件のコメント
Hacker News の意見
BERTが最初に登場したとき、誰もがテキスト生成を試みたが、たいていうまくいかなかったことについての参考論文はこちら。DeepMindのSavinovらは、学習中に2段階だけ適用し、マスキング確率をランダム化すれば、かなり実用的に動かせることを示した
先週、BERTで対話させようといろいろ試したあとで、この投稿を通じてこうした内容を知った。まだ試してみたいアイデアがあるので、引き続き楽しく実験している 関連ブログ
私は、BERTはテキスト生成ではなく、純粋なエンコーダー、たとえば意味類似度の測定や分類にしか使えないと考えていた大きな集団の一員だった
2021年のこの論文(5ページ)で、このつながりが初めて指摘されたと理解している。テキスト拡散で単語を意味的に近い別の単語で汚染するアプローチを試していたが、モデルにマスクされた単語を見つけさせるほうが簡単だった。このアプローチの歴史はさらに前のこの論文までさかのぼり、拡散の数学で説明せずに生成的MLMを作っていた
それよりさらに古い。2014年にLi Yaoらがこの論文で、自己回帰(次トークン予測)モデルと生成的確率ネットワーク(デノイジング・オートエンコーダー、拡散モデルの先駆け)の等価性を示した。並列サンプリングが逐次サンプリングをうまく近似すると主張しており、2016年の私の研究 Counterpoint by Convolution(リンク)でもこの方法を適用し、むしろより良い性能が得られると説明した。残念ながら応用論文だったため、拡散分野からの引用はあまり得られなかった。これよりさらに古いルーツもあると思う
私も、その論文が最初に公式に扱ったリファレンスだったと記憶している(BERTを見たあと、NLP研究者が拡散の概念を思いつきそうだ)。3年前に集めたテキスト拡散の初期リファレンス集はここにある
参考までに、関連論文はこちらでも確認できる
拡散ベースのアプローチのほうが、動物の脳で起きていることにより近いように感じる。私は話すとき、前の単語に基づいて1語ずつ作っているというより、まず頭の中にぼんやりした全体のアイデアがあり、それを言語として整理するプロセスのほうが大きいと感じる
自己回帰LLMも、実際にはそんな単純に次の単語だけを予測しているわけではない。1回のフォワードパスで1トークンずつ出てくるとはいえ、潜在空間では長期的な計画や推論が明確に現れている。私たちが実際に話すときも、頭の中で全体像を描いてから順に話すので、拡散方式のほうが私たちに似ているとは言えない
個人的には、拡散を一種の「スペクトル自己回帰」として説明したブログ記事が気に入った。拡散は低周波の特徴を先に予測し、高周波の特徴を後で予測する傾向がある
Slackでメッセージを5回くらい修正するたびに、自分は拡散モデルだと感じることがよくある
話したり文章を書いたりするとき、実際には単語を順番に出している。ただし文を始める前に、頭の中ではすでにおおまかな要点を先に考えている。LLMも、トークンを出力する前に潜在空間で全体の方向性を先に計画しているのと似ている
自己回帰LLMも、何を話すかを事前に内部で計画していることが解釈研究で示されている
個人的には、マスキングして埋めることだけを許すより、Levenshteinスタイルの挿入・削除型編集をサポートするほうが自然だと感じる。たとえば最終段階で、より長い同義語に単語を置き換えようとすると、右方向にずらせないので、この種の拡散構造では簡単ではない
こうした単純な実験は、原理を素早く把握できて興味深い。テキスト拡散モデルで惜しいのは、トークンを連続値ではなく離散値として扱わなければならない点。画像ではピクセルが連続的なので自然にノイズを加えられるが、テキストトークンはそうではない。たいていは完全な置換として処理することになり、埋め込み空間にノイズを加えたり、埋め込み自体を直接学習したりする複数のアプローチがあるものの、どれも画像拡散より複雑だ
私は、この論文のようなアプローチにより期待している。連続的な潜在拡散と自己回帰Transformerベースのテキスト生成を組み合わせている。オートエンコーダーとTransformerを同時に、あるいは独立して学習できる
テキスト拡散モデルを最初に見たとき、「これってただのMLMでは?」と思った。私が注目していたのはMaskGITだ。この方式が本当に拡散と呼べるのは、モデルが誤ったトークンを正しいトークンに置き換えるよう学習されている場合だろう。というのも、連続拡散の強みはノイズ耐性にあるからだ。しかし実際には、誤ったトークンを生成するというアイデアは簡単ではなく、まだあまり試されていないように思う
私はUTF8バイト単位でMLMを実験し、非定型な翻訳モデルを試してみた。カリキュラム学習と段階的なランダム汚染方式を使った。単にノイズを加えるのが目的なら、ランダムなインデックスをランダムなバイト値に置き換えるだけで十分だ。たとえば次のようなパターンを入力する
ターゲット系列全体を反復的に評価して、難易度段階が上がるたびに学習ドメインで不連続な変化が起きないようにする。結果として、用語や分類は気にしなくなり、「これが拡散なのか」が重要なのではなく、「実際にうまく動くのか」だけを気にするようになった
ELECTRAとの比較がどうなるのかも気になる
インラインコード補完モデルをファインチューニングして、Cursorに近い性能が出るか試してみたかった(実際には難しいだろうが面白そう)。ただ、土台として使えそうなオープンな拡散モデル自体がまだないので期待している
画像拡散モデルの魅力的な点は、ランダムノイズから画像を生成することだった。では、なぜテキスト拡散モデルはランダムトークンではなく、すべてマスクされた空トークンから始めるのだろうか
それはやりたいタスク次第だ。単純なテキスト補完なら、入力文はマスクせずに与え、続くマスクトークンをモデルに埋めさせればいい。コード編集などをしたいなら、編集する部分だけをマスクして順次埋めさせることもできる。テキスト拡散モデルの強みは、コード編集のような実用的な編集を実行できることだ(自己回帰LLMは、命令で編集するための別システムが必要になる)。十分に適したコード編集ヘッドでモデルを学習させれば、多くのコーディング作業の速度と効率を高められる。今後はGPT Codexのようなところに、こうした機能が自然に統合されていく気がする
みなが空トークンから始める方式だけを使っているわけではない。ランダムトークンから始めるもの、マスクから始めるもの、埋め込みベクトルから始めるものなど、さまざまな試みがある