Using CAI Outputs as Input for RLAIF

The Constitutional AI (CAI) process, particularly its supervised learning (SL) phase, generates significant artifacts beyond just the fine-tuned model. This phase involves prompting an LLM, having an AI critique the response based on a constitution, and then having an AI revise the response according to the critique. This results in a dataset rich with examples of constitutionally-aligned behavior refinement. Instead of treating the CAI SL phase and the RLAIF phase as separate steps, we can strategically use the outputs of the former to enhance the latter. This integration aims to infuse the principle-based guidance from CAI into the preference-based optimization of RLAIF, potentially leading to more effective alignment.

Starting RLAIF with a CAI-Tuned Model

One straightforward integration method is to use the LLM fine-tuned during the CAI SL phase as the starting point for RLAIF training. Recall that the CAI SL phase fine-tunes the model on pairs of critiques and revised responses, effectively teaching the model to internalize the constitutional principles through self-correction examples.

Instead of initializing RLAIF with a generic pre-trained or instruction-tuned model, starting with the CAI-tuned model offers potential advantages:

1. Better Initialization: The model already possesses some degree of constitutional alignment. Its initial outputs during RLAIF's online generation phase are less likely to drastically violate the principles, potentially focusing the RL exploration on refining helpfulness and harmlessness within already acceptable bounds.
2. Faster Convergence: Because the starting policy is closer to the desired aligned state (as implicitly defined by the constitution and refined by RLAIF's AI preferences), the RL process might converge more quickly or reach a higher level of alignment compared to starting from a less aligned model. The Proximal Policy Optimization (PPO) algorithm used in RLAIF penalizes large deviations from the reference policy; starting with a better reference policy can streamline optimization.
3. Reduced Catastrophic Forgetting: While still a concern, starting with a CAI-tuned model might slightly mitigate the risk of the RL process unlearning the constitutional constraints, as these are already partially baked into the model's weights.

The implementation involves simply taking the final checkpoint from the CAI SL fine-tuning stage and loading these weights as the initial policy $\pi_{\theta}$ and potentially the reference policy $\pi_{ref}$ at the beginning of the RLAIF PPO training loop.

Transforming CAI Data into Preference Pairs

The dataset generated during the CAI SL phase typically contains tuples like: (prompt, initial_response, critique, revised_response). This structured data can be repurposed to create preference pairs suitable for training the RLAIF preference model (PM).

The core idea is to treat the constitutionally revised response as preferred over the initial response for the given prompt. The critique-and-revision process, guided by the constitution, acts as an implicit preference judgment: the model should have produced something like the revised response instead of the initial one.

So, for a given prompt $p$, initial response $r_{initial}$, and revised response $r_{revised}$, we can generate a preference pair $(r_{chosen}, r_{rejected})$ where $r_{chosen} = r_{revised}$ and $r_{rejected} = r_{initial}$.

This creates a dataset $D_{CAI\_prefs} = \{ (p, r_{revised}, r_{initial}) \}$. This dataset can then be combined with preference data generated by the standard RLAIF AI labeler (which compares two outputs, $r_A$ and $r_B$, from the current policy).

Considerations:

• Data Augmentation: $D_{CAI\_prefs}$ can augment the primary RLAIF preference dataset, providing examples grounded explicitly in constitutional principles.
• Filtering: Not all CAI revisions necessarily represent strong preferences relevant to the RLAIF objective. Minor stylistic edits might be less important than corrections related to safety or bias. You might filter the CAI data to include only revisions triggered by specific constitutional principles (e.g., those related to harm reduction).
• Weighting: When training the preference model on a combined dataset, you might assign different weights to samples originating from $D_{CAI\_prefs}$ versus those from the standard AI labeler, depending on the confidence in each source or the desire to emphasize constitutional adherence.
• Preference Strength: The binary preference $(r_{revised} \succ r_{initial})$ might be too simplistic. The severity of the critique or the magnitude of the revision could potentially inform a preference strength score, although this adds complexity.

The preference model is then trained to maximize the likelihood of the chosen responses and minimize the likelihood of the rejected responses, using a loss function like:

$$\mathcal{L}_{PM} = - \mathbb{E}_{(p, r_{chosen}, r_{rejected}) \sim D_{combined}} [\log(\sigma(f_{PM}(p, r_{chosen}) - f_{PM}(p, r_{rejected})))]$$

where $f_{PM}$ is the preference model's scoring function, $\sigma$ is the sigmoid function, and $D_{combined}$ is the union of RLAIF-generated and CAI-derived preference data.

Seeding the RLAIF Experience Buffer

Another technique involves directly injecting successful CAI examples into the experience buffer used during the PPO phase of RLAIF. The PPO algorithm learns by sampling prompts, generating responses using the current policy $\pi_{\theta}$, evaluating these responses using the reward model (trained on preferences), and updating the policy.

We can augment the experience buffer with tuples of $(p, r_{revised})$ derived from the CAI SL dataset. When the PPO algorithm samples experiences for policy updates, it will occasionally draw these high-quality, constitutionally-vetted examples.

How it helps:

• Guidance: Provides direct examples of desired behavior early in the RL process.
• Stability: Can potentially stabilize training by ensuring the policy updates are influenced by known-good trajectories, especially if the initial policy or reward model is noisy.

Implementation requires modifying the experience collection or sampling mechanism in the PPO trainer to include these pre-computed examples alongside the online generations. Care must be taken to assign appropriate reward scores and potentially advantage estimates to these injected samples, likely using the trained RLAIF preference/reward model to score $r_{revised}$ for prompt $p$.

Integration points for leveraging CAI outputs in RLAIF: (1) Initializing the RLAIF policy with the CAI-tuned model. (2) Using CAI (prompt, initial, revised) data to generate preference pairs for PM training. (3) Seeding the PPO experience buffer with CAI examples. (4) Potentially using the CAI-tuned model itself within the AI Preference Labeler.

By strategically reusing the outputs of the CAI SL phase, RLAIF training can be initialized with a more aligned model, benefit from constitutionally grounded preference data, and potentially converge faster or to a better state. This integration transforms the CAI process from just a pre-training step into a source of valuable data and initialization for the subsequent reinforcement learning phase. However, careful consideration must be given to data selection, weighting, and potential conflicts between the implicit preferences learned by RLAIF and the explicit principles encoded during CAI.