Mrt5 Large

🚀 MrT5 Largeモデルカード

MrT5は、エンコーダにトークン削除メカニズムを統合し、入力シーケンスの長さを動的に短縮する、ByT5のより効率的なバリアントです。既存のバイトレベルモデルの実用的な制限に対する解決策を提供します。

🚀 クイックスタート

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

from transformers import AutoModelForSeq2SeqLM
import torch

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

input_ids = torch.tensor([list("Life is like a box of chocolates.".encode("utf-8"))]) + 3  # add 3 for special tokens
labels = torch.tensor([list("La vie est comme une boîte de chocolat.".encode("utf-8"))]) + 3  # add 3 for special tokens

# Forward pass with hard deletion
loss = model(input_ids, labels=labels, hard_delete=True).loss

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

from transformers import AutoModelForSeq2SeqLM, AutoTokenizer

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

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

# Forward pass with hard deletion
loss = model(**model_inputs, labels=labels, hard_delete=True).loss

✨ 主な機能

• エンコーダにトークン削除メカニズムを統合し、入力シーケンスの長さを動的に短縮。
• 削除されたトークンの重要な情報をよりコンパクトなシーケンスに「マージ」することで、既存のバイトレベルモデルの制限を解消。
• 注意機構にsoftmax1を使用。

📚 ドキュメント

引用

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

@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}
}

@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",
}

モデルの詳細

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

モデルアーキテクチャ

MrT5 Largeは、標準のByT5 Largeのモデル構成を使用しています。フィードフォワード次元数は3840、モデル次元数は1536、エンコーダ層は36層、デコーダ層は12層、各層に16個のアテンションヘッドがあり、合計12.3億のパラメータを持っています。

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

MrT5 LargeはByT5 Largeから初期化され、同じトレーニング目的でファインチューニングされます。トレーニング前にランダムに初期化されるのは、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％のトークンがマスクされ、ByT5論文で規定された仕様に従います。

最適化

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

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

環境への影響

• ハードウェアタイプ：NVIDIA A100 - SXM4 - 80GB
• GPU数：4
• 使用時間：約73時間
• クラウドプロバイダー：Stanford NLP Cluster

モデルカードの作成者

Julie Kallini
kallini@stanford.edu