--- library_name: transformers tags: [] language: - en - fr - es - de - el - bg - ru - tr - ar - vi - th - zh - ai - sw - ur datasets: - allenai/c4 ---
# Model Card for MrT5 Large [**MrT5: Dynamic Token Merging for Efficient Byte-level Language Models**](https://arxiv.org/pdf/2410.20771)\ (Kallini et al., 2024)
![](MrT5.png) **MrT5** (**M**e**r**ge**T5**) is a more efficient variant of [ByT5 (Xue et al., 2022)](https://arxiv.org/abs/2105.13626) that integrates a token deletion mechanism in its encoder to *dynamically* shorten the input sequence length. After processing through a fixed number of encoder layers, a learned *delete gate* determines which tokens are to be removed and which are to be retained for subsequent layers. By effectively "merging" critical information from deleted tokens into a more compact sequence, MrT5 presents a solution to the practical limitations of existing byte-level models. ## Citation If you use this model, please cite the MrT5 paper: ```bibtex @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} } ``` Also cite the ByT5 paper: ```bibtex @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", } ``` ## Model Details This is the model card for the 1.23B-parameter **MrT5 Large** (`mrt5-large`), a more efficient variant of ByT5 Large (`google/byt5-large`). This model is trained to reduce sequence lengths by ~50% on average. - **Developed by:** Julie Kallini, Shikhar Murty, Christopher D. Manning, Christopher Potts, Róbert Csordás - **Model type:** MrT5 - **Languages:** English, French, Spanish, German, Greek, Bulgarian, Russian, Turkish, Arabic, Vietnamese, Thai, Chinese, Hindi, Swahili, and Urdu - **Fine-tuned from model:** [google/byt5-large](https://huggingface.co/google/byt5-large) - **Sources for more information**: - [GitHub Repository](https://github.com/jkallini/mrt5) - [Paper](https://arxiv.org/abs/2410.20771) ### Model Architecture MrT5 Large uses the model configuration of the standard ByT5 Large, which has a feed-forward dimensionality of 3840, a model dimensionality of 1536, 36 encoder layers, 12 decoder layers, 16 attention heads in each layer, and 1.23B total parameters. MrT5 has an additional *delete gate*, which dynamically reduces the encoder sequence length. In this model, it is placed after the third encoder layer, and all subsequent layers operate on a reduced sequence. This model was trained with a deletion rate of δ=0.5, which means that the model reduces its encoder sequence length by ~50% after the third layer. MrT5’s gating mechanism only introduces an additional 3,000 parameters. MrT5 Large is initialized from ByT5 Large and fine-tuned on the same training objective. Only MrT5's delete gate is randomly initialized before training. The other distinguishing feature of MrT5 is that it uses [softmax1](https://www.evanmiller.org/attention-is-off-by-one.html) in its attention mechanism. ## Uses This model is an encoder-decoder architecture designed primarily for sequence-to-sequence tasks. While it can be used as-is for exploratory or academic purposes, fine-tuning is recommended to achieve optimal performance on specific downstream tasks. To leverage the model’s deletion feature, please use the custom **MrT5Trainer** available in the [accompanying repository](https://github.com/jkallini/mrt5). This specialized trainer ensures that the deletion mechanism is properly maintained and integrated during fine-tuning. Because this is a base model built for academic and research explorations, it is not intended for production-grade deployments. Users should carefully evaluate the model’s outputs, especially in any setting where reliability and robustness are critical. ## Bias, Risks, and Limitations Language models are known to exhibit various forms of social bias and may produce harmful or offensive content ([Bender et al., 2021](https://dl.acm.org/doi/pdf/10.1145/3442188.3445922); [Bommasani et al., 2022](https://arxiv.org/abs/2108.07258); [Liang et al., 2022](https://arxiv.org/abs/2211.09110)). Like other language models, this model may produce biased or harmful outputs. It has not been fine-tuned for safety and should be used with caution, especially in sensitive contexts. ## How to Get Started with the Model Like ByT5, MrT5 works on raw UTF-8 bytes and can be used without a tokenizer. Make sure to set `trust_remote_code=True` to load the MrT5 code: ```python 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 ``` For batched inference and training, you can use ByT5's tokenizer class: ```python 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 ``` ## Training Details ### Training Data For continued pre-training, we use the [multilingual C4 (mC4) corpus](https://huggingface.co/datasets/allenai/c4) ([Raffel et al., 2020](https://arxiv.org/abs/1910.10683); [Xue et al., 2021](https://arxiv.org/abs/2010.11934)). MrT5 is trained on 15 typologically diverse languages: English, French, Spanish, German, Greek, Bulgarian, Russian, Turkish, Arabic, Vietnamese, Thai, Chinese, Hindi, Swahili, and Urdu. To avoid training models for multiple epochs, we ensure that the samples drawn from the mC4 corpus are sufficiently large. Additionally, we extract equal-sized samples for each language (in terms of bytes) from the mC4 training split. ### Training Procedure MrT5 is trained on the ByT5 span corruption pre-training objective. In this task, spans of tokens in unlabeled text data are replaced with a single *sentinel token* ID per span, and the model must fill in the missing tokens. For ByT5 and MrT5, these are spans of bytes, and the masks can potentially interfere with word boundaries. #### Preprocessing When training on the span corruption objective, we calculate the corrupted spans such that the average masked span length is 20 tokens with a noise density of 15%—that is, 15% of tokens in the sequence are masked out, following the specification outlined in the ByT5 paper. #### Optimization MrT5 is trained for 5,000 gradient steps over batches of 2^20 tokens (i.e., an encoder sequence length of 1024 with an effective batch size of 1024). We use the AdamW optimizer with an initial learning rate of 1e-4 with linear decay and no warmup. To achieve a specific sequence length reduction rate, we use a PI controller with a target deletion ratio of δ=0.5, as described in Section 3.2 of the paper. We also use attention score regularization, as described in Appendix D of the paper. ## Environmental Impact - **Hardware Type:** NVIDIA A100-SXM4-80GB - **GPU Count**: 4 - **Hours used:** ~73 hours - **Cloud Provider:** Stanford NLP Cluster ## Model Card Authors Julie Kallini \ kallini@stanford.edu