Mrt5 Small

🚀 MrT5 Small モデルカード

MrT5は、エンコーダにトークン削除メカニズムを組み込み、入力シーケンスの長さを動的に短縮する、ByT5 (Xue et al., 2022)のより効率的なバリアントです。一定数のエンコーダ層を通過した後、学習された削除ゲートが、どのトークンを削除し、どのトークンを後続の層で保持するかを決定します。削除されたトークンの重要な情報をよりコンパクトなシーケンスに「マージ」することで、MrT5は既存のバイトレベルモデルの実用的な制限に対する解決策を提供します。

🚀 クイックスタート

MrT5は、ByT5と同様に生のUTF - 8バイトで動作し、トークナイザーを使用せずに利用できます。MrT5のコードをロードするには、trust_remote_code=Trueを設定してください。

from transformers import AutoModelForSeq2SeqLM
import torch

model = AutoModelForSeq2SeqLM.from_pretrained('stanfordnlp/mrt5-small', trust_remote_code=True)

input_ids = torch.tensor([list("Life is like a box of chocolates.".encode("utf-8"))]) + 3  # 特殊トークン用に3を追加
labels = torch.tensor([list("La vie est comme une boîte de chocolat.".encode("utf-8"))]) + 3  # 特殊トークン用に3を追加

# ハード削除を伴う順伝播
loss = model(input_ids, labels=labels, hard_delete=True).loss

バッチ推論とトレーニングには、ByT5のトークナイザークラスを使用できます。

from transformers import AutoModelForSeq2SeqLM, AutoTokenizer

model = AutoModelForSeq2SeqLM.from_pretrained('stanfordnlp/mrt5-small', trust_remote_code=True)
tokenizer = AutoTokenizer.from_pretrained('google/byt5-small')

model_inputs = tokenizer(["Life is like a box of chocolates.", "Today is Monday."], padding="longest", return_tensors="pt")
labels = tokenizer(["La vie est comme une boîte de chocolat.", "Aujourd'hui c'est lundi."], padding="longest", return_tensors="pt").input_ids

# ハード削除を伴う順伝播
loss = model(**model_inputs, labels=labels, hard_delete=True).loss

✨ 主な機能

• MrT5は、エンコーダにトークン削除メカニズムを組み込み、入力シーケンスの長さを動的に短縮することができます。
• 既存のバイトレベルモデルの実用的な制限に対する解決策を提供します。
• シーケンス-to-シーケンスタスクに適したエンコーダ - デコーダアーキテクチャです。

📦 インストール

このモデルはtransformersライブラリを使用してロードできます。以下のコードでモデルをロードできます。

from transformers import AutoModelForSeq2SeqLM
model = AutoModelForSeq2SeqLM.from_pretrained('stanfordnlp/mrt5-small', trust_remote_code=True)

📚 ドキュメント

引用

このモデルを使用する場合は、MrT5の論文を引用してください。

@inproceedings{
    kallini2025mrt,
    title={MrT5: Dynamic Token Merging for Efficient Byte-level Language Models},
    author={Julie Kallini and Shikhar Murty and Christopher D Manning and Christopher Potts and R{\'o}bert Csord{\'a}s},
    booktitle={The Thirteenth International Conference on Learning Representations},
    year={2025},
    url={https://openreview.net/forum?id=VYWBMq1L7H}
}

また、ByT5の論文も引用してください。

@article{xue-etal-2022-byt5,
    title = "{B}y{T}5: Towards a Token-Free Future with Pre-trained Byte-to-Byte Models",
    author = "Xue, Linting  and
      Barua, Aditya  and
      Constant, Noah  and
      Al-Rfou, Rami  and
      Narang, Sharan  and
      Kale, Mihir  and
      Roberts, Adam  and
      Raffel, Colin",
    editor = "Roark, Brian  and
      Nenkova, Ani",
    journal = "Transactions of the Association for Computational Linguistics",
    volume = "10",
    year = "2022",
    address = "Cambridge, MA",
    publisher = "MIT Press",
    url = "https://aclanthology.org/2022.tacl-1.17",
    doi = "10.1162/tacl_a_00461",
    pages = "291--306",
}

モデルの詳細

これは3億パラメータのMrT5 Small (mrt5-small)のモデルカードです。これはByT5 Small (google/byt5-small)のより効率的なバリアントで、シーケンス長を平均で約50％削減するようにトレーニングされています。

モデルアーキテクチャ

MrT5 Smallは、標準のByT5 Smallのモデル構成を使用しています。これは、フィードフォワード次元が3584、モデル次元が1472、エンコーダ層が12層、デコーダ層が4層、各層に6つのアテンションヘッドがあり、合計3億のパラメータを持っています。

MrT5には追加の削除ゲートがあり、これによりエンコーダのシーケンス長を動的に削減します。このモデルでは、3番目のエンコーダ層の後に配置され、それ以降のすべての層は削減されたシーケンスで動作します。このモデルは削除率δ = 0.5でトレーニングされており、つまり、3番目の層の後でエンコーダのシーケンス長を約50％削減します。MrT5のゲートメカニズムは、わずか3000個の追加パラメータを導入するだけです。

MrT5 SmallはByT5 Smallから初期化され、同じトレーニング目的でファインチューニングされます。トレーニング前にランダムに初期化されるのは、MrT5の削除ゲートのみです。MrT5のもう1つの特徴は、アテンションメカニズムでsoftmax1を使用することです。

用途

このモデルは主にシーケンス-to-シーケンスタスク用に設計されたエンコーダ - デコーダアーキテクチャです。探索的または学術的な目的でそのまま使用することもできますが、特定の下流タスクで最適なパフォーマンスを達成するには、ファインチューニングが推奨されます。

モデルの削除機能を利用するには、付属のリポジトリにあるカスタムのMrT5Trainerを使用してください。この専用のトレーナーは、ファインチューニング中に削除メカニズムが適切に維持および統合されることを保証します。

これは学術的および研究的な探索のために構築されたベースモデルであり、本番環境でのデプロイを目的としていません。ユーザーは、特に信頼性と堅牢性が重要な設定において、モデルの出力を注意深く評価する必要があります。

バイアス、リスク、制限

言語モデルは、さまざまな形の社会的バイアスを示し、有害または不快なコンテンツを生成することが知られています（Bender et al., 2021; Bommasani et al., 2022; Liang et al., 2022）。他の言語モデルと同様に、このモデルもバイアスまたは有害な出力を生成する可能性があります。安全性のためにファインチューニングされておらず、特に敏感なコンテキストでは慎重に使用する必要があります。

トレーニング詳細

トレーニングデータ

継続的な事前トレーニングには、多言語C4 (mC4)コーパスを使用しています（Raffel et al., 2020; Xue et al., 2021）。MrT5は、15の類型的に多様な言語でトレーニングされています。mC4コーパスから十分に大きなサンプルを抽出し、mC4のトレーニング分割から各言語の等サイズのサンプル（バイト単位）を抽出しています。

トレーニング手順

MrT5は、ByT5のスパン破損事前トレーニング目的でトレーニングされています。このタスクでは、ラベルのないテキストデータのトークンのスパンが、各スパンに1つのセンチネルトークン IDで置き換えられ、モデルは欠落したトークンを埋める必要があります。ByT5とMrT5の場合、これらはバイトのスパンであり、マスクが単語境界と干渉する可能性があります。

前処理

スパン破損目的でトレーニングする際には、平均マスクスパン長が20トークン、ノイズ密度が15％となるように破損スパンを計算します。つまり、シーケンス内のトークンの15％がマスクされます。

最適化

MrT5は、2^20トークンのバッチ（つまり、エンコーダシーケンス長が1024、有効バッチサイズが1024）で5,000回の勾配ステップでトレーニングされています。AdamWオプティマイザを使用し、初期学習率を1e - 4として線形減衰させ、ウォームアップは行っていません。

特定のシーケンス長削減率を達成するために、目標削除率δ = 0.5のPIコントローラを使用しています。また、論文の付録Dに記載されているように、アテンションスコア正則化を使用しています。

環境への影響

• ハードウェアタイプ：NVIDIA RTX 6000 Ada Generation
• GPU数：1
• 使用時間：約63時間
• クラウドプロバイダー：Stanford NLP Cluster

モデルカード作成者

Julie Kallini
kallini@stanford.edu