Source-linked AI summary
TIAO: Token Importance-Aware Policy Optimization for Text Summarization
Qixiu Li, Chenlong Bao, Xiang Zhu, Xiaoyong Li, Ruixin Cao, Shukai Chen, Zhenxiong Zhou
TL;DR
Summarization RL applies sequence-level rewards across tokens despite their differing relevance to source evidence. TIAO estimates token dependency using source masking, then reweights trajectory advantages and focuses updates on source-sensitive tokens. On CNN/DailyMail, its 7B policy achieves competitive quality against strong open-source and proprietary baselines, exceeding GPT-4 in overall score by 0.039.
Problem
Summarization RL commonly broadcasts sequence-level rewards across tokens, although tokens differ in source dependence and factual-evidence use.
Source: §2.3 · p. 2Method
TIAO estimates output-token dependency by comparing probabilities under original and masked sources, then uses dependency to reweight trajectory advantages and focus token updates.
Source: §1 · p. 1Results
On CNN/DailyMail, the 7B TIAO policy achieves competitive quality against strong open-source and proprietary baselines, exceeding GPT-4 in overall score by 0.039.
Source: §4 · p. 4, §3.2 · p. 4Takeaways & Limitations
The results suggest source-aware RL can be more effective than simply increasing model scale under the same evaluation setting.
Source: §3.2 · p. 4
Abstract
from arXiv · showhide
Text summarization requires models to condense content while preserving key qualities such as consistency and coherence. Large language models (LLMs) have shown strong performance on this task and can be further improved through reinforcement learning (RL). However, most existing methods apply reward signals directly to undifferentiated token sequences, overlooking the varying importance of individual tokens to word and sentence level quality in summarization. In this paper, we propose Token Importance-Aware Policy Optimization (TIAO), a novel reinforcement learning strategy that explicitly leverages token-importance awareness. Specifically, TIAO identifies core tokens based on token dependency and reweights a trajectory's advantage according to its overall dependencies. Experiments on the real world dataset show that our TIAO achieves highly competitive results, and that a 7B foundation model enhanced by TIAO performs comparably to GPT-4 and GPT-5-nano. Code is available at https://github.com/TechCloud-x/TIAO
1. INTRODUCTION
Summarization RL can broadcast sequence-level rewards across tokens despite their differing dependence on source evidence. TIAO addresses this mismatch by estimating token dependency and using it to reweight trajectory advantages and focus policy updates.
1. INTRODUCTION
- Summarization compresses a source while preserving salient content and qualities such as factual consistency; RL methods can optimize non-differentiable quality signals.Fluent summaries may still hallucinate or distort source-supported content.Source: §1 · p. 1
- Sequence-level summary rewards are commonly broadcast across token positions, leaving the policy unable to identify which tokens depend on the source or determine factual coverage.This can dilute useful gradients over low-information tokens.Source: §1 · p. 1
- Prior token-level RL work includes reasoning-trace feedback [13, 14], reward redistribution [15], entropy-based analysis [16], and optimization methods DAPO [17] and SAPO [18].The paper distinguishes summarization's source-evidence question from methods developed largely for mathematical or general reasoning.Source: Figure 1 · p. 1
- The paper targets source-dependency-aware credit assignment for multi-dimensional summarization rewards without training an additional token critic.Its goal is to distinguish source-grounded trajectories from language-prior shortcuts and apply stronger learning pressure to source-sensitive tokens.Source: §1 · p. 1
- TIAO estimates token dependency from source masking, reweights trajectory advantages by overall dependency, and focuses policy updates on source-sensitive tokens.The method targets improved summary quality and training stability on CNN/DailyMail.Source: §1 · p. 1
2. METHODOLOGY
TIAO redirects sequence-level summarization feedback toward source-dependent trajectories and tokens by estimating token dependency from source perturbations. It combines trajectory-level advantage shaping with token-level gradient selection while retaining the existing reward signal and avoiding an additional token critic.
2.1. Overview
- Source perturbations estimate how strongly each sampled output token depends on source evidence; TIAO uses those dependencies to reshape trajectory advantages and focus policy gradients on source-sensitive positions.The same generated tokens are rescored under original and perturbed sources.Source: §2.1 · p. 2
2.2. Problem Formulation
- Summarization is modeled as token-level autoregressive decisions with sequence-level feedback, valid-token masking, and scalar rewards aggregated from multidimensional evaluator scores.TIAO is independent of the specific evaluator and aggregation rule and adds no token critic.Source: §2.2 · p. 2
2.3. Multi-Dimensional Group-Relative Policy Optimization
- Group-relative optimization standardizes sampled summaries’ multidimensional rewards, aggregates them into scalar rewards, and broadcasts one trajectory advantage across valid tokens.Conventional GRPO-style training therefore gives factual content tokens, discourse markers, and low-information function words equal responsibility for sequence-level reward.Source: §2.3 · p. 2
2.4. Source-Dependency Estimation for Output Tokens
- Token dependency is estimated by masking source evidence while keeping the generated prefix fixed, then measuring changes in the sampled tokens’ probabilities under the rollout policy.Larger dependency scores indicate greater source sensitivity; clipping stabilizes extreme shifts without introducing a trainable critic.Source: §2.4 · pp. 2–3
2.5. Trajectory–Token Credit Assignment
- TIAO aggregates token dependencies into a trajectory score to shape advantages, then selects the top-κ proportion of valid tokens for policy-gradient updates.The token gate concentrates updates on source-sensitive positions rather than weakly grounded or generic positions.Source: §2.5 · p. 3
2.6. Training Objective
- The final objective combines trajectory advantage shaping and token focusing, excluding inactive positions from the gradient sum while preserving reward ordering and full-completion-length normalization.Dependency estimation, reward normalization, and token gating are recomputed for each rollout batch.Source: §2.6 · p. 3
3. EXPERIMENTS
On CNN/DailyMail, TIAO improves overall summarization scores over several RL baselines and achieves competitive quality against proprietary models. Ablations and a case study report stronger results with moderate source masking and later-stage gains over GRPO.
3.1. Experimental Settings
- Experiments use CNN/DailyMail, UniEval's four dimensions, and Qwen2.5-7B-Instruct as the initial policy for all RL methods.TIAO estimates token importance by comparing probabilities under original and 50%-masked sources, reweights trajectory advantages, and updates only the top 40% most important valid output tokens.Source: §3.1 · p. 3
3.2. Performance Comparison
- TIAO's overall score exceeds GRPO, HVO, DAPO, and SAPO by 0.038, 0.017, 0.051, and 0.050, respectively.Table 1 reports multi-dimensional CNN/DailyMail evaluation results.Source: Table 1 · p. 4, §3.2 · p. 4
- Against GPT-4, TIAO improves consistency, fluency, relevance, and overall score by 0.106, 0.022, 0.032, and 0.039, respectively; GPT-4 has the highest coherence score at 0.967.TIAO's overall score is also 0.063 higher than Qwen2.5 32B zero-shot.Source: §3.2 · p. 4, Table 1 · p. 4
3.3. Ablation Study
- The 50% masking setting achieves the best overall score (0.960) and highest consistency (0.946); 20% and 80% masking score 0.959 and 0.957 overall.The authors report that moderate perturbation provides more balanced source-dependency signals than weaker or more aggressive masking.Source: §3.3 · p. 4
3.4. Case Study
- GRPO improves overall reward faster early in training, but TIAO later has a stronger pre-clip gradient signal and reaches higher overall reward.TIAO estimates source dependency, reshapes trajectory advantages, and selects source-sensitive output tokens for updates.Source: §3.4 · p. 4, Figure 3 · p. 4
4. CONCLUSION
The paper addresses sequence-level reward being applied to undifferentiated tokens in RL summarization. It proposes TIAO, which uses source masking to estimate token importance and focuses updates on important output tokens; CNN/DailyMail experiments report competitive quality for its 7B policy.
4. CONCLUSION
- TIAO addresses the mismatch between sequence-level rewards and undifferentiated token updates, and achieves competitive quality against strong open-source and proprietary baselines with a 7B policy on CNN/DailyMail.It estimates source-dependent token importance through source masking, reweights trajectory advantages, and focuses updates on important output tokens.Source: §4 · p. 4