Los Transformadores Basados en Energía son Aprendices y Pensadores Escalables

Resumen

Las técnicas de computación en tiempo de inferencia, análogas al Pensamiento de Sistema 2 humano, han ganado recientemente popularidad para mejorar el rendimiento de los modelos. Sin embargo, la mayoría de los enfoques existentes presentan varias limitaciones: son específicos de una modalidad (por ejemplo, funcionan solo en texto), específicos de un problema (por ejemplo, dominios verificables como matemáticas y programación), o requieren supervisión/entrenamiento adicional además del preentrenamiento no supervisado (por ejemplo, verificadores o recompensas verificables). En este artículo, nos planteamos la pregunta: "¿Es posible generalizar estos enfoques de Pensamiento de Sistema 2 y desarrollar modelos que aprendan a pensar únicamente a partir del aprendizaje no supervisado?" Curiosamente, encontramos que la respuesta es sí, al aprender a verificar explícitamente la compatibilidad entre las entradas y las predicciones candidatas, y luego reformular los problemas de predicción como una optimización con respecto a este verificador. Específicamente, entrenamos Transformadores Basados en Energía (EBTs, por sus siglas en inglés) —una nueva clase de Modelos Basados en Energía (EBMs)— para asignar un valor de energía a cada par de entrada y predicción candidata, permitiendo predicciones mediante la minimización de energía basada en descenso de gradiente hasta la convergencia. En ambas modalidades, discreta (texto) y continua (visual), encontramos que los EBTs escalan más rápido que el enfoque dominante Transformer++ durante el entrenamiento, logrando una tasa de escalado hasta un 35% mayor con respecto a datos, tamaño de lote, parámetros, FLOPs y profundidad. Durante la inferencia, los EBTs mejoran el rendimiento con Pensamiento de Sistema 2 en un 29% más que el Transformer++ en tareas de lenguaje, y los EBTs superan a los Transformadores de Difusión en la eliminación de ruido de imágenes mientras utilizan menos pasos hacia adelante. Además, encontramos que los EBTs obtienen mejores resultados que los modelos existentes en la mayoría de las tareas posteriores, dado el mismo o peor rendimiento de preentrenamiento, lo que sugiere que los EBTs generalizan mejor que los enfoques existentes. En consecuencia, los EBTs representan un nuevo paradigma prometedor para escalar tanto las capacidades de aprendizaje como de pensamiento de los modelos.

English

Inference-time computation techniques, analogous to human System 2 Thinking, have recently become popular for improving model performances. However, most existing approaches suffer from several limitations: they are modality-specific (e.g., working only in text), problem-specific (e.g., verifiable domains like math and coding), or require additional supervision/training on top of unsupervised pretraining (e.g., verifiers or verifiable rewards). In this paper, we ask the question "Is it possible to generalize these System 2 Thinking approaches, and develop models that learn to think solely from unsupervised learning?" Interestingly, we find the answer is yes, by learning to explicitly verify the compatibility between inputs and candidate-predictions, and then re-framing prediction problems as optimization with respect to this verifier. Specifically, we train Energy-Based Transformers (EBTs) -- a new class of Energy-Based Models (EBMs) -- to assign an energy value to every input and candidate-prediction pair, enabling predictions through gradient descent-based energy minimization until convergence. Across both discrete (text) and continuous (visual) modalities, we find EBTs scale faster than the dominant Transformer++ approach during training, achieving an up to 35% higher scaling rate with respect to data, batch size, parameters, FLOPs, and depth. During inference, EBTs improve performance with System 2 Thinking by 29% more than the Transformer++ on language tasks, and EBTs outperform Diffusion Transformers on image denoising while using fewer forward passes. Further, we find that EBTs achieve better results than existing models on most downstream tasks given the same or worse pretraining performance, suggesting that EBTs generalize better than existing approaches. Consequently, EBTs are a promising new paradigm for scaling both the learning and thinking capabilities of models.

Los Transformadores Basados en Energía son Aprendices y Pensadores Escalables

Energy-Based Transformers are Scalable Learners and Thinkers

Resumen

Support