Reinforcement Learning for Language Models: Techniques and Implementations

Introduction to Reinforcement Learning for Language Models

Reinforcement learning (RL) has been making waves in the field of artificial intelligence, particularly in the domain of large language models (LLMs). Recent advancements by OpenAI with their O1 and O3 series, as well as DeepSeek's V3 reasoning models, have showcased the potential of RL techniques in improving model performance. This article delves into the theory and practical implementation of various RL methods, including GRPO (Group Relative Policy Optimization), PPO (Proximal Policy Optimization), ORPO (Odds Ratio Policy Optimization), and standard supervised fine-tuning.

We'll walk through examples of improving the performance of a small model, specifically the Llama 1B, on the GSM-8K dataset (Grade School Math 8K). This dataset consists of grade school math problems and serves as an excellent benchmark for assessing reasoning capabilities in language models.

Setting Up the Environment

Before diving into the RL techniques, it's crucial to set up a proper environment for experimentation. Here's a step-by-step guide to getting started:

1. Choose a suitable GPU: For this experiment, an H100 GPU was selected for its speed. When using a single GPU, opt for the PCIe version as it's more cost-effective than the NVLink version.

2. Set up the development environment: Use Visual Studio Code with SSH capabilities to connect to your GPU instance.

3. Clone the repository: For this experiment, a private repository containing advanced fine-tuning scripts was used. However, you can create your own scripts or use open-source alternatives.

4. Install dependencies: Key libraries include:

   • datasets for managing datasets
   • requests for making API calls
   • light_llm for interfacing with various LLM APIs
   • sentencepiece for tokenization
   • sgLang for efficient batch generation

5. Set up environment variables: Create a .env file to store API keys and other sensitive information.

Understanding the GSM-8K Dataset

The GSM-8K dataset is a collection of grade school math questions, each accompanied by a step-by-step solution and a final answer. It consists of approximately 7,500 training examples and 1,300 test examples. For our experiments, we'll use a subset of 100 test examples to speed up the evaluation process.

Here's an example from the dataset:

Question: Natalia sold clips to 48 of her friends in April, and then she sold half as many clips in May. How many clips did Natalia sell altogether in April and May?

Answer: Let's approach this step-by-step:
1) First, we need to find out how many clips Natalia sold in April.
   - We're told she sold clips to 48 friends in April.
   - So, in April, Natalia sold 48 clips.

2) Next, we need to calculate how many clips she sold in May.
   - We're told she sold half as many clips in May as she did in April.
   - Half of 48 is 48 ÷ 2 = 24
   - So, in May, Natalia sold 24 clips.

3) Finally, we need to add the number of clips sold in both months.
   - April: 48 clips
   - May: 24 clips
   - Total: 48 + 24 = 72 clips

Therefore, Natalia sold 72 clips altogether in April and May.

#### 72

Baseline Performance Evaluation

Before applying any RL techniques, it's essential to establish a baseline performance for our model. We'll use the Llama 3.2B instruct model as our starting point. Here's how we evaluate the model:

1. Run inference on the test set (100 examples)
2. Calculate two key metrics:
   • Pass@K: The percentage of questions where the model gets the correct answer at least once in K attempts.
   • Majority@K: The percentage of questions where the model gets the correct answer in the majority of K attempts.

For our baseline, we'll use K=8, meaning we'll sample 8 responses for each question.

Here's a snippet of the evaluation script:

def check_answer(question, model_answer, ground_truth, enforce_think_tags=False):
    # Extract the numerical answer from the model's response
    extracted_answer = extract_number(model_answer)
    extracted_ground_truth = extract_number(ground_truth)
    
    if extracted_answer == extracted_ground_truth:
        return True, "Exact numerical match"
    
    # If exact match fails, use Gemini API as a backup
    return check_with_gemini(question, model_answer, ground_truth)

def evaluate_model(model, test_set, num_samples=8):
    correct_count = 0
    majority_count = 0
    
    for question in test_set:
        correct_answers = 0
        for _ in range(num_samples):
            answer = model.generate(question)
            if check_answer(question, answer, question['answer']):
                correct_answers += 1
        
        if correct_answers > 0:
            correct_count += 1
        if correct_answers > num_samples // 2:
            majority_count += 1
    
    pass_at_k = correct_count / len(test_set)
    majority_at_k = majority_count / len(test_set)
    
    return pass_at_k, majority_at_k

After running the baseline evaluation, we obtained the following results:

• Pass@8: 79%
• Majority@8: 28%

These numbers serve as our reference point for improvement using RL techniques.

Supervised Fine-Tuning (SFT)

The first technique we'll explore is supervised fine-tuning. This method involves training the model on a dataset of correct answers, aiming to improve its performance through traditional cross-entropy loss.

Here's an overview of the SFT process:

1. Generate a dataset of correct answers from the training set
2. Format the data for fine-tuning
3. Apply LoRA (Low-Rank Adaptation) to the model
4. Train the model using cross-entropy loss
5. Evaluate the fine-tuned model on the test set

Here's a simplified version of the SFT training script:

from transformers import Trainer, TrainingArguments
from peft import get_peft_model, LoraConfig

def format_sft_data(example):
    prompt = f"System: You are a helpful math assistant.\nUser: {example['question']}\nAssistant:"
    response = f"<think>{example['answer']}</think>The answer is {extract_number(example['answer'])}"
    return {"input_ids": tokenizer(prompt + response, truncation=True, max_length=max_seq_length)}

# Load and format the dataset
dataset = load_dataset("gsm8k", "main")
train_dataset = dataset["train"].map(format_sft_data)

# Apply LoRA
config = LoraConfig(r=64, lora_alpha=32, target_modules=["q_proj", "v_proj"])
model = get_peft_model(model, config)

# Set up training arguments
training_args = TrainingArguments(
    output_dir="./results",
    num_train_epochs=1,
    per_device_train_batch_size=4,
    gradient_accumulation_steps=8,
    learning_rate=1e-5,
    fp16=True,
    logging_steps=100,
    save_strategy="epoch",
)

# Train the model
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    data_collator=lambda data: {"input_ids": torch.stack([f["input_ids"] for f in data])},
)
trainer.train()

# Save the fine-tuned model
model.save_pretrained("./sft_model")

After applying SFT, we evaluated the model again and obtained the following results:

• Pass@8: 79% (no significant change)
• Majority@8: 34% (6% improvement)

While we didn't see an improvement in the Pass@8 metric, the increase in Majority@8 suggests that the model became more consistent in providing correct answers.

Odds Ratio Policy Optimization (ORPO)

ORPO is a more advanced RL technique that combines supervised fine-tuning with a preference optimization approach. It aims to shift the model's probabilities towards chosen (correct) answers and away from rejected (incorrect) answers.

Here's an overview of the ORPO process:

1. Generate pairs of chosen and rejected answers from the training set
2. Format the data for ORPO training
3. Apply LoRA to the model
4. Train the model using a combined loss function (cross-entropy + odds ratio)
5. Evaluate the ORPO-trained model on the test set

Here's a simplified version of the ORPO training script:

from transformers import Trainer, TrainingArguments
from peft import get_peft_model, LoraConfig

def format_orpo_data(example):
    prompt = f"System: You are a helpful math assistant.\nUser: {example['question']}\nAssistant:"
    chosen = f"<think>{example['chosen_answer']}</think>The answer is {extract_number(example['chosen_answer'])}"
    rejected = f"<think>{example['rejected_answer']}</think>The answer is {extract_number(example['rejected_answer'])}"
    return {
        "prompt": prompt,
        "chosen": chosen,
        "rejected": rejected
    }

# Load and format the dataset
dataset = load_dataset("gsm8k", "main")
train_dataset = dataset["train"].map(format_orpo_data)

# Apply LoRA
config = LoraConfig(r=64, lora_alpha=32, target_modules=["q_proj", "v_proj"])
model = get_peft_model(model, config)

# Set up training arguments
training_args = TrainingArguments(
    output_dir="./results",
    num_train_epochs=1,
    per_device_train_batch_size=2,
    gradient_accumulation_steps=16,
    learning_rate=5e-6,
    fp16=True,
    logging_steps=100,
    save_strategy="epoch",
)

# Custom ORPO loss function
def orpo_loss(logits_chosen, logits_rejected):
    loss_ce = F.cross_entropy(logits_chosen, torch.argmax(logits_chosen, dim=-1))
    odds_ratio = torch.exp(logits_chosen - logits_rejected)
    loss_orpo = -torch.log(odds_ratio / (1 + odds_ratio)).mean()
    return loss_ce + 0.1 * loss_orpo

# Train the model
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    data_collator=lambda data: {
        "input_ids": torch.stack([tokenizer(d["prompt"] + d["chosen"])['input_ids'] for d in data]),
        "labels": torch.stack([tokenizer(d["prompt"] + d["rejected"])['input_ids'] for d in data]),
    },
    compute_loss=orpo_loss,
)
trainer.train()

# Save the ORPO-trained model
model.save_pretrained("./orpo_model")

After applying ORPO, we evaluated the model once again and obtained the following results:

• Pass@8: 75% (4% decrease)
• Majority@8: 35.5% (1.5% improvement over SFT, 7.5% over baseline)

The results from ORPO training are mixed. While we see a slight improvement in the Majority@8 metric compared to SFT, there's a small decrease in the Pass@8 metric. This suggests that ORPO might be making the model more consistent in its correct answers but potentially at the cost of some exploration capability.

Analysis and Future Directions

Based on our experiments with SFT and ORPO, we can draw several conclusions and identify areas for further investigation:

1. Consistency vs. Exploration: Both SFT and ORPO showed improvements in the Majority@8 metric, indicating that these techniques are effective at making the model more consistent in providing correct answers. However, the lack of improvement (or slight decrease) in the Pass@8 metric suggests that we might be limiting the model's ability to explore diverse solutions.

2. Hyperparameter Tuning: The performance of both SFT and ORPO can be sensitive to hyperparameters such as learning rate, batch size, and the balance between different loss components. Further experimentation with these parameters might yield better results.

3. Model Size Limitations: It's possible that the relatively small size of the Llama 1B model is limiting the potential gains from these RL techniques. Experimenting with larger models might reveal more significant improvements.

4. Multi-step Reinforcement Learning: Our current approach uses a single step of RL. Implementing multiple rounds of data generation and fine-tuning could lead to more substantial improvements over time.

5. Verifiable Self-correction: An interesting direction for future work is to implement a mechanism for the model to attempt self-correction. This could involve training the model to recognize when it might be wrong and to generate alternative answers.

6. Exploration of Other RL Techniques: While we focused on SFT and ORPO, other techniques like GRPO (Group Relative Policy Optimization) and PPO (Proximal Policy Optimization) could be explored and compared.

7. Application to More Complex Datasets: The GSM-8K dataset, while useful, is relatively simple. Applying these techniques to more complex datasets like ARC (AI2 Reasoning Challenge) could provide insights into their effectiveness on harder reasoning tasks.

8. Analysis of Generated Responses: A more detailed analysis of the model's outputs, including the length and structure of responses, could provide insights into how the RL techniques are affecting the model's behavior.

Conclusion

Reinforcement learning techniques offer promising avenues for improving the performance of language models on reasoning tasks. Our experiments with supervised fine-tuning and ORPO on the Llama 1B model and the GSM-8K dataset have shown modest improvements, particularly in the consistency of correct answers.

However, these results also highlight the complexities involved in applying RL to language models. The trade-offs between consistency and exploration, the sensitivity to hyperparameters, and the potential limitations of smaller models all present challenges that require further investigation.

As the field of AI continues to evolve, the refinement and development of RL techniques for language models will likely play a crucial role in advancing their capabilities. By building on the foundations laid out in this article and exploring new directions, researchers and practitioners can continue to push the boundaries of what's possible with language models and reinforcement learning.

Future work in this area should focus on scaling these techniques to larger models, exploring multi-step RL processes, and tackling more complex reasoning tasks. Additionally, the development of more sophisticated evaluation metrics that can capture nuanced improvements in model performance will be crucial for guiding future research and development efforts in this exciting field.

Article created from: https://youtu.be/C4HxJQ2QzWo?si=XBbAcuLbrTjvtPKV