El muestreo top-p, también conocido como muestreo por núcleo, es una estrategia de decodificación estocástica utilizada para generar secuencias a partir de modelos probabilísticos autorregresivos, particularmente en la generación de lenguaje natural. Fue propuesto originalmente por Ari Holtzman, Yejin Choi y sus colegas en 2019 para abordar el problema del texto repetitivo y sin sentido producido por métodos de decodificación deterministas como la búsqueda de haz. Desde entonces, la técnica se ha aplicado en campos como la ingeniería de proteínas y la geofísica.
En el muestreo top-p, se establece un umbral de probabilidad p, y el siguiente elemento de una secuencia se muestrea solo del conjunto más pequeño posible de candidatos de alta probabilidad cuya probabilidad acumulada supera p. Este método adapta el tamaño del grupo de candidatos según la certeza del modelo, lo que lo hace más flexible que el muestreo top-k, que muestrea de un número fijo de candidatos. Debido a su efectividad, el muestreo top-p se usa ampliamente en muchas aplicaciones de modelos de lenguaje grandes.
Técnica
En cada paso de la generación de texto, un modelo de lenguaje calcula una distribución de probabilidad sobre todo su vocabulario para el siguiente token. Aunque es posible simplemente elegir el token con mayor probabilidad (búsqueda codiciosa) o un conjunto limitado de secuencias de alta probabilidad (búsqueda de haz), estos métodos deterministas a menudo producen texto aburrido, repetitivo o sin sentido. El muestreo top-p introduce aleatoriedad para evitar estos problemas mientras mantiene la calidad.
La idea central es muestrear de un conjunto más pequeño y más creíble de tokens en cada paso, llamado núcleo. Este núcleo contiene los tokens siguientes más probables cuya probabilidad combinada, o acumulada, apenas supera el umbral p. Al muestrear solo de este grupo de tamaño dinámico, el modelo puede adaptarse a diferentes situaciones. Cuando el modelo está seguro sobre el siguiente token (por ejemplo, un token tiene una probabilidad muy alta), el núcleo será pequeño. Cuando el modelo es incierto (las probabilidades están más distribuidas uniformemente), el núcleo será más grande, permitiendo más diversidad.
El proceso en cada paso es el siguiente:
- El modelo calcula las probabilidades para todos los posibles tokens siguientes.
- Los tokens se ordenan por su probabilidad en orden descendente.
- El núcleo se forma seleccionando tokens desde la parte superior de la lista hasta que su probabilidad acumulada supere el umbral predefinido, p.
- Las probabilidades de los tokens dentro de este núcleo se reescalan para que sumen 1. Todos los tokens fuera del núcleo se descartan (se les asigna una probabilidad de 0).
- El siguiente token final se muestrea aleatoriamente de esta nueva distribución más pequeña.
Formalmente, el núcleo, V^(p) ⊆ V, se define como el conjunto más pequeño de tokens que satisface: la suma de P(x | x_1, ..., x_{t-1}) para todos los x en V^(p) es mayor o igual a p. Aquí, P(x | x_1, ..., x_{t-1}) representa la probabilidad de un token x dado los tokens precedentes x_1, ..., x_{t-1}.
Ejemplo
Imagina que en un cierto paso, un modelo de lenguaje tiene un vocabulario de cinco palabras: [the, a, cat, dog, eats] y produce las siguientes probabilidades:
- the: 0.5
- a: 0.2
- cat: 0.1
- dog: 0.1
- eats: 0.1
Si establecemos p = 0.8:
- Los tokens se ordenan por probabilidad: [the, a, cat, dog, eats].
- La probabilidad acumulada se calcula:
- the: 0.5
- the + a: 0.5 + 0.2 = 0.7
- the + a + cat: 0.7 + 0.1 = 0.8
- El núcleo es el conjunto más pequeño con probabilidad acumulada ≥ 0.8, que es V^(0.8) = {the, a, cat}.
- Las probabilidades para este conjunto se reescalan para sumar 1:
- P(the) = 0.5 / 0.8 = 0.625
- P(a) = 0.2 / 0.8 = 0.25
- P(cat) = 0.1 / 0.8 = 0.125
- El siguiente token se muestrea entonces de esta nueva distribución, lo que significa que dog y eats tienen un 0% de probabilidad de ser elegidos.
Muestreo top-k
El muestreo top-k es una técnica similar donde el grupo de tokens candidatos se restringe a los k tokens más probables. La principal ventaja de top-p es su adaptabilidad. Cuando el modelo está muy seguro sobre el siguiente token (una distribución puntiaguda), el núcleo V^(p) puede ser muy pequeño. Cuando el modelo es incierto (una distribución plana), el núcleo puede ser mucho más grande, permitiendo más diversidad. En contraste, top-k siempre muestrea de un número fijo de tokens, que puede ser demasiado restrictivo o demasiado amplio dependiendo del contexto.
Aplicaciones
Aunque el muestreo top-p es más famoso por su uso como estrategia de decodificación para modelos de lenguaje grandes, la técnica también se ha adaptado para su uso en otros dominios científicos que implican generar o analizar datos secuenciales de modelos probabilísticos.
Generación de lenguaje natural
En su dominio original de generación de lenguaje natural, el muestreo top-p es valorado por su capacidad para producir texto más diverso y coherente en comparación con métodos deterministas. Se ha demostrado que es beneficioso en tareas como la generación automática de preguntas, donde la diversidad de muestras es importante para crear datos de entrenamiento efectivos para modelos de respuesta a preguntas.
Diseño de fármacos y proteínas
El muestreo top-p se usa en biología computacional para generar nuevas secuencias moleculares y de proteínas a partir de modelos de lenguaje especializados. En el diseño de fármacos de novo, los modelos de lenguaje químico entrenados en estructuras moleculares utilizan el muestreo por núcleo para generar bibliotecas enfocadas de nuevos candidatos a fármacos válidos. De manera similar, en ingeniería de proteínas, los modelos de lenguaje entrenados en secuencias de proteínas emplean el muestreo top-p para proponer nuevas secuencias con propiedades deseadas, expandiendo el espacio de búsqueda más allá de las variantes naturales.
Geofísica
En geofísica, el muestreo top-p se ha aplicado para generar datos sísmicos sintéticos o modelar estructuras del subsuelo. Al muestrear de modelos probabilísticos de secuencias geológicas, los investigadores pueden crear diversos escenarios plausibles que ayudan en la cuantificación de incertidumbre y la interpretación de estudios sísmicos.
Relación con otros métodos de decodificación
El muestreo top-p es una de varias estrategias de decodificación estocástica utilizadas con modelos autorregresivos. Complementa otras técnicas como el escalado de temperatura, que ajusta la nitidez de la distribución de probabilidad antes del muestreo, y la penalización por repetición, que desalienta al modelo de repetir tokens. En la práctica, top-p a menudo se combina con el escalado de temperatura para ajustar la diversidad y la calidad del texto generado. Por ejemplo, una temperatura más baja hace que la distribución sea más puntiaguda, y top-p entonces selecciona un núcleo más pequeño, resultando en salidas más conservadoras.
Consideraciones de implementación
En la práctica, el muestreo top-p requiere ordenar la distribución de probabilidad en cada paso de generación, lo que añade una sobrecarga computacional en comparación con la decodificación codiciosa. Sin embargo, para muchas aplicaciones, los beneficios en la calidad de salida superan el costo. Las implementaciones eficientes en marcos como PyTorch y TensorFlow optimizan este proceso utilizando operaciones de suma acumulada y enmascaramiento. El umbral p es un hiperparámetro que se puede ajustar; los valores comunes varían de 0.9 a 0.95 para muchas tareas de generación de lenguaje, pero el valor óptimo depende del modelo específico y la aplicación.
Impacto y adopción
Desde su introducción en 2019, el muestreo top-p se ha convertido en un componente estándar en las estrategias de decodificación de muchos modelos de lenguaje grandes, incluidos los desarrollados por OpenAI, Anthropic y Google DeepMind. A menudo es el método de muestreo predeterminado en las API de generación de texto y en las bibliotecas de código abierto. La adaptabilidad de la técnica la ha convertido en una herramienta clave para equilibrar la creatividad y la coherencia en los sistemas de IA generativa, influyendo en cómo los chatbots, los generadores de contenido y otras aplicaciones producen texto similar al humano.