simple_dpo.py

import math

def main():
    print("Ultra-Simple DPO (Direct Preference Optimization) Example")
    print("=" * 60)
    
    # Our "model" is just a single weight (parameter)
    # The model calculates: score = input * weight
    weight = 0.1  # Start with a small random weight
    
    # Our preference data: (preferred, rejected) pairs
    # We want the model to learn that larger numbers are better
    preferences = [
        (8, 2),  # We prefer 8 over 2
        (6, 3),  # We prefer 6 over 3  
        (4, 1),  # We prefer 4 over 1
    ]
    
    print(f"Starting weight: {weight:.3f}")
    print("Training data (preferred, rejected):", preferences)
    print()
    
    # Training parameters
    learning_rate = 0.1
    beta = 1.0  # DPO temperature parameter
    
    print("Training with DPO...")
    print("Step | Weight | Loss")
    print("-" * 25)
    
    # Train for a few steps
    for step in range(10):
        total_loss = 0
        gradient = 0
        
        # Process each preference pair
        for preferred, rejected in preferences:
            # Calculate scores using our simple model
            preferred_score = preferred * weight
            rejected_score = rejected * weight
            
            # DPO loss calculation (simplified)
            # The magic: we want P(preferred > rejected) to be high
            logits = beta * (preferred_score - rejected_score)
            prob = 1 / (1 + math.exp(-logits))  # Sigmoid function
            
            # Loss: negative log probability of preferring the preferred item
            loss = -math.log(prob + 1e-8)  # Small epsilon to avoid log(0)
            total_loss += loss
            
            # Calculate gradient (how to update the weight)
            # This pushes preferred scores up and rejected scores down
            # Calculate gradient contribution for this preference pair:
            grad_contribution = beta * (preferred - rejected) * (prob - 1)
            gradient += grad_contribution
        
        # Update the weight using gradient descent
        weight = weight - learning_rate * gradient
        
        print(f"{step:4d} | {weight:6.3f} | {total_loss:6.3f}")
    
    print()
    print("Training complete!")
    print(f"Final weight: {weight:.3f}")
    print()
    
    # Test the trained model
    print("Testing the trained model:")
    test_inputs = [1, 2, 3, 4, 5, 6, 7, 8]
    
    for x in test_inputs:
        score = x * weight
        print(f"Input: {x} -> Score: {score:.3f}")
    
    print()
    print("Notice: The model learned to give higher scores to larger numbers!")
    print("This is exactly what DPO does - it learns preferences from comparisons.")

if __name__ == "__main__":
    main()