Fine-tune LLM with ROUGE

You've fine-tuned a model and it produces summaries, but how do you know if they're actually good? Without a rigorous, automated evaluation, you're flying blind — tweaking hyperparameters and hoping for the best. This lesson solves that pain by teaching you how to fine-tune for summarization with ROUGE, the standard metric for measuring summary quality, and how to use it to guide your training loop toward better results in a measurable, repeatable way.

The problem this lesson solves

Manual evaluation doesn't scale. When you fine-tune a model for summarization, you need a way to compare candidate summaries against a reference (human-written) summary — and you need that comparison to be objective, fast, and automated.

Without a clear metric, you might ship a model that sounds fluent but misses key facts, or one that's overly verbose and does no useful abstraction. ROUGE (Recall-Oriented Understudy for Gisting Evaluation) fills this gap by measuring overlap between the generated summary and the reference, giving you a numeric score you can track during training.

Why this matters now: As you progress in the LLM Finetuning track, you'll soon tune hyperparameters and compare runs. A reliable metric like ROUGE is what makes those comparisons meaningful.

Core concept / mental model

Think of ROUGE as a string-matching score with a purpose. It compares n-grams (sequences of words) between your generated summary and the reference, then reports precision, recall, and F1.

• ROUGE-1: overlap of unigrams (single words) — captures word presence.
• ROUGE-2: overlap of bigrams — captures phrase order.
• ROUGE-L: longest common subsequence — captures sentence-level fluency.

The core idea: the more overlapping important words and phrases your summary has with the reference, the better aligned it is with the human-written gold standard. ROUGE is not a perfect judge — it's a proxy for quality, but it's a useful proxy because it's cheap, deterministic, and correlates reasonably with human judgment for extractive-style summarization.

Analogy: grading essays by keyword matching

Imagine a teacher grading an essay by counting how many key terms from the textbook appear in the student's answer. That's ROUGE in essence — it rewards using the "right" words and phrases, even if the student's sentence structure is different. It doesn't penalize for style, but it does reward coverage of the source material.

How it works step by step

1. Prepare your dataset with document and summary (reference) pairs.
2. Load a pretrained model (e.g., t5-small) and its tokenizer.
3. Tokenize the input documents and summaries, handling truncation and padding.
4. Define a training loop using Hugging Face Trainer with Seq2SeqTrainingArguments.
5. Add a compute_metrics function that calls rouge_score to calculate ROUGE scores on the validation set.
6. Train and evaluate — watch ROUGE scores improve as training progresses.
7. Iterate — adjust learning rate, batch size, or epochs based on ROUGE performance.

Hands-on walkthrough

Let's build a complete example. First, ensure you have the necessary libraries installed:

pip install transformers datasets evaluate rouge_score

Now create a script to fine-tune a T5-small model on a small summarization dataset (we'll use the cnn_dailymail subset for illustration).

from datasets import load_dataset
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, Seq2SeqTrainingArguments, Seq2SeqTrainer
import evaluate
import numpy as np

# 1. Prepare data
dataset = load_dataset("cnn_dailymail", "3.0.0", split="train[:1%]")
val_split = load_dataset("cnn_dailymail", "3.0.0", split="validation[:10%]")

model_name = "t5-small"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)

# 2. Tokenize
def preprocess_function(examples):
    inputs = [doc for doc in examples["article"]]
    model_inputs = tokenizer(inputs, max_length=512, truncation=True, padding="max_length")
    labels = tokenizer(examples["highlights"], max_length=64, truncation=True, padding="max_length")
    model_inputs["labels"] = labels["input_ids"]
    return model_inputs

tokenized_train = dataset.map(preprocess_function, batched=True, remove_columns=dataset.column_names)
tokenized_val = val_split.map(preprocess_function, batched=True, remove_columns=val_split.column_names)

# 3. Define ROUGE metric
rouge = evaluate.load("rouge")

def compute_metrics(eval_pred):
    predictions, labels = eval_pred
    decoded_preds = tokenizer.batch_decode(predictions, skip_special_tokens=True)
    labels = np.where(labels != -100, labels, tokenizer.pad_token_id)
    decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)
    result = rouge.compute(predictions=decoded_preds, references=decoded_labels, use_stemmer=True)
    return {k: v.mid.fmeasure for k, v in result.items()}

# 4. Training args
training_args = Seq2SeqTrainingArguments(
    output_dir="./results",
    evaluation_strategy="epoch",
    learning_rate=3e-5,
    per_device_train_batch_size=4,
    per_device_eval_batch_size=4,
    weight_decay=0.01,
    num_train_epochs=3,
    predict_with_generate=True,
    report_to="none",
)

# 5. Trainer
trainer = Seq2SeqTrainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_train,
    eval_dataset=tokenized_val,
    compute_metrics=compute_metrics,
)

# 6. Train
trainer.train()

Expected output (abbreviated):

Epoch 1: Loss 2.1, Rouge1 0.35, Rouge2 0.15, RougeL 0.30
Epoch 2: Loss 1.8, Rouge1 0.38, Rouge2 0.17, RougeL 0.32
Epoch 3: Loss 1.6, Rouge1 0.40, Rouge2 0.18, RougeL 0.33

Notice how the ROUGE scores climb as the loss drops — that's the signal you're improving.

Compare options / when to choose what

ROUGE is not the only evaluation metric. Here's a quick comparison:

When to choose what: - For quick, automated training feedback, use ROUGE. - If your summaries are highly abstractive, consider BERTScore as a complement. - For final evaluation before release, use LLM-as-judge or human eval.

Pro tip: In production, combine ROUGE with a small human-labeled test set — ROUGE can be gamed by repeating the source, so add a faithfulness check like factuality metrics.

Troubleshooting & edge cases

• ROUGE scores don't improve: You may be overfitting — check your training loss vs eval loss. Try lower learning_rate, increase warmup_steps, or add weight_decay.
• Zero scores on ROUGE-2: If your summaries are highly abstractive, bigram overlap may be zero. Consider using ROUGE-L or BERTScore instead.
• Slow evaluation: ROUGE computation on a large eval set is slow. Use a subset (e.g., 10%) for validation.
• Decoded labels contain padding tokens: Ensure you replace -100 with tokenizer.pad_token_id before decoding.
• predict_with_generate=True is required — otherwise the trainer uses model.generate() which is standard for summarization.
• Memory errors on long documents: Truncate with max_length, use gradient accumulation, or switch to LoRA (check the LoRA lesson in this track).

What you learned & what's next

You now understand the core idea behind fine-tune for summarization with ROUGE: it's a metric-driven approach that lets you quantify how well your summaries align with human-written references. You applied it in a hands-on exercise, adding ROUGE evaluation to a training loop. You also learned how to interpret the scores and use them to guide hyperparameter tuning.

Next step: In the next lesson, you'll learn how to interpret and troubleshoot ROUGE scores — diving deeper into diagnosing bad summarization models and when to consider alternative metrics like BERTScore.