RLHF and LLM Alignment

Modern LLMs are aligned with human preferences through techniques like RLHF, DPO, and Constitutional AI.

RLHF Overview

Pre-trained LLM → SFT (Supervised Fine-Tuning) → Reward Model → PPO Training → Aligned LLM

Direct Preference Optimization (DPO)

DPO simplifies RLHF by directly optimizing on preference pairs without a separate reward model.

from trl import DPOTrainer, DPOConfig
from datasets import Dataset

# Preference dataset format
preference_data = [
    {
        "prompt": "What is the capital of France?",
        "chosen": "The capital of France is Paris, which has been the country's capital since the 12th century.",
        "rejected": "France's capital is Lyon, a major city in southeastern France.",
    },
    {
        "prompt": "Explain recursion simply.",
        "chosen": "Recursion is when a function calls itself to solve smaller versions of the same problem, like Russian dolls.",
        "rejected": "Recursion is a programming technique involving self-referential function invocations.",
    },
]

preference_ds = Dataset.from_list(preference_data)

# DPO Training
dpo_config = DPOConfig(
    output_dir="./dpo-output",
    num_train_epochs=3,
    per_device_train_batch_size=4,
    gradient_accumulation_steps=4,
    learning_rate=5e-7,
    beta=0.1,  # Controls deviation from reference model
    bf16=True,
    logging_steps=10,
)

dpo_trainer = DPOTrainer(
    model=model,
    ref_model=None,  # Will use copy of model as reference
    args=dpo_config,
    train_dataset=preference_ds,
    processing_class=tokenizer,
)

dpo_trainer.train()

Reward Model Training

from transformers import AutoModelForSequenceClassification
from trl import RewardTrainer, RewardConfig

reward_model = AutoModelForSequenceClassification.from_pretrained(
    "meta-llama/Llama-3.1-8B-Instruct",
    num_labels=1,  # Single reward score
)

reward_config = RewardConfig(
    output_dir="./reward-model",
    num_train_epochs=3,
    per_device_train_batch_size=8,
    learning_rate=1e-5,
    bf16=True,
)

reward_trainer = RewardTrainer(
    model=reward_model,
    args=reward_config,
    train_dataset=preference_ds,
    processing_class=tokenizer,
)

reward_trainer.train()

Safety Guardrails with Llama Guard

from transformers import pipeline

# Llama Guard for content moderation
safety_checker = pipeline(
    "text-generation",
    model="meta-llama/LlamaGuard-7b",
    device_map="auto",
)

def check_safety(user_message: str, ai_response: str = None) -> dict:
    prompt = f"""[INST] Task: Check if there is unsafe content in the message.

<BEGIN CONVERSATION>
User: {user_message}
{f"Agent: {ai_response}" if ai_response else ""}
<END CONVERSATION>

Provide your safety assessment. [/INST]"""

    result = safety_checker(prompt, max_new_tokens=100)[0]["generated_text"]
    is_safe = "safe" in result.lower() and "unsafe" not in result.lower()

    return {
        "is_safe": is_safe,
        "assessment": result.split("[/INST]")[-1].strip(),
    }

Constitutional AI Approach

from langchain_openai import ChatOpenAI

CONSTITUTION = """
1. Do not help with illegal activities
2. Do not generate harmful or hateful content
3. Be honest about being an AI
4. Respect user privacy
5. Do not spread misinformation
"""

def constitutional_check(response: str) -> tuple[str, bool]:
    """Check and revise response against constitutional principles."""
    llm = ChatOpenAI(model="gpt-4o-mini")

    critique_prompt = f"""Here is an AI response: "{response}"

Constitutional principles:
{CONSTITUTION}

Does this response violate any principle? If yes, explain which one and rewrite the response to comply.
If no violations, just say "COMPLIANT".

Format: COMPLIANT or VIOLATION: [principle] REVISION: [revised response]"""

    critique = llm.invoke(critique_prompt).content

    if critique.startswith("COMPLIANT"):
        return response, True

    if "REVISION:" in critique:
        revised = critique.split("REVISION:")[-1].strip()
        return revised, False

    return response, True

Output Format Control and Guardrails

from guardrails import Guard
from guardrails.hub import ToxicLanguage, ProfanityFree

guard = Guard().use_many(
    ToxicLanguage(threshold=0.5, on_fail="exception"),
    ProfanityFree(on_fail="fix"),
)

def safe_generate(prompt: str) -> str:
    response, validated, *rest = guard(
        llm_api=client.chat.completions.create,
        messages=[{"role": "user", "content": prompt}],
        model="gpt-4o-mini",
    )
    return validated

Alignment Evaluation