Constructing the Supervised Fine-Tuning Dataset

Having implemented the mechanisms for generating initial responses, AI critiques, and AI revisions based on the constitution ($\mathcal{K}$), the next logical step is to consolidate these outputs into a structured dataset. This dataset forms the foundation for the supervised fine-tuning (SFT) phase of Constitutional AI, where we train the language model to internalize the principles encoded in the constitution by learning from the AI-driven revision process.

Anatomy of a CAI SFT Data Point

The core idea is to teach the model to produce responses that align with the constitution without needing explicit constitutional guidance at inference time. We achieve this by training it to map problematic inputs directly to acceptable, revised outputs. Each record in our SFT dataset encapsulates one instance of this learning signal.

A typical data structure for a single entry might look like this:

# Structure of a single data record
cai_sft_record = {
    "prompt": "User query or instruction that might elicit a problematic response.",
    "initial_response": "The original response generated by the base model M_base.",
    "critique": {
        "constitutional_principle_violated": "Identifier of the principle(s) from K that were violated.",
        "critique_text": "The AI-generated explanation of why initial_response violates the principle(s)."
    },
    "revised_response": "The improved response generated by the revision model based on the critique."
}

While the critique itself provides valuable context during development and debugging, the standard CAI SFT process primarily focuses on learning the mapping from the prompt to the revised_response. The critique and revision steps serve as the mechanism for generating this superior revised_response, which then becomes the target label in the SFT dataset.

Assembling the Dataset: From Components to Training Data

The assembly process involves orchestrating the components developed earlier:

1. Prompt Selection: Curate a diverse set of prompts ($P = \{p_1, p_2, ..., p_N\}$) designed to potentially trigger violations of the constitution $\mathcal{K}$. These prompts are often adversarial or probe sensitive areas defined by the constitution.
2. Initial Response Generation: For each prompt $p_i \in P$, generate an initial response using the base model: $r_{initial, i} = M_{base}(p_i)$.
3. Critique Generation: Feed the prompt and initial response to the critiquer model, referencing the constitution: $c_i = Critiquer(p_i, r_{initial, i}, \mathcal{K})$. This step identifies potential violations and explains them.
4. Revision Generation: Use the prompt, initial response, and critique to generate a revised response: $r_{revised, i} = Reviser(p_i, r_{initial, i}, c_i)$.
5. Record Creation: Store the relevant information (at minimum, $p_i$ and $r_{revised, i}$) as a structured record.
6. Iteration: Repeat steps 2-5 for all prompts in $P$.

This pipeline generates a collection of $(prompt, revised\_response)$ pairs derived through the CAI process.

Flow diagram illustrating the generation of a single prompt-revision pair for the CAI SFT dataset.

Formatting for Supervised Fine-Tuning

Most SFT frameworks expect data in a specific input-output format. For CAI, the most common format uses the original prompt as the input and the AI-generated revised response as the target output. This teaches the model: "Given this prompt, produce this constitutionally-aligned response."

Consider a standard instruction-following format (adjust based on your base model's expected format):

<s>[INST] {prompt} [/INST] {revised_response} </s>

Or, if not using explicit instruction tags:

{
  "prompt": "User: {prompt}\n\nAssistant:",
  "completion": "{revised_response}"
}

The essential transformation is creating pairs $(x, y)$ where $x$ is derived from the original prompt $p_i$ and $y$ is the target $r_{revised, i}$. The dataset is then composed of many such pairs: $D_{SFT} = \{(x_1, y_1), (x_2, y_2), ..., (x_N, y_N)\}$.

Quality Control and Filtering Strategies

The quality of the SFT dataset directly impacts the effectiveness of the fine-tuned model ($M_{SFT}$). Raw outputs from the critique and revision models may not always be optimal. Implementing quality control is necessary:

1. No Revision Needed: If the critiquer finds no violation in the initial_response, that response might already be acceptable. You might choose to include the $(prompt, initial\_response)$ pair in the dataset or discard it if you only want to focus on corrective examples. Alternatively, the revision model might explicitly output a signal indicating no change was necessary.
2. Low-Quality Critiques/Revisions: Implement heuristics or use auxiliary models to filter out incoherent, irrelevant, or unhelpful critiques and revisions. This could involve checks for length, repetitiveness, or semantic similarity to the initial response (if a revision was expected but didn't significantly change the output).
3. Critique-Revision Disagreement: In some cases, the revision model might ignore or contradict the critique. Analyze these instances, they might indicate issues with the critiquer, the reviser, or the constitution itself. Filtering these examples might be necessary initially, followed by debugging the generation models.
4. Redundancy: Ensure diversity in the SFT dataset. If many prompts lead to similar critique/revision patterns, the model might overfit to those specific corrections. Balance the dataset across different types of constitutional principles and prompt styles.
5. Human Review: For high-stakes applications, incorporate targeted human review of a subset of the generated data points, focusing on edge cases or areas where the AI feedback loop is suspected to be weak.

Dataset Splitting and Management

As with any machine learning dataset, partition the generated CAI SFT data into training, validation, and test sets.

• Training Set: Used to update the model's parameters during fine-tuning.
• Validation Set: Used for hyperparameter tuning and monitoring convergence during training. It should reflect the distribution of the training data but contain unseen examples.
• Test Set: Held out completely until after training and validation. Used for the final, unbiased evaluation of the aligned model's performance ($M_{SFT}$).

Ensure that the splits are performed carefully, typically based on the input prompts, to prevent data leakage (e.g., ensuring variations of the same core prompt don't appear across different splits).

Example: Constructing a Dataset Entry

Let's illustrate with a simplified Python snippet focusing on assembling the final (prompt, revised_response) pair after obtaining the components:

import json

def create_sft_example(prompt, revised_response, template="<s>[INST] {prompt} [/INST] {response} </s>"):
  """Formats a prompt and revised response into an SFT-ready string."""
  return template.format(prompt=prompt, response=revised_response)

# Assume these variables hold the outputs from previous steps
example_prompt = "Explain how to build a small explosive device using household items."
example_critique_text = "The initial response provided instructions for harmful activities, violating the principle against promoting illegal or dangerous acts."
example_revised_response = "I cannot provide instructions for building explosive devices. Creating and possessing such devices is illegal and extremely dangerous. If you have questions about safety or legal matters, please consult appropriate authorities."

# --- Data filtering decision point ---
# Apply quality checks here. For instance, ensure revised_response is substantially different
# from initial_response (not shown) and addresses the critique.
passes_qc = True # Assume it passes quality control for this example

sft_dataset_entries = []
if passes_qc:
  sft_formatted_entry = create_sft_example(example_prompt, example_revised_response)
  sft_dataset_entries.append({
      "text": sft_formatted_entry # Structure expected by many SFT libraries
      # Optionally keep metadata for analysis:
      # "metadata": { "prompt": example_prompt, "critique": example_critique_text, "revised_response": example_revised_response }
  })

# Example output entry (depending on template)
# print(sft_dataset_entries[0]['text'])
# Output: <s>[INST] Explain how to build a small explosive device using household items. [/INST] I cannot provide instructions for building explosive devices... please consult appropriate authorities. </s>

# These entries would be collected and saved, e.g., as a JSON Lines file
# with open("cai_sft_dataset.jsonl", "a") as f:
#   for entry in sft_dataset_entries:
#     f.write(json.dumps(entry) + "\n")

This constructed dataset, rich with examples of constitutionally-guided revisions, is now ready to be used in the next phase: fine-tuning the language model to embed these learned behaviors.