Explorative Modeling añade búsqueda best-of-K al preentrenamiento de modelos generativos

Why it matters

XM reframes generative scaling around training-time search. If its efficiency gains survive larger models, labs could trade extra pretraining work for cheaper inference and better data use.

Alexi Gladstone (@AlexiGlad) detalló un nuevo método de entrenamiento de modelos generativos el 31 de julio que busca varios outputs candidatos para cada ejemplo y entrena con la coincidencia más cercana, desplazando el cómputo hacia el preentrenamiento en un esfuerzo por producir mejores modelos y una inferencia más barata.

Gladstone, Becario de Investigación de Posgrado de la NSF y estudiante de doctorado en informática en la University of Illinois Urbana-Champaign, encabezó el proyecto Explorative Modeling con el profesor de UIUC Heng Ji (@hengjinlp) y la profesora asistente de Harvard Yilun Du (@du_yilun). El BLENDER Lab de UIUC lista las áreas de investigación de Gladstone como modelos del mundo, pensamiento de sistema 2, aprendizaje auto-supervisado y aprendizaje multimodal, junto con pasantías previas en Meta y Palantir.

Los investigadores sostienen que la exploración puede servir como un tercer eje de preentrenamiento junto con los parámetros del modelo y los datos. Su evidencia proviene de experimentos que abarcan generación de imágenes, generación de video, modelos de lenguaje de difusión enmascarada, políticas de robots y modelos del mundo condicionados por objetivos. Los resultados se publicaron en un preprint de arXiv del 29 de julio, con código PyTorch liberado en GitHub.

Mover la búsqueda al entrenamiento

Los modelos generativos enfrentan un problema básico: un prompt o input puede tener muchas salidas válidas. Un modelo entrenado para hacer una única predicción directa bajo una pérdida de reconstrucción puede converger hacia un promedio de esas posibilidades, produciendo un resultado que no coincide con ninguna de ellas. Los sistemas de difusión y autoregresivos evitan gran parte de ese promediado al dividir la generación en pasos más pequeños de eliminación de ruido o predicción de tokens.

Explorative Modeling, o XM, toma otra ruta. En cada paso de entrenamiento, el modelo genera K candidatos, mide cada uno contra el objetivo y retropropaga a través del candidato con la pérdida más baja. La implementación más sencilla es un bucle "best-of-K". En el repositorio publicado, los investigadores pueden habilitar el método con la bandera --xm_best_of_k K; K=1 es la línea base sin exploración.

Cada candidato extra conlleva un costo de cómputo. El artículo estima que un candidato adicional en Forward XM cuesta aproximadamente un tercio de un paso de entrenamiento estándar de un transformer porque requiere otra pasada hacia adelante sin otra pasada completa hacia atrás. Los candidatos también pueden incorporarse en la dimensión del lote y procesarse en paralelo. Por lo tanto, el método hace un intercambio directo: gastar cómputo adicional durante el entrenamiento para encontrar un mejor emparejamiento entre el ruido latente y los datos, y luego potencialmente recuperar ese gasto mediante una convergencia más rápida o menos pasos de generación.

Gladstone escribió en un ensayo del proyecto que el trabajo surgió de un intento desde primeros principios para explicar por qué el modelado generativo es difícil. "Aumentamos el tamaño de los modelos generativos y la cantidad de datos con los que los entrenamos... ¿entonces por qué no hemos escalado lo que pueden generar?" escribió.

Las ganancias reportadas

Los autores informan que agregar exploración a una receta de generación de imágenes alcanzó el mejor resultado de la línea base con 6.2 veces menos muestras de entrenamiento y 4.1 veces menos FLOPs totales. Un modelo grande explorando cinco candidatos también superó a una línea base extra-grande con 47% más parámetros, según el artículo.

Esas cifras miden los experimentos realizados por los autores, en lugar de una reducción general que se pueda asumir para otras arquitecturas. El artículo informa que las ganancias relativas de la exploración aumentaron de 7% a 36% a medida que los datos escalaron y de 13% a 23% a medida que creció el tamaño del modelo. Las métricas de imagen y video mejoraron conforme K aumentó en el rango probado, mientras que un modelo de lenguaje de difusión enmascarada explorativo produjo una mejor compensación entre perplexidad y diversidad que su línea base.

Los resultados de inferencia podrían resultar más trascendentes si se reproducen a mayor escala. Una Explorative Policy igualó o superó la línea base Diffusion Policy del artículo en cinco tareas de manipulación robótica mientras usaba una sola pasada hacia adelante de la red en lugar de 100. En tareas de planificación Maze2D, los investigadores reportaron una puntuación promedio más alta que una línea base Diffuser mientras usaban entre 16 y 256 veces menos pasos de inferencia.

La afirmación de escalado aún enfrenta una prueba mayor

La afirmación central del artículo va más allá de la escala de sus experimentos actuales. El trabajo de imágenes usa ImageNet condicional por clase a 256 por 256 píxeles, mientras que los experimentos de video modelan 10 cuadros a resolución de 128 por 128. Los investigadores dicen que los experimentos de video de mayor resolución excedieron su presupuesto de cómputo. La replicación independiente y las pruebas en modelos más grandes determinarán si la exploración sigue siendo eficiente en cómputo cuando los costos de entrenamiento alcanzan la escala de frontera.

El lenguaje es otro límite. El artículo reporta resultados más claros para modelos de lenguaje de difusión enmascarada y señala que los modelos de lenguaje autoregresivos fueron más difíciles de mejorar porque introducir una variable latente buscable es menos natural. Los autores describen sus resultados autoregresivos como modestos y dejan experimentos más completos para trabajos posteriores. Esa distinción importa porque la arquitectura dominante de grandes modelos de lenguaje sigue siendo autoregresiva.

La generación de imágenes completamente de extremo a extremo también crea un problema de costos: el número de candidatos puede necesitar crecer con el número de modos en una distribución compleja. Los investigadores proponen una versión inversa de XM que busca objetivos de datos para cada generación a un costo de cómputo mucho menor, aunque ese enfoque requiere una restricción de entropía o cobertura para evitar el colapso.

Explorative Modeling sigue siendo un proyecto académico y de código abierto afiliado a UIUC y Harvard. El Laude Institute lo incluyó en su tercera cohorte de subvenciones Slingshots en junio, y el artículo reconoce el apoyo de una beca Flapping Airplanes y del National Science Foundation Graduate Research Fellowship Program.

El valor a corto plazo es la verificabilidad. XM puede añadirse a las canalizaciones de entrenamiento de difusión y flow existentes sin reemplazar su arquitectura subyacente, dando a los equipos de investigación una manera concreta de medir si la búsqueda de candidatos adicional durante el preentrenamiento compra suficiente convergencia, calidad o ahorro de inferencia como para justificar su factura de cómputo.