BERT は単なる単一のテキスト拡散ステップです

少し前に、Google DeepMind が発表されました ジェミニの拡散、拡散を使用してテキストを生成する実験的な言語モデル。一度に 1 つの単語を生成する従来の GPT スタイルのモデルとは異なり、Gemini Diffusion はランダム ノイズを段階的に調整することでテキストのブロック全体を作成します。

新聞を読みました 大規模言語普及モデル そして、離散言語の普及が、私たちがそれ以来取り組んでいるマスク言語モデリング (MLM) を一般化したものにすぎないことを知って驚きました。 2018年。私が最初に考えたのは、「テキスト生成を行うために BERT のようなモデルを微調整できないだろうか?」ということでした。私は好奇心から簡単な概念実証を試してみることにしました。

注: 記事を書いた後、偶然この論文を見つけました。 ブロードキャストBERT これは本質的に同じことを行いますが、より厳密なテストが行われます。この投稿に興味があればチェックしてください。

トランスフォーマーの歴史#

で導入されたオリジナルの Transformer アーキテクチャ 2017年、エンコーダ デコーダ モデルでした。 2018 年、研究者たちは、モデルのエンコーダーとデコーダーのコンポーネントを分離できることに気づきました（ バート そして GPT)、2 つの異なるモデル ファミリが作成されました。

1. エンコーダ専用モデル (BERT スタイル、双方向)

エンコーダー モデルは、トレーニング目標としてマスク言語モデリング (MLM) を使用しました。つまり、各入力のトークンのサブセットをランダムにマスクし、不足しているトークンを再構築する (空白を埋める) ようにエンコーダーをトレーニングします。モデルは (部分的にマスクされた) コンテキスト全体を一度に認識し、双方向の表現を学習します。このアーキテクチャは、全文 (または段落) 表現を必要とするタスク (分類や検索など) に優れています。

2. デコーダ専用モデル (GPT スタイル、自己回帰)

デコーダ モデルは、次のトークンの予測をトレーニング目標として使用しました。つまり、各位置 $t$ で、コンテキストとして $t$ までのすべてのトークンを与えて、位置 $t + 1$ のトークンを予測します。将来の値を予測するために左側のコンテキストのみが使用されます (一方向)。このアーキテクチャは、オープンエンドの生成、要約、翻訳など、一度に 1 トークンずつテキストを生成する生成タスクに優れています。

当初、BERT は分類などのタスクですぐに使用されていましたが、GPT スタイルのモデルが普及したのは後になってからでした (初期の機能が限られていたため)。最終的に、自己回帰 (デコーダ) トランスフォーマーの生成能力は大幅に向上しました。 「次のトークンの予測」という一般的なトレーニング目標は、エンコーダー モデルと比較した場合、ユースケースのスペースがはるかに広いことを意味します。

離散言語の普及モデル#

拡散モデルは、画像生成において初めて普及しました。画像生成では、拡散モデルは徐々にガウス ノイズを画像に追加し (順方向プロセス)、次にニューラル ネットワークをトレーニングして繰り返しノイズを除去します (逆方向プロセス)。画像による継続的拡散の概要は次のとおりです。

1. 前進プロセス: きれいな画像から始めます ×₀次に、ほぼ純粋なノイズが得られるまで、各タイムステップに少量の (通常はガウス) ノイズを追加します。
2. 逆のプロセス: モデル (多くの場合 U-Net) をトレーニングして各タイムステップでノイズを予測し、離散的なノイズ除去ステップで徐々に元の画像を回復します。

このアイデアを言語に適用すると、テキストにノイズを追加し、それを段階的に除去する方法が必要になります。これを行う最も簡単な方法は、 マスキングベースのノイズ処理:

1. フォワード（マスキング）処理:

   • タイムステップで t = 0、完全に破損していないテキストシーケンスが得られます。
   • 後続の各タイムステップで t > 0、トークンの一部を特別なトークンにランダムに置き換えます。 事前に定義されたスケジュールに従ってトークンを生成します (たとえば、マスクされる割合を 0% から 100% まで徐々に増加させます)。
   • 最終タイムステップまでに T、シーケンス全体がマスクされる場合があります（すべてのトークンは ）。

2. 逆（ノイズ除去）プロセス:

   • モデル (多くの場合、標準の Transformer エンコーダー) をトレーニングして、タイムステップで部分的にマスクされたシーケンスが与えられた場合に元のトークン ID を予測します。 t。
   • これは、マスクされた言語モデリングをさまざまなマスク レートで実行することに似ています。初期のタイムステップでは、少数のトークンのみがマスクされます (予測が容易です)。後のタイムステップでは、多くのトークンがマスクされます (より難しくなります)。
   • 高マスクレートからゼロまで予測を連鎖させることで、完全なシーケンスを復元 (または生成) できます。

この離散テキスト拡散フレームワークでは、モデルは、固定マスク確率での単一の MLM 目標ではなく、すべてのタイムステップにわたるノイズ除去損失の合計を最適化することによって、データ分布の尤度限界を学習します。

ご覧のとおり、BERT のマスクされた言語モデリングの目的は、 テキストの拡散と同じトレーニング目的ですが、マスキング率のサブセットのみが対象です。可変マスキング レート (0 から 1) と、スケジュールされた一連のノイズ除去ステップ (拡散理論にヒントを得た) を導入することで、BERT のマスクされた言語モデリングの目的を完全な生成手順に変換できます。

ロベルタのディフュージョン#

2019年には、 ロベルタ リリースされました。これは本質的に元の BERT モデルを強化したもので、ハイパーパラメータ、データ トレーニング サイズ、およびより単純なトレーニング目標 (MLM のみ、次の文の予測を削除) が改善されました。

ここではHuggingFaceを使用します transformers そして dataset オリジナルの RoBERTa 重み、トークナイザー、および Trainer クラスを取り込むためのライブラリを使用して、WikiText データセット上のモデルを簡単に微調整します。メインコード (完全なコードはこちら) 以下のようになります。

# Load and tokenize dataset and instantiate the model
dataset = load_dataset("wikitext", "wikitext-2-raw-v1")
tokenizer = RobertaTokenizerFast.from_pretrained("roberta-base")
model = RobertaForMaskedLM.from_pretrained("roberta-base")

# Create the training args and Trainer instance
training_args = TrainingArguments(
    output_dir="finetuned-roberta-diffusion",
    overwrite_output_dir=True,
    num_train_epochs=NUM_EPOCHS,
    per_device_train_batch_size=BATCH_SIZE,
    save_strategy="epoch",
    save_total_limit=1,
    logging_steps=200,
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized["train"],
    eval_dataset=tokenized["validation"],
    data_collator=diffusion_collator, # custom implementation
    tokenizer=tokenizer,
)

# Train & save
trainer.train()
trainer.save_model("finetuned-roberta-diffusion")

現在、10 の拡散ステップがあるため、そのうちのパーセント $p$ をランダムにサンプリングします。 mask_probs (1.0、0.9、0.9、…、0.1) として、各バッチのトークンのパーセントをマスクします。カスタム diffusion_collator 関数 （ここのコードを参照してください) 1 つのマスク確率をサンプリングします。 p から mask_probs バッチごとに各トークンを次のように設定します と p 確率。

「プロンプト」で生成を条件付けできるようにするために、現在、最初の 16 個のトークンをマスクすることはありません。つまり、トレーニング中、各ステップには生成のコンテキストとして常に最初の 16 個のトークンが含まれます。

の簡略化されたコード diffusion_collator 次のようになります:

  def diffusion_collator(examples):
      batch = tokenizer.pad(examples, return_tensors="pt")

      # Randomly select masking probability for this batch
      mask_prob = random.choice([1.0, 0.9, 0.8, 0.7, 0.6, 0.5, 0.4, 0.3, 0.2, 0.1])

      # Never mask the first PREFIX_LEN tokens (preserved context)
      maskable_positions = batch.input_ids[:, PREFIX_LEN:]

      # Create random mask for the chosen probability
      mask = torch.rand(maskable_positions.shape)  mask_prob

      # Apply masking
      batch.input_ids[:, PREFIX_LEN:][mask] = tokenizer.mask_token_id
      batch.labels = batch.input_ids.clone()

      return batch

推論のために、サイズ 256 のテンソルである入力から開始します (256 個のトークンのブロックを生成しているため)。最初の 16 桁はプロンプトに対応するトークン ID で、最後の 240 桁は単なる トークン。ノイズ除去スケジュールと各ステップを繰り返し、予測を生成し、シーケンスを再度マスクします。プロセスは次のようになります。

Step 0: [PREFIX]      ...     (100% masked)
Step 1: [PREFIX] will  over  control ...        (90% masked)
Step 2: [PREFIX] will begin  greater control ...      (80% masked)
...
Step 10: [PREFIX] will begin to assert greater control ...  (0% masked - DONE)

生成用の簡略化されたコードは次のようになります。

# Generate text through iterative denoising
for step, mask_prob in enumerate(mask_probs):
    # Forward pass: predict masked tokens
    with torch.no_grad():
        outputs = model(input_ids=input_ids, attention_mask=attention_mask)
        predictions = outputs.logits  # shape: (1, MAX_LEN, vocab_size)

    # For each masked position, sample from top-k/top-p filtered distribution
    for pos in range(PREFIX_LEN, MAX_LEN):
        if input_ids[0, pos] == tokenizer.mask_token_id:
            logits = predictions[0, pos, :]
            # Apply top-k and top-p filtering
            filtered_logits = top_k_top_p_filtering(logits, top_k=TOP_K, top_p=TOP_P)
            probs = F.softmax(filtered_logits, dim=-1)
            # Sample token
            sampled_token = torch.multinomial(probs, 1)
            input_ids[0, pos] = sampled_token

    # Re-mask a portion of non-prefix tokens for next iteration
    if mask_prob > 0:
        mask_indices = torch.rand(MAX_LEN - PREFIX_LEN)  mask_prob
        input_ids[0, PREFIX_LEN:][mask_indices] = tokenizer.mask_token_id

H200 で 30 分間トレーニングした後の微調整モデルの出力生成例を次に示します (最初の行は最初のプロンプトです)。

Following their victory in the French and Indian War, Britain began to assert
greater...

...dominion over Europe beginning about the early 19th. There conflict took
place on the island, between British and Irish Ireland. British officials 
administered British Ireland, a Celtic empire under the control of the Irish 
nationalist authorities, defined as a dominion of Britain. As the newly Fortic 
states acquired independent and powerful status, many former English colonies
played their part in this new, British @-@ controlled colonial system. Following
this period the Non @-@ Parliamentaryist Party won its influence in Britain in 
1890, led by the support of settlers from the Irish colonies. Looking inwards, 
Sinclair, Lewis questioned, and debated the need to describe " The New Britain "

出力は驚くほど一貫しているように見えます。存在する癖のほとんどは、実際には WikiText の書式設定 (句読点の周りのスペース) からの癖に過ぎません。 "、ハイフンを反転 - の中へ @-@）。

以下は、私たちの拡散モデルと GPT-2 の比較です。

GPT-2 の出力はより一貫性があり、わずかに高速であることがわかります (~9 秒対 ~13 秒)。しかし、私の単純な実装がどれほど優れていたかに嬉しい驚きを感じています。これは優れた概念実証であり、AR 拡散やスキップステップ拡散などの新しいアプローチ (およびより最適化された実装) を使用すると、品質と速度を大幅に向上できます。

結論#

RoBERTa のようなマスクされた言語モデルは、元々は穴埋めタスク用に設計されましたが、可変レート マスキングを離散拡散プロセスとして解釈することで、完全な生成エンジンに再利用できることがわかりました。テキストを徐々に破損させることで、 トークンを使用し、マスク強度を増加させながら反復的にノイズを除去するようにモデルをトレーニングすることで、標準的な MLM 目標を段階的な生成手順に効果的に変換します。

アーキテクチャを変更しなくても、微調整された RoBERTa は、トレーニング目標をわずかに変更した後、一貫した外観のテキストを生成できます。これにより、BERT スタイルのモデルは本質的には 1 つのマスキング レートでトレーニングされた単なるテキスト拡散モデルであるという考えが検証されます。