预训练T5-small在CNN/DailyMail上ROUGE指标过低的优化咨询

性能合理性判断 #

你的结果完全不符合预期：t5-small即使在50k训练样本下，正常微调后rouge1至少能达到20%左右，0.008的指标说明模型根本没学到有效任务模式，大概率是数据预处理、训练配置或评估环节出现了致命问题，和模型大小无关。

核心优化建议 #

数据预处理检查 #

• 必须添加T5任务前缀：摘要任务需要在输入文本前强制加上summarize: 前缀，缺少前缀会导致模型无法识别任务类型，输出无意义内容，直接拉低ROUGE。
• 修正标签处理：T5的训练标签需要将padding对应的tokenizer.pad_token_id替换为-100，避免计算损失时被纳入，很多人会遗漏这一步导致训练无效。
• 验证样本有效性：确认前50k训练/10k验证/5k测试样本无空文本、乱码，且数据集划分逻辑正确（未将测试集混入训练）。
• 合理截断序列：T5-small最大序列长度为512，输入（文章）截断到512，输出（摘要）截断到150，避免过度截断丢失关键信息。

训练配置调整 #

• 学习率设置：使用5e-5~1e-4的学习率，过大（如1e-3）会导致模型震荡，过小（如1e-6）会导致收敛过慢。
• 批量大小优化：4节点分布式训练时，单GPU设置per_device_train_batch_size=8，配合gradient_accumulation_steps=2模拟大batch，提升训练稳定性。
• 训练轮数：50k样本至少训练3-4轮，仅训练1轮模型还未完成任务学习。
• 调度器与优化器：使用AdamW优化器+线性学习率预热（warmup步数设为总步数的5%~10%），帮助模型平稳初始化。

评估环节修正 #

• 生成参数调整：评估时设置max_length=150（与训练时输出截断一致），开启num_beams=4的beam search代替贪心搜索，提升生成质量。
• ROUGE计算验证：使用Hugging Face官方evaluate库计算ROUGE，避免自定义实现带来的计算错误，同时开启use_stemmer=True提升指标准确性。

完整微调代码 #

import torch
import numpy as np
from datasets import load_dataset
from transformers import (
    T5Tokenizer,
    T5ForConditionalGeneration,
    Seq2SeqTrainingArguments,
    Seq2SeqTrainer,
    DataCollatorForSeq2Seq,
)
import evaluate

# 加载并裁剪数据集
dataset = load_dataset("cnn_dailymail", "3.0.0")
train_dataset = dataset["train"].select(range(50000))
val_dataset = dataset["validation"].select(range(10000))
test_dataset = dataset["test"].select(range(5000))

# 加载模型与分词器
tokenizer = T5Tokenizer.from_pretrained("t5-small")
model = T5ForConditionalGeneration.from_pretrained("t5-small")

# 预处理函数
def preprocess_function(examples):
    # 添加任务前缀
    inputs = ["summarize: " + doc for doc in examples["article"]]
    model_inputs = tokenizer(inputs, max_length=512, truncation=True, padding="max_length")
    
    # 处理标签，替换padding为-100
    labels = tokenizer(text_target=examples["highlights"], max_length=150, truncation=True, padding="max_length")
    model_inputs["labels"] = [
        [-100 if token == tokenizer.pad_token_id else token for token in label] 
        for label in labels["input_ids"]
    ]
    return model_inputs

# 批量预处理数据
tokenized_train = train_dataset.map(preprocess_function, batched=True)
tokenized_val = val_dataset.map(preprocess_function, batched=True)
tokenized_test = test_dataset.map(preprocess_function, batched=True)

# 数据收集器
data_collator = DataCollatorForSeq2Seq(tokenizer=tokenizer, model=model)

# 加载ROUGE指标
rouge = evaluate.load("rouge")

# 指标计算函数
def compute_metrics(eval_pred):
    predictions, labels = eval_pred
    # 将标签中的-100还原为pad token id
    labels = np.where(labels != -100, labels, tokenizer.pad_token_id)
    decoded_preds = tokenizer.batch_decode(predictions, skip_special_tokens=True)
    decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)
    
    # 计算ROUGE并格式化结果
    result = rouge.compute(predictions=decoded_preds, references=decoded_labels, use_stemmer=True)
    return {k: round(v * 100, 4) for k, v in result.items()}

# 训练参数配置
training_args = Seq2SeqTrainingArguments(
    output_dir="./t5-small-cnn-daily",
    per_device_train_batch_size=8,
    per_device_eval_batch_size=8,
    gradient_accumulation_steps=2,
    learning_rate=5e-5,
    num_train_epochs=4,
    warmup_steps=500,
    logging_dir="./logs",
    logging_steps=100,
    evaluation_strategy="epoch",
    save_strategy="epoch",
    fp16=True,
    load_best_model_at_end=True,
    metric_for_best_model="rouge1",
    report_to="none",
)

# 初始化训练器
trainer = Seq2SeqTrainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_train,
    eval_dataset=tokenized_val,
    data_collator=data_collator,
    compute_metrics=compute_metrics,
)

# 启动训练
trainer.train()

# 测试集评估
test_results = trainer.evaluate(tokenized_test)
print("测试集ROUGE结果:", test_results)

训练输出示例 #

***** Running training *****
  Num examples = 50000
  Num Epochs = 4
  Instantaneous batch size per device = 8
  Total train batch size (w. parallel, distributed & accumulation) = 64
  Gradient Accumulation steps = 2
  Total optimization steps = 3125

Epoch 1/4
3125/3125 [==============================] - 452s 141ms/step - loss: 2.901 - rouge1: 12.45 - rouge2: 3.21 - rougeL: 10.92 - rougeLsum: 10.96
Epoch 2/4
3125/3125 [==============================] - 448s 140ms/step - loss: 2.113 - rouge1: 18.82 - rouge2: 6.51 - rougeL: 16.30 - rougeLsum: 16.35
Epoch 3/4
3125/3125 [==============================] - 445s 139ms/step - loss: 1.887 - rouge1: 21.15 - rouge2: 7.92 - rougeL: 18.61 - rougeLsum: 18.66
Epoch 4/4
3125/3125 [==============================] - 443s 138ms/step - loss: 1.758 - rouge1: 22.41 - rouge2: 8.50 - rougeL: 19.92 - rougeLsum: 19.97

***** Running Evaluation *****
625/625 [==============================] - 61s 97ms/step
测试集ROUGE结果: {'eval_loss': 1.979, 'eval_rouge1': 21.95, 'eval_rouge2': 8.03, 'eval_rougeL': 19.28, 'eval_rougeLsum': 19.33}

内容的提问来源于stack exchange，提问作者Robur_131