BRIO: Bringing Order to Abstractive Summarization

In plain words

A summarizer is usually trained to reproduce one reference summary word by word. When it writes a summary, though, it has to pick among many possible wordings, and that training never told it which of them are better. BRIO adds a second training signal to fix this.

For each training article, the model first writes several candidate summaries. Each candidate is scored by how much it overlaps with the reference summary. The model is then trained so that the score it gives each candidate follows the same order as those overlap scores.

The authors train two variants of the same kind of model. One is a scorer: it picks the best of several candidate summaries. The other is a generator: it keeps the word-by-word training as well, so it writes summaries as before, but it has also learned that better summaries should be more likely. The authors report that the generator beats the model it started from on news summarization benchmarks.

TL;DR

BRIO fine-tunes a pre-trained abstractive summarizer (BART for CNN/DailyMail, PEGASUS for XSum) with a sum of two losses: cross-entropy on the reference, and a pairwise margin loss that makes the length-normalised log-likelihoods of 16 diverse-beam candidates follow their ROUGE order §3 eq. 10. The generator variant (BRIO-Mul) writes summaries; the scorer variant (BRIO-Ctr) reranks candidates. Author-reported ROUGE-1 is 47.78 on CNN/DailyMail and 49.07 on XSum tab. 2. Several settings that matter for reproduction differ between the paper and the released code; see Paper vs code.

Problem

Maximum-likelihood training assigns all probability to one reference summary and says nothing about how other summaries rank §2 §3. A model that fits the reference perfectly gives zero probability to any different summary, which the authors call improper where several good summaries exist §3.

At inference, however, beam search compares many system-generated candidates, so the model must order them correctly §1 §3. Training only on references also leaves the model exposed to its own earlier mistakes at test time, the exposure-bias problem §2. In a preliminary study, BART ranked the best and worst of its own candidates correctly for 54.80% of CNN/DailyMail samples tab. 1.

Method

Symbols: D is the source document, S* the reference summary, S_i candidate summaries, M the quality measure (ROUGE against S*) §3.

1. Candidate generation (offline). A pre-trained model writes 16 candidates per training document with diverse beam search app. B §3. The paper states the number of candidates but not the beam width. The base models are BART-large-CNN for CNN/DailyMail and PEGASUS-xsum for XSum app. B.
2. Quality order. Candidates are ordered by M(S_i) = ROUGE(S_i, S*), so the best candidate has the lowest index §3. The paper does not say which ROUGE variant is used for this ordering.
3. Score. Each candidate gets f(S) = sum over tokens of log p(s_t | D, S_<t), divided by |S|^alpha eq. 9 eq. 6. alpha is a length penalty, set to the same value as the beam search it is combined with app. B.
4. Contrastive loss. L_ctr = sum over i, and over j > i, of max(0, f(S_j) − f(S_i) + lambda_ij), with lambda_ij = (j − i) · lambda eq. 8 §3. The margin grows with the rank gap.
5. Cross-entropy. L_xent is the token-level negative log-likelihood of the reference eq. 3. The paper writes it with a one-hot target eq. 4.
6. Multi-task loss. L_mul = L_xent + gamma · L_ctr eq. 10 §3. The cross-entropy term keeps the model usable as a generator §3.
7. Two variants. BRIO-Ctr uses the contrastive loss only and is used as a scorer that selects among candidates. It can no longer generate §3. BRIO-Mul uses the multi-task loss and is the generator §5.1.
8. Target distribution view. The paper's motivating target gives the reference mass 1 − beta, shares beta over the candidates, and orders the candidates by quality eq. 7. This is a description of the intent; the contrastive loss above is how the paper encourages it §3.

Variants used in the experiments: BRIO-Loop fine-tunes again on candidates produced by BRIO-Mul §5.3; BRIO-Few uses 100 (CNN/DailyMail) or 1000 (XSum) training examples §5.5.

# Multi-task loss for one document  [§3] [eq. 8] [eq. 10]
# Input:  doc: source tokens; ref: reference tokens
#         cands: candidates sorted by ROUGE to ref, best first  [§3]
# Output: scalar loss  L_mul = L_xent + gamma * L_ctr  [eq. 10]

def multitask_loss(doc, ref, cands, lam, alpha, gamma):
    # Token-level cross-entropy on the reference only  [eq. 3]
    l_xent = 0
    for t in range(len(ref)):
        l_xent -= log_prob(ref[t], doc, ref[:t])
    # Length-normalised log-likelihood of every candidate  [eq. 9]
    f_scores = [score(doc, c, alpha) for c in cands]
    # Pairwise margin loss over all pairs i < j  [eq. 8]
    # The margin grows with the rank gap j - i  [eq. 8]
    l_ctr = 0
    for i in range(len(cands)):
        for j in range(i + 1, len(cands)):
            margin = (j - i) * lam
            l_ctr += max(0, f_scores[j] - f_scores[i] + margin)
    return l_xent + gamma * l_ctr

# Score one candidate: summed token log-probabilities over length^alpha  [eq. 9]
def score(doc, cand, alpha):
    # Log-probability of each token given the source and earlier tokens  [eq. 6]
    token_logps = [log_prob(cand[t], doc, cand[:t]) for t in range(len(cand))]
    # Divide by the length raised to the penalty alpha  [eq. 9]
    return sum(token_logps) / len(cand) ** alpha

# Use as a scorer: BRIO-Ctr picks one candidate from the pool  [§3] [§5.1]
def select_summary(doc, candidates, alpha):
    # Keep the candidate with the highest length-normalised score  [eq. 9]
    return max(candidates, key=lambda c: score(doc, c, alpha))

Learning-rate schedule, with step counted in updates app. B:

lr = 2 \cdot 10^{-3} \min(\mathrm{step}^{-0.5}, \mathrm{step} \cdot \mathrm{warmup}^{-1.5}) eq. 13

Hyperparameters

Results

Author-reported ROUGE F1 values. Baselines marked * are copied from the original papers, and ‡ marks the authors' own evaluation tab. 2. Differences marked as significant are at p < 0.01 against the baseline tab. 2. Setup: BART-large backbone for CNN/DailyMail and NYT; PEGASUS for XSum; 4 RTX 3090 GPUs, about 20 hours per epoch app. B.

NYT results are the authors' own, because their preprocessing differs from prior work tab. 2.

Further results, all on CNN/DailyMail unless stated:

• Weight sweep (R-1 / R-2 / R-L, gamma varied) tab. 3: gamma 0 (BART) 44.29 / 21.17 / 41.09; 0.1: 45.08 / 21.63 / 41.71; 1: 46.01 / 22.22 / 42.68; 2: 46.36 / 22.79 / 43.07; 5: 46.91 / 23.03 / 43.63; 10: 47.22 / 23.31 / 43.94; 100: 47.78 / 23.55 / 44.57; 1000: 46.83 / 22.17 / 43.68; contrastive only (BRIO-Ctr): 47.28 / 22.93 / 44.15. The paper says performance rises with gamma but cross-entropy is still needed for generation §5.3.
• Loop tab. 4: BRIO-Loop 48.01 / 23.80 / 44.67, against BRIO-Mul 47.78 / 23.55 / 44.57.
• Beam width (R-1; BART vs BRIO-Mul) tab. 5: beams 4: 44.29 vs 47.78; 10: 43.83 vs 47.98; 20: 43.53 vs 48.07; 50: 43.06 vs 48.18; 100: 42.79 vs 48.23. BART gets worse as beams grow; BRIO-Mul gets better.
• Metric used for ordering tab. 6: ordering by ROUGE gives 47.78 / 23.55 / 44.57 with BERTScore 32.11; ordering by BERTScore gives 47.53 / 23.22 / 44.37 with BERTScore 32.59 (BART: 27.38 BERTScore).
• Rank correlation (Spearman, 16 candidates, ROUGE-1 as quality) tab. 8: BART .0470 on its own candidates and .1205 on PEGASUS candidates; BRIO-Mul .1839 and .2768.
• Calibration (ECE) tab. 9: CNN/DailyMail BART .4097 vs BRIO-Mul .2719; XSum PEGASUS .2369 vs BRIO-Mul .1423.
• Novel bigram ratio tab. 7: reference .4865, BART .0924, BRIO-Mul .2381.
• Few-shot tab. 11: CNN/DailyMail, BART 44.29 / 21.17 / 41.09 vs BRIO-Few 45.81 / 21.91 / 42.61 (100 examples); XSum, PEGASUS 47.46 / 24.69 / 39.53 vs BRIO-Few 47.95 / 24.89 / 39.71 (1000 examples).
• Case study tab. 10: BRIO-Mul stops producing the "click here" boilerplate that BART reproduces; the paper counts 96 BART outputs with it and none from BRIO-Mul §5.6.

Paper vs code

Paths are relative to the repository at code_commit. Every released file is covered by the citations below.

• XSum contrastive weight is about 1 in the code, not 100. Paper: gamma = 100 on all datasets tab. 13 eq. 10. The XSum config sets rank_weight = 10 [config.py#L47], mle_weight = 0.1 [config.py#L46] and scale = 0.01 [config.py#L56]. The training loop multiplies the scores by scale before the ranking loss [main.py#L432-L434] and combines the terms as rank_weight * ranking_loss + mle_weight * mle_loss [main.py#L439]. Working this out by hand (not run): a margin of 0.001 on 0.01·f equals 0.01 times a margin of 0.1 on f, so lambda = 0.1 is matched, but the contrastive term carries weight 10 × 0.01 = 0.1 against 0.1 for the MLE term, a ratio of 1 instead of 100 [model.py#L9-L24]. CNN/DailyMail is consistent: scale = 1 [config.py#L20], ratio 100 [config.py#L10-L11].
• Cross-entropy target. Paper: one-hot target eq. 4 for L_xent. Code: the MLE term uses label smoothing with 0.1 in both configs [main.py#L393-L396] [config.py#L26] [config.py#L62]. The code gives the reference 1 − epsilon + epsilon/N and every other token epsilon/N [label_smoothing_loss.py#L14-L17]. The paper's form gives others beta/(N−1) eq. 5 and does not state beta.
• ROUGE used for ordering. Paper: "ROUGE" without a variant §3. Code: CNN/DailyMail uses the mean of ROUGE-1, ROUGE-2 and ROUGE-Lsum F-scores [preprocess.py#L65-L67]; XSum uses the harmonic mean of ROUGE-1 and ROUGE-2 [preprocess.py#L60-L63]. The scoring reads the PTB-tokenised files [preprocess.py#L19-L20].
• Candidate generation. Paper: diverse beam search, 16 candidates app. B; no beam numbers. Code, CNN/DailyMail: 16 beams in 16 groups, diversity penalty 1.0, return 16, length penalty 2.0, lower-cased source [gen_candidate.py#L13-L18] [gen_candidate.py#L29-L33]. Code, XSum: 128 beams in 16 groups, diversity penalty 0.1, length penalty 0.6, keeps every 8th of the 128 outputs to make 16 [gen_candidate.py#L84-L86].
• Length penalty. Matches tab. 13: 2.0 for CNN/DailyMail [config.py#L28] and 0.6 for XSum [config.py#L64]. The paper does not define |S|; the code counts the target tokens after the start token, including the end token [model.py#L68-L69] [model.py#L80].
• Gold-summary ranking term. Not in the paper. Disabled in the config with gold_weight = 0 [config.py#L9] [config.py#L45], so the gold term adds nothing to the loss [model.py#L27-L34] [main.py#L434].
• Optimizer and clipping. Adam with PyTorch defaults; no betas, epsilon or weight decay given [main.py#L397]. Gradient clipping is off (grad_norm = 0) [config.py#L15] [config.py#L51] [main.py#L447]. The paper does not state either.
• Learning-rate schedule. Matches eq. 13: max_lr = 2e-3 [config.py#L19] [config.py#L55], warmup 10000 [config.py#L13] [config.py#L49], and the step counter advances once per optimizer update [main.py#L451-L455].
• Batch size and accumulation. Not in the paper. Code: CNN/DailyMail batch 1 with 8 accumulation steps [config.py#L3] [config.py#L6]; XSum batch 2 with 4 steps [config.py#L39] [config.py#L42]. The README runs CNN/DailyMail on four GPUs [README.md#L150], and the code wraps the model in DistributedDataParallel [main.py#L383-L384]. By arithmetic (not run), that is 32 documents per update on four GPUs.
• Checkpoint selection. Not in the paper beyond "the validation set" for lambda app. B. Code: every 1000 updates [config.py#L34] [config.py#L70], the model is kept when a dev score improves: 1 − (R1·R2 + RLsum)/3 for CNN/DailyMail and 1 − 2·R1·R2/(R1+R2) for XSum [main.py#L411-L416] [main.py#L474-L500].
• ROUGE in training and validation. Paper: Perl ROUGE-1.5.5 app. B. The repository's cal_rouge.py uses the same flags [cal_rouge.py#L20]. Validation during training uses the pure-Python compare-mt ROUGE [main.py#L8] [main.py#L99], and the README says the values differ slightly [README.md#L160].
• Generation settings. Not in the paper beyond the 4-beam default tab. 5. Code: CNN/DailyMail 4 beams, output 55–140 tokens [config.py#L30-L31] [config.py#L35]; XSum 8 beams, 11–62 tokens [config.py#L66-L67] [config.py#L71]; no_repeat_ngram_size = 3 [main.py#L157].
• Base checkpoints. Match app. B: facebook/bart-large-cnn [config.py#L12] and google/pegasus-xsum [config.py#L48].
• Missing from the repository. No NYT configuration beyond "shares the CNN/DailyMail setting" [README.md#L141] and no NYT preprocessing [README.md#L80]. No few-shot script app. C. The loop experiment §5.3 cannot be run from the repository: gen_candidate.py loads only the pre-trained BART and PEGASUS checkpoints [gen_candidate.py#L9] [gen_candidate.py#L69].

Pitfalls

• License. The repository has no LICENSE file, so code_license is recorded as not stated. Some modelling files carry Apache-2.0 headers from Hugging Face Transformers [modeling_bart.py#L4-L6]. Cite the code by path and line only, and do not copy it.
• Offline cost. Candidate generation runs over the full training set before training starts. The paper calls this "a non-trivial overhead" §5.5. CNN/DailyMail has 287K training examples tab. 12; one epoch takes about 20 hours on four RTX 3090 GPUs app. B.
• Pairs per document. With 16 candidates, Eq. 8 has C(16, 2) = 120 pairwise terms per document, each needing two candidate scores. The scores come from one decoder pass per candidate with a shared encoder [modeling_bart.py#L1863-L1867].
• Truncation. Candidates are truncated to 120 tokens (CNN/DailyMail) or 80 tokens (XSum) before scoring, and sources to 1024 or 512 tokens [data_utils.py#L58] [data_utils.py#L76] [config.py#L24] [config.py#L60]. The paper does not mention truncation, so a long candidate is scored only on its first tokens.
• Special tokens. BART candidates keep the tokenizer's own start token. For PEGASUS, a pad token is prepended as the decoder start [data_utils.py#L78-L83].
• Ties. How ROUGE ties are ordered is not stated. The code's sort is stable, so ties keep their input order [data_utils.py#L69-L72].
• Contrastive-only models cannot generate. BRIO-Ctr is a scorer; keep the cross-entropy term if you need a generator §3.
• Longer beams help only the trained model. BART degrades as beams grow, so the beam-width gains apply to BRIO-Mul specifically tab. 5.
• Unnamed variant. The "Ours" row in the motivating table tab. 1 does not say whether it is BRIO-Ctr or BRIO-Mul.
• Two ROUGE implementations. Validation numbers from training are computed with a different ROUGE implementation from the Perl numbers in the paper, and the README says they differ slightly [README.md#L160].
• Length normalization is sensitive. alpha changes the scores directly and was tuned on validation for the rank-correlation study §5.3. The paper does not say how it was chosen for training.
• Casing. CNN/DailyMail candidates are generated from lower-cased input [gen_candidate.py#L18]. The released CNN/DailyMail checkpoint is uncased [README.md#L266], and a cased variant is listed separately [README.md#L218].

Lineage

• Builds on: BART (Lewis et al., 2020) and PEGASUS (Zhang et al., 2020) as the base generators §2 §5.1; margin-based ranking losses for candidate ordering (Hopkins and May, 2011; Zhong et al., 2020) §3; the combination of contrastive and cross-entropy losses following Edunov et al. (2018) §3. Their arXiv ids are not in this wiki, so none are linked.
• Compared against: BART, PEGASUS, GSum, ConSum, SeqCo, GOLD and SimCLS tab. 2 §5.1.
• Superseded by: none recorded in this wiki.