DiffusEdit T5edit

t5edit is google/flan-t5-large, a sequence-to-sequence model with 780M parameters, fine-tuned to rewrite an outdated article from evidence. It follows the design of EDIT5 from the FRUIT paper by Iv et al. (2022), in which the model writes the updated article but copies an unchanged sentence through a short marker instead of spelling it out again.

The model was trained as the upper bound for DiffusEdit, a system that keeps a frozen diffusion language model and trains only small classifiers over its hidden states. The classifiers are at gamemaker1/diffusedit and the data at gamemaker1/wiki-wrong. Every parameter of t5edit was trained, so it bounds what a fully trained generator reaches on the same data. The training script is code/pipeline/experiments/editor.py in github.com/gamemaker1/diffusedit.

Input and output

The input names each sentence of the outdated article with an index and appends the evidence passages.

update: [0] <sentence 0> [1] <sentence 1> ... evidence: (0) <title>: <text> (1) <title>: <text> ...

The output names each source sentence in order. [j] means that sentence j is copied unchanged, #j <text> means that sentence j is replaced by the text, and a sentence whose index is absent is deleted.

[0] #1 <the rewritten sentence 1> [2]

The input is cut to 1,024 tokens from the end, so the evidence is shortened before the article, and the output is cut at 384 tokens. The T5 vocabulary cannot write some letters, such as the ø in a Norwegian name, so after decoding, a word of the output that equals a word of the input with its missing letters removed is replaced by that input word.

Training

The model was fit on the train split of wiki-wrong. The validation loss was computed after each epoch, the checkpoint with the lowest validation loss was kept, and training stopped after one epoch without improvement. The learning rate was 1e-4 with a linear decay and a warm-up over the first 5% of the steps. Each optimiser step used 32 examples, as eight accumulated micro-batches of four, with gradient clipping at 1.0, bfloat16 autocast and seed 42, on one A100 with 40 GB of memory. The validation loss per epoch and the time taken are in run_metrics.json.

Use

from transformers import AutoModelForSeq2SeqLM, AutoTokenizer

tok = AutoTokenizer.from_pretrained("gamemaker1/t5edit")
model = AutoModelForSeq2SeqLM.from_pretrained("gamemaker1/t5edit")

In the training script, input_text builds the input string from a wiki-wrong record and parse maps the generated text back to one output sentence per source sentence. run_config.json and run_metrics.json in this repository are the configuration and the metrics that the training run wrote.

Licence

Apache 2.0, the licence of the base model.