|
--- |
|
library_name: transformers |
|
tags: [] |
|
language: |
|
- en |
|
- fr |
|
- es |
|
- de |
|
- el |
|
- bg |
|
- ru |
|
- tr |
|
- ar |
|
- vi |
|
- th |
|
- zh |
|
- ai |
|
- sw |
|
- ur |
|
datasets: |
|
- allenai/c4 |
|
--- |
|
<div align="center"> |
|
|
|
# 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) |
|
|
|
</div> |
|
|
|
 |
|
|
|
<!-- Provide a quick summary of what the model is/does. --> |
|
|
|
**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 \ |
|
[email protected] |
|
|
