Technical Summary: Semi-Offline Reinforcement Learning for Optimized Text Generation

Problem Statement

Reinforcement Learning (RL) for text generation currently faces a dichotomy between online and offline settings, each with significant trade-offs:

• Online RL: Allows the language model (LM) to fully explore the environment by sampling from its output distribution. While this provides a comprehensive understanding of the final objective (e.g., ROUGE, human feedback), it incurs prohibitive optimization costs. Specifically, optimizing with $K$ samples requires $K \times T$ forward propagations (FPs) through the model, where $T$ is the sequence length. This is often impractical for large pretrained models.
• Offline RL: Utilizes a static dataset (e.g., ground truth or beam-search outputs) to eliminate the need for autoregressive generation during training. This reduces the cost from $KT$ FPs to $K$ FPs. However, it sacrifices exploration capability; the model only sees reward signals for specific static data points, preventing it from discovering optimal generation directions outside the static distribution.

The fundamental research question addressed is: Can we refine the RL setting to achieve effective exploration with minimum optimization cost?

Methodology: Semi-Offline RL

The authors propose Semi-Offline RL, a novel paradigm that bridges the gap between online and offline settings by defining a continuous design space controlled by a hyperparameter $p_m \in [0, 1]$.

1. Semi-Offline MDP Formulation

The method is grounded in a Markov Decision Process (MDP) where a sample is composed by mixing tokens generated by the language model and tokens from a static dataset.

• State ($s_t$): Includes the input $x$, the partially derived output, and a binary mask sequence $M_0, \dots, M_t$.
• Transition: At each step $t$, a Bernoulli variable $M_{t+1} \sim \text{Bernoulli}(p_m)$ determines the source of the next token:
  • If $M_{t+1} = 1$: The agent samples a token $\hat{y}_t$ from the LM (exploration).
  • If $M_{t+1} = 0$: The agent takes a token $y_t$ from the static dataset (exploitation).
• Cost Analysis: In a fully observable scenario, the expected number of FPs is $E(C_t) = t \cdot p_m$. By setting $p_m \in (0, 1)$, the method balances exploration capability with an intermediate optimization cost.

2. Optimal RL Setting: Masked Observations

To achieve the minimum optimization cost (specifically 1 FP per input regardless of sample count) while maintaining exploration, the authors move to a Partially Observable MDP (POMDP) formulation.

• Constraint: To enable parallel generation (1 FP), the agent's action cannot depend on tokens generated simultaneously by the model.
• Optimal Observation ($o^*_t$): The authors prove that the optimal observation satisfies three desirable properties (DP1-DP3):
  1. Minimum Optimization Cost: Achieved by 1 FP.
  2. Minimum Asymptotic Bias: Requires retaining all information necessary to predict the state.
  3. Minimum Overfitting Error Bound: Requires minimizing redundant information in the observation.
• Implementation: The optimal setting is realized by Masked Observations. The input to the LM consists of the source $x$ and the target sequence where tokens generated by the model are replaced by a special mask token $[M]$, while tokens from the static dataset are visible.
  • $y^M_t = [M]$ if $M_t=1$ (model generates).
  • $y^M_t = y_t$ if $M_t=0$ (static data used).
  • This allows the model to explore $K$ different trajectories simultaneously using a single forward pass, as the observation remains constant across samples.

3. Optimization Strategy

• Loss Function: The model is optimized using a policy gradient approach (REINFORCE with baseline). The loss is decomposed into token-wise rewards, allowing efficient estimation of the advantage of generating specific tokens under the masked observation.
• Training Protocol:
  1. Pre-training: Fine-tune the pretrained LM on ground truth with some tokens masked (M-FT) to adapt the model to the masked input format.
  2. RL Training: Minimize a combined loss $L = L_{MLE} + \lambda L_{RL}$, where $L_{MLE}$ prevents policy drift from the original dataset, and $L_{RL}$ optimizes the specific objective (e.g., ROUGE).

Key Contributions

1. Definition of Semi-Offline RL: The paper defines a theoretical design space that smoothly transitions between offline and online RL, providing a foundation for comparing different settings based on exploration capability and optimization cost.
2. Theoretical Optimality: The authors derive the theoretically optimal RL setting (Masked Observations) that satisfies minimum optimization cost, minimum asymptotic bias, and minimum overfitting error bounds.
3. Efficient Implementation: They demonstrate that this optimal setting can be implemented by mixing static data with mask tokens, enabling the exploration of $K$ samples with only 1 forward propagation.
4. Empirical Validation: Extensive experiments show the method yields comparable or superior performance to state-of-the-art methods (including online Actor-Critic and offline BRIO) while significantly reducing training time and computational cost.

Experimental Results

The method was evaluated on four text generation tasks: CNN/DM and SAMSum (summarization), SQuAD (question generation), and XSum (extreme summarization).

• Performance vs. Cost: The proposed method ("OURS") achieved the best or comparable results to state-of-the-art methods (e.g., BRIO, AVG, SC, AC) while requiring only $N$ forward propagations (where $N$ is the number of instances), whereas multi-sample online methods required $N \times T$ or $N \times K \times T$ propagations.
• Efficiency: On the SQuAD and SAMSum datasets, the semi-offline method demonstrated the fastest optimization speed and lowest training time consumption compared to BRIO and AVG.
• Ablation Studies:
  • Removing mask information (violating Lemma 1) led to performance drops due to inaccurate reward estimation.
  • Adding redundant information (violating Lemma 2) led to overfitting and lower test set performance.
  • Using "harder" static datasets (DATA-, lower metric candidates) for exploration yielded better results than using high-quality candidates (DATA+), suggesting that exploring the vicinity of sub-optimal points helps identify improvement directions.

Significance and Claims

The paper claims that Semi-Offline RL provides a theoretical foundation for comparing RL settings and offers a practical solution to the high computational cost of online RL. By utilizing a masked observation strategy, the method enables large pretrained language models to explore the environment and optimize for complex objectives (like human feedback or ROUGE) with the efficiency of offline training. The authors assert that this approach allows models to identify better improvement directions than purely offline methods while avoiding the prohibitive time costs of fully online exploration.