Accueil › Glossaire › Rejection Sampling

Rejection Sampling (LLM)

Le rejection sampling, dans le contexte des LLM, est une technique d’alignement qui génère plusieurs réponses candidates à partir du modèle, évalue chacune avec un reward model, ne conserve que les meilleures, et utilise ces réponses filtrées comme données d’entraînement pour un fine-tuning supervisé du modèle.

Le concept

Le rejection sampling est né en statistique computationnelle comme méthode pour échantillonner à partir d’une distribution complexe quand on ne peut pas le faire directement. L’idée : on échantillonne à partir d’une distribution plus simple, puis on accepte ou rejette chaque échantillon selon un critère qui rapproche la distribution résultante de la distribution cible.

Appliqué aux LLM, le principe est le même. La distribution « simple » est le LLM courant (policy model). La distribution « cible » est le LLM idéal qui produit toujours des réponses alignées avec les préférences humaines. Le critère d’acceptation est un reward model entraîné sur des préférences humaines. On génère de nombreuses réponses, on ne garde que celles ayant un score élevé, et on réentraîne le modèle sur ces réponses filtrées pour qu’il apprenne à les produire directement.

Le pipeline en détail

Le rejection sampling pour l’alignement des LLM suit ces étapes :

Étape 1 : Modèle SFT de départ. On part d’un modèle ayant déjà subi un instruction tuning (SFT). Ce modèle sait suivre des instructions mais n’est pas encore aligné avec les préférences humaines.

Étape 2 : Génération de réponses candidates. Pour chaque prompt d’un dataset, le modèle SFT génère K réponses différentes (typiquement K = 4 à 64). La diversité est obtenue en utilisant une température supérieure à 0 lors du sampling.

Étape 3 : Scoring par le reward model. Chaque réponse candidate est évaluée par un reward model entraîné sur des données de préférences humaines. Le reward model attribue un score scalaire à chaque paire (prompt, réponse).

Étape 4 : Sélection des meilleures réponses. Pour chaque prompt, on ne conserve que la réponse (ou les réponses) ayant les scores les plus élevés. C’est l’étape de « rejet » : les réponses à faible score sont éliminées.

Étape 5 : Fine-tuning supervisé. Le modèle est réentraîné par SFT standard (même loss function que l’instruction tuning, à savoir la cross-entropy) sur les réponses filtrées. Le modèle apprend à produire directement des réponses de haute qualité.

Étape 6 (optionnel) : Itération. Le processus peut être répété de façon itérative. Le modèle fine-tuné à l’étape 5 devient le nouveau modèle de génération à l’étape 2. Chaque itération améliore la qualité des réponses candidates, car le modèle de base est meilleur.

Formulation mathématique

Si on note R(x, y) la fonction de récompense pour un prompt x et une réponse y, le rejection sampling sélectionne les réponses selon le critère :

y* = argmax_{y ∈ {y_1, ..., y_K}} R(x, y)

Dans la version la plus simple (top-1), on garde uniquement la meilleure réponse parmi les K candidates. Dans des versions plus souples, on accepte toute réponse dont le score dépasse un seuil τ :

Accepter y si R(x, y) > τ

Le seuil τ peut être fixe ou adaptatif (par exemple, le score moyen du batch). Plus τ est élevé, plus les réponses retenues sont de haute qualité, mais moins il y a de données d’entraînement (plus de rejets). C’est un compromis qualité/quantité classique.

Rejection sampling vs. Best-of-N

Les deux termes sont souvent confondus, mais ils désignent des usages différents de la même idée :

Le Best-of-N est la version « à l’inférence » du rejection sampling. Il est souvent utilisé comme baseline dans les papiers de recherche car il offre un moyen simple de mesurer l’impact d’un reward model sans la complexité d’un entraînement RL. Si le Best-of-N avec N=16 donne de meilleurs résultats que le modèle avec PPO, cela suggère que le reward model capture bien les préférences mais que le PPO n’exploite pas pleinement cette information.

Pourquoi le rejection sampling est populaire

Simplicité. Le rejection sampling ne nécessite aucun algorithme de reinforcement learning. La seule complexité technique est la génération de K réponses par prompt (inférence batch) et le scoring par le reward model. Le fine-tuning lui-même est un SFT standard. Tout praticien sachant faire du SFT peut implémenter le rejection sampling.

Stabilité. Contrairement au PPO, le rejection sampling n’a pas de problèmes d’instabilité d’entraînement, de reward hacking pendant le training, ni de sensibilité aux hyperparamètres RL. Le processus est entièrement déterministe (hors la génération des candidats).

Efficacité mémoire. Le rejection sampling ne nécessite que deux modèles : le LLM pour la génération et le reward model pour le scoring. Le PPO en nécessite quatre (LLM courant, modèle de référence, reward model, value network).

Résultats compétitifs. De nombreuses comparaisons montrent que le rejection sampling, surtout en mode itératif, atteint des performances comparables au PPO sur les benchmarks d’alignement courants. Pour Llama 2, Meta a constaté que le rejection sampling seul représentait une grande partie de l’amélioration, le PPO ajoutant un gain marginal supplémentaire.

Limites

Coût de génération. Générer K réponses par prompt multiplie le coût d’inférence par K. Pour K=64 sur un dataset de 100 000 prompts, c’est 6,4 millions d’inférences. Le coût est significatif, surtout pour les grands modèles.

Plafonné par le modèle courant. Le rejection sampling ne peut sélectionner que parmi les réponses que le modèle actuel est capable de produire. Si le modèle ne sait pas du tout résoudre un type de problème, augmenter K ne servira à rien. Le PPO, en revanche, peut pousser le modèle à explorer des zones de l’espace de réponses qu’il n’aurait pas atteintes naturellement.

Dépendance au reward model. La qualité du rejection sampling est directement plafonnée par la qualité du reward model. Si le reward model a des biais (préférence pour la longueur, le formalisme, etc.), le rejection sampling amplifiera ces biais en sélectionnant les réponses qui les exploitent.

Offline par nature. Le rejection sampling est une méthode offline : les données sont générées une fois, filtrées, puis utilisées pour l’entraînement. Il ne s’adapte pas en temps réel aux faiblesses découvertes pendant le fine-tuning. C’est l’une des raisons pour lesquelles les pipelines avancés combinent rejection sampling (pour l’alignement initial) et PPO (pour l’affinage online).

Variantes et extensions

RSO (Rejection Sampling Optimization)

Proposé à ICLR 2024, RSO utilise le rejection sampling pour générer des paires de préférences à partir de la politique optimale (pas juste de la politique SFT). L’idée : au lieu de simplement prendre la meilleure réponse pour du SFT, on crée des paires (meilleure réponse, pire réponse) et on les utilise avec DPO ou SLiC. RSO surpasse systématiquement le DPO et le SLiC standard sur de multiples tâches, car les paires de préférences sont échantillonnées d’une distribution plus proche de la politique optimale.

RS-DPO

RS-DPO combine explicitement le rejection sampling et le DPO. Le pipeline : générer K réponses par prompt, créer des paires de préférences (haute récompense vs. basse récompense) par rejection sampling, puis entraîner avec la loss DPO au lieu de la loss SFT. Les résultats montrent que RS-DPO est plus robuste aux variations de qualité du reward model que le DPO pur ou le PPO, et surpasse les deux sur MT-Bench et AlpacaEval.

STARS (Inference-Time Rejection Sampling)

STARS (2025) applique le rejection sampling au niveau des blocs de génération pendant l’inférence, pas seulement sur la réponse complète. La génération est segmentée en blocs de taille fixe, et chaque bloc est accepté ou rejeté selon un reward model. Cela permet une correction précoce des erreurs et surpasse le SFT et le DPO de 14,9 et 4,3 points de pourcentage respectivement en win-rate, tout en étant moins coûteux que le Best-of-N sur la réponse complète.

RFT (Rejection Sampling Fine-Tuning)

RFT est la version la plus directe du rejection sampling pour les tâches à réponse vérifiable (mathématiques, code). Les réponses candidates sont générées, puis filtrées par un script de vérification automatique (pas besoin de reward model appris). Seules les réponses correctes sont conservées pour le fine-tuning. C’est l’approche la plus simple et la plus robuste pour les tâches où la correction est objectivement vérifiable.

Le RFT est particulièrement pertinent pour les modèles de raisonnement. Si votre modèle ne résout que 40 % des problèmes mathématiques d’un dataset, générer 64 réponses par problème et ne garder que les correctes vous donne un dataset de haute qualité (réponses toutes correctes) qui couvre potentiellement 90 %+ des problèmes. Le modèle réentraîné sur ces données corrige ses lacunes sans avoir besoin d’un reward model coûteux. C’est aussi l’approche utilisée dans le guide DPO de HuggingFace (2025) pour créer des paires de préférences sur des tâches mathématiques : un script regex vérifie la correction de la réponse, et les paires correcte/incorrecte alimentent ensuite un entraînement DPO.

Implémentation pratique

# Pipeline de rejection sampling simplifié
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

# 1. Charger le modèle SFT et le reward model
sft_model = AutoModelForCausalLM.from_pretrained("model-sft")
reward_model = AutoModelForSequenceClassification.from_pretrained("reward-model")
tokenizer = AutoTokenizer.from_pretrained("model-sft")

K = 16  # Nombre de réponses candidates par prompt

def rejection_sample(prompts, k=K, temperature=0.8):
    """Générer K réponses, scorer, garder la meilleure."""
    selected_data = []
    for prompt in prompts:
        inputs = tokenizer(prompt, return_tensors="pt")
        # Générer K réponses avec sampling
        candidates = sft_model.generate(
            **inputs,
            num_return_sequences=k,
            temperature=temperature,
            do_sample=True,
            max_new_tokens=512,
        )
        # Scorer chaque réponse
        scores = []
        for candidate in candidates:
            text = tokenizer.decode(candidate, skip_special_tokens=True)
            reward_input = tokenizer(prompt + text, return_tensors="pt")
            score = reward_model(**reward_input).logits.item()
            scores.append((text, score))

        # Garder la meilleure réponse
        best = max(scores, key=lambda x: x[1])
        selected_data.append({"prompt": prompt, "response": best[0]})

    return selected_data

# 2. Générer le dataset filtré
filtered_dataset = rejection_sample(training_prompts)

# 3. Fine-tuner avec SFT standard sur les réponses sélectionnées
# (même pipeline que l'instruction tuning classique)

Verdict

Le rejection sampling est la technique d’alignement la plus sous-estimée dans l’écosystème LLM. Pendant que le RLHF et le DPO captent l’attention médiatique, le rejection sampling fait le gros du travail dans les pipelines de production réels. Llama 2 l’utilise sur 4 rounds avant même d’appliquer le PPO. Llama 3 continue cette approche. C’est souvent le premier levier d’alignement que les équipes implémentent car il ne nécessite aucune expertise en RL.

Pour le praticien : si vous avez un modèle SFT et un reward model, commencez par le rejection sampling avant de passer au DPO ou au PPO. Générez 8 à 16 réponses par prompt, gardez la meilleure, et fine-tunez. Répétez 2 à 4 fois. Vous obtiendrez probablement 80 % du bénéfice d’un pipeline RLHF complet avec 20 % de la complexité. Ensuite seulement, ajoutez du DPO ou du PPO pour le gain marginal restant.