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

Alineación en Tiempo de Prueba de LLMs mediante Control Óptimo Basado en Muestreo en Espacio Pre-logit

Información Básica

ID del Artículo: 2510.26219
Título: Test-Time Alignment of LLMs via Sampling-Based Optimal Control in pre-logit space
Autores: Sekitoshi Kanai, Tsukasa Yoshida, Hiroshi Takahashi (NTT, Inc.), Haru Kuroki, Kazumune Hashimoto (The University of Osaka)
Clasificación: cs.LG cs.AI
Fecha de Publicación: 30 de octubre de 2025 (preimpresión arXiv)
Enlace del Artículo: https://arxiv.org/abs/2510.26219v1

Resumen

La alineación en tiempo de prueba de modelos de lenguaje grande (LLM) ha recibido atención por evitar los costosos gastos de ajuste fino. Este artículo propone un nuevo método de alineación en tiempo de prueba denominado Muestreo Adaptativo de Importancia Pre-logit (AISP), basado en control predictivo de modelo con entrada de control estocástico. AISP aplica perturbación gaussiana a la salida de la penúltima capa (pre-logits), logrando alineación mediante la maximización de la recompensa esperada de la media perturbada. El artículo demuestra que la media óptima puede obtenerse mediante muestreo de importancia adaptativo de recompensas muestreadas. AISP supera al muestreo best-of-n en eficiencia de uso de muestras y supera otros métodos de alineación en tiempo de prueba basados en recompensas en valores de recompensa.

Antecedentes y Motivación de la Investigación

Problema a Resolver

La alineación de modelos de lenguaje grande es una tecnología clave para garantizar la seguridad de los LLM y su aplicación generalizada. Los métodos tradicionales de aprendizaje por refuerzo con retroalimentación humana (RLHF) requieren ajuste fino de parámetros del LLM, generando enormes costos computacionales. La alineación en tiempo de prueba (test-time alignment) tiene como objetivo hacer que los LLM generen respuestas alineadas con las preferencias humanas sin actualizar los parámetros del modelo.

Importancia del Problema

Costo Computacional: El ajuste fino de LLM a gran escala requiere recursos GPU significativos y tiempo de entrenamiento
Flexibilidad: La alineación en tiempo de prueba permite ajustar dinámicamente el comportamiento del modelo durante la fase de inferencia
Practicidad: Elimina la necesidad de reentrenar el modelo para cada tarea específica

Limitaciones de Métodos Existentes

Muestreo Best-of-N (BoN): Aunque es simple y efectivo, no explora activamente respuestas óptimas, con baja eficiencia de muestras
RE-Control: Requiere entrenar una función de valor, necesitando grandes conjuntos de datos (como 349,000 muestras de entrenamiento) y costos de almacenamiento
Control Óptimo Tradicional: No es aplicable a sistemas no lineales y a gran escala como los LLM

Motivación de la Investigación

¿Puede controlarse un LLM para explorar respuestas óptimas mediante un método que no requiera entrenamiento? Este artículo parte de la teoría de control, adoptando la técnica de control predictivo de modelo basado en muestreo (MPPI), proponiendo un método de alineación en tiempo de prueba sin necesidad de entrenamiento.

Contribuciones Principales

Propuesta del Método AISP: Primera aplicación del control predictivo de modelo basado en muestreo (MPPI) a la alineación de LLM, logrando alineación en tiempo de prueba sin entrenamiento mediante perturbación gaussiana en el espacio pre-logit
Contribuciones Teóricas:
- Demostración de que la distribución pre-logit óptima puede obtenerse mediante el límite de energía libre
- Derivación de solución de forma cerrada basada en muestreo de importancia adaptativo
- Revelación de la conexión teórica entre AISP y BoN (AISP se degrada a BoN bajo parámetros específicos)
Análisis de Razonabilidad de la Suposición Gaussiana: Argumentación de la conexión entre la suposición de distribución gaussiana pre-logit y la conexión inherente de la capa softmax de redes neuronales
Mejora de Rendimiento:
- Superioridad significativa sobre BoN en eficiencia de muestras (mayor recompensa con el mismo número de muestras)
- Superación de RE-Control sin necesidad de entrenamiento
- Propuesta de AISP por Lotes para aceleración paralela

Explicación Detallada del Método

Definición de Tarea

Dado un prompt de entrada $x = [x_1, ..., x_{T_x}]$ , el LLM genera una respuesta $y = [y_1, ..., y_{T_y}]$ . El objetivo es maximizar la recompensa esperada mientras se mantiene una restricción de divergencia KL con respecto al LLM base, dado un modelo de recompensa $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)$

Arquitectura del Modelo

1. Diseño de Entrada de Control Estocástico

A diferencia de RE-Control que utiliza entrada de control determinista, AISP utiliza entrada de control estocástico $v_t \sim \mathcal{N}(u_t, \sigma^2I)$ :

undefined