Contrastive Decoding with PyTorch

An implementation of Contrastive Decoding using PyTorch and HuggingFace Transformers.

Overview

Contrastive Decoding is a technique that improves text generation quality by leveraging the differences between a large "expert" model and a smaller "amateur" model. The key insight is that by subtracting the amateur model's token probabilities from the expert's, we can emphasize tokens that represent the unique knowledge of the larger model.

This implementation uses:

• Qwen/Qwen2.5-3B-Instruct as the expert model
• Qwen/Qwen2.5-Coder-0.5B-Instruct as the amateur model

Features

• Token-level contrastive decoding implementation
• Plausibility threshold to ensure high-quality outputs
• Configurable contrast strength (alpha parameter)
• Optional context length limiting for the amateur model to improve efficiency

Installation

# Clone the repository
git clone https://github.com/FrigaZzz/contrastive-generation.git
cd contrastive-generation

# Install dependencies
pip install torch transformers

Usage

import torch
import transformers as tr
from main import contrastive_generation

# Load models
amateur = tr.AutoModelForCausalLM.from_pretrained('Qwen/Qwen2.5-Coder-0.5B-Instruct', torch_dtype=torch.float16)
expert = tr.AutoModelForCausalLM.from_pretrained('Qwen/Qwen2.5-3B-Instruct', torch_dtype=torch.float16)

# Move to GPU if available
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
amateur = amateur.to(device)
expert = expert.to(device)

# Create prompt
tokenizer = tr.AutoTokenizer.from_pretrained('Qwen/Qwen2.5-Coder-0.5B-Instruct')
prompt = tokenizer.apply_chat_template([
    {'role': 'system', 'content': 'You are a helpful assistant'},
    {'role': 'user', 'content': 'Write a short poem about machine learning.'}
], add_generation_prompt=True, tokenize=False)

# Generate text
generated_text = contrastive_generation(
    amateur=amateur,
    expert=expert,
    prompt=prompt,
    max_tokens=200,
    alpha=0.1,  # Contrast strength
    plausibility_threshold=0.05,  # Minimum probability for tokens
    amateur_context_length=1  # Use only last token for amateur context
)

print(generated_text)

How It Works

1. Expert Knowledge: The large model (expert) has more knowledge but can still produce repetitive or low-quality text
2. Amateur Mistakes: The small model (amateur) makes similar mistakes but more frequently
3. Contrastive Scoring: By computing expert_prob - alpha * amateur_prob, we emphasize tokens that the expert prefers but the amateur doesn't
4. Plausibility Constraint: We only consider tokens with a minimum probability from the expert model

Parameters

• alpha: Controls the contrast strength (default: 0.1)
• plausibility_threshold: Minimum probability for tokens from the expert model (default: 0.05)
• amateur_context_length: Number of tokens to use as context for the amateur model (default: 1)

Problem statement (from the Cursor ML Challenge)

Your goal is to implement Contrastive Decoding with HuggingFace transformers and PyTorch.

Your code should use Qwen/Qwen2.5-3B-Instruct as the large model and Qwen/Qwen2.5-Coder-0.5B-Instruct as the small model and be implemented in main.py.

Your code should be correct first, but also efficient. Implement the token-level algorithm, rather than the beam search algorithm.

In addition to implementing main.py, please answer the following questions:

1. What should you do if the two models have different tokenizers?

When models have different tokenizers, you need to align their vocabularies:

1. Choose one tokenizer space (typically the expert's)
2. Map tokens between vocabularies by converting to strings and back
3. Align probability distributions in the chosen token space Alternatively, use models from the same family that share tokenizers.

2. Do you think contrastive decoding is used in practice?

From what I've researched, contrastive decoding isn't super common in everyday applications, but it definitely has its place. Here's why:

Why it's not everywhere:

• It's expensive - Running two models at once eats up a lot more GPU memory and compute
• It's a bit complicated - Getting two different models to work together, especially with different tokenizers, is tricky
• Newer stuff exists - Techniques like speculative decoding can be faster while still improving quality

Where it probably makes sense:

• When you're generating longer content where staying on topic matters
• For premium applications where quality trumps cost
• In research settings where you're pushing for the best possible outputs

I couldn't find solid evidence of big companies using it in production, though since the paper came from Stanford, CMU, and Meta researchers, Meta might be using some version internally.

The quality improvement is real though - it does help reduce those annoying repetitions and topic drifts that plague even large models. For applications where that extra quality matters more than saving on compute costs, it's probably worth the trade-off.