2025-11-16T11:43:12.671286

Test-Time Alignment of LLMs via Sampling-Based Optimal Control in pre-logit space

Kanai, Yoshida, Takahashi et al.

Test-time alignment of large language models (LLMs) attracts attention because fine-tuning LLMs requires high computational costs. In this paper, we propose a new test-time alignment method called adaptive importance sampling on pre-logits (AISP) on the basis of the sampling-based model predictive control with the stochastic control input. AISP applies the Gaussian perturbation into pre-logits, which are outputs of the penultimate layer, so as to maximize expected rewards with respect to the mean of the perturbation. We demonstrate that the optimal mean is obtained by importance sampling with sampled rewards. AISP outperforms best-of-n sampling in terms of rewards over the number of used samples and achieves higher rewards than other reward-based test-time alignment methods.

academic

Alignement au Moment du Test des LLM via Contrôle Optimal Basé sur l'Échantillonnage dans l'Espace Pré-Logit

Informations Fondamentales

ID de l'article: 2510.26219
Titre: Test-Time Alignment of LLMs via Sampling-Based Optimal Control in pre-logit space
Auteurs: Sekitoshi Kanai, Tsukasa Yoshida, Hiroshi Takahashi (NTT, Inc.), Haru Kuroki, Kazumune Hashimoto (The University of Osaka)
Classification: cs.LG cs.AI
Date de Publication: 30 octobre 2025 (prépublication arXiv)
Lien de l'article: https://arxiv.org/abs/2510.26219v1

Résumé

L'alignement au moment du test des grands modèles de langage (LLM) suscite l'intérêt en raison de l'évitement des coûts élevés de l'ajustement fin. Cet article propose une nouvelle méthode d'alignement au moment du test — l'échantillonnage par importance adaptative pré-logit (AISP), basée sur le contrôle prédictif de modèle avec entrées de contrôle stochastiques. AISP applique une perturbation gaussienne aux sorties de l'avant-dernière couche (pré-logits), réalisant l'alignement en maximisant la récompense attendue de la moyenne perturbée. L'article démontre que la moyenne optimale peut être obtenue par échantillonnage par importance des récompenses échantillonnées. AISP surpasse l'échantillonnage best-of-n en efficacité d'utilisation des échantillons et dépasse les autres méthodes d'alignement au moment du test basées sur les récompenses en valeur de récompense.

Contexte et Motivation de la Recherche

Problème à Résoudre

L'alignement des grands modèles de langage est une technologie clé pour assurer la sécurité des LLM et leur application généralisée. Les méthodes traditionnelles d'apprentissage par renforcement avec retours humains (RLHF) nécessitent l'ajustement fin des paramètres du LLM, entraînant des coûts de calcul énormes. L'alignement au moment du test (test-time alignment) vise à générer des réponses conformes aux préférences humaines sans mettre à jour les paramètres du modèle.

Importance du Problème

Coût de calcul: L'ajustement fin des LLM à grande échelle nécessite d'importantes ressources GPU et temps d'entraînement
Flexibilité: L'alignement au moment du test permet d'ajuster dynamiquement le comportement du modèle au moment de l'inférence
Praticité: Élimine le besoin de réentraîner le modèle pour chaque tâche spécifique

Limitations des Méthodes Existantes

Échantillonnage Best-of-N (BoN): Bien que simple et efficace, il n'explore pas activement les réponses optimales, avec une faible efficacité d'échantillonnage
RE-Control: Nécessite l'entraînement d'une fonction de valeur, demandant de grands ensembles de données (par exemple, 349 000 échantillons d'entraînement) et des coûts de stockage
Contrôle optimal traditionnel: Inadapté aux systèmes non linéaires et à grande échelle des LLM

Motivation de la Recherche

Peut-on contrôler les LLM pour explorer les réponses optimales via une méthode sans entraînement? Cet article part d'une perspective de théorie du contrôle et adopte la technique du contrôle prédictif de modèle basée sur l'échantillonnage (MPPI), proposant une méthode d'alignement au moment du test sans entraînement.

Contributions Principales

Proposition de la méthode AISP: Application pour la première fois du contrôle prédictif de modèle basé sur l'échantillonnage (MPPI) à l'alignement des LLM, réalisant l'alignement au moment du test sans entraînement en appliquant une perturbation gaussienne dans l'espace pré-logit
Contributions théoriques:
- Preuve que la distribution pré-logit optimale peut être obtenue via la limite d'énergie libre (free energy)
- Dérivation d'une solution en forme fermée basée sur l'échantillonnage par importance adaptative
- Révélation du lien théorique entre AISP et BoN (AISP dégénère en BoN sous certains paramètres)
Analyse de la justification de l'hypothèse gaussienne: Argumentation du lien entre l'hypothèse de distribution gaussienne du pré-logit et les propriétés intrinsèques de la couche softmax du réseau de neurones
Amélioration des performances:
- Surpasse significativement BoN en efficacité d'échantillonnage (récompense plus élevée avec le même nombre d'échantillons)
- Dépasse RE-Control sans entraînement
- Propose AISP par lots pour l'accélération parallèle

Détails de la Méthode

Définition de la Tâche

Étant donné une invite d'entrée $x = [x_1, ..., x_{T_x}]$ , le LLM génère une réponse $y = [y_1, ..., y_{T_y}]$ . L'objectif est de maximiser la récompense attendue tout en maintenant une contrainte de divergence KL avec le LLM de base, étant donné un modèle de récompense $r(x,y)$ :

$\min_U J(x, U) = -\mathbb{E}_{V\sim Q_{U,\sigma^2}}[r(x,y(V))] + \lambda D_{KL}(Q_{U,\sigma^2}|P)$

Architecture du Modèle

1. Conception de l'Entrée de Contrôle Stochastique

Contrairement à RE-Control qui utilise des entrées de contrôle déterministes, AISP utilise des entrées de contrôle stochastiques $v_t \sim \mathcal{N}(u_t, \sigma^2I)$ :

undefined