Transformadores Baseados em Energia são Aprendizes e Pensadores Escaláveis

Resumo

Técnicas de computação em tempo de inferência, análogas ao Pensamento do Sistema 2 humano, têm se tornado recentemente populares para melhorar o desempenho dos modelos. No entanto, a maioria das abordagens existentes sofre de várias limitações: são específicas de modalidade (por exemplo, funcionando apenas com texto), específicas de problema (por exemplo, domínios verificáveis como matemática e programação) ou exigem supervisão/treinamento adicional além do pré-treinamento não supervisionado (por exemplo, verificadores ou recompensas verificáveis). Neste artigo, fazemos a pergunta: "É possível generalizar essas abordagens de Pensamento do Sistema 2 e desenvolver modelos que aprendam a pensar apenas a partir de aprendizado não supervisionado?" Curiosamente, descobrimos que a resposta é sim, aprendendo a verificar explicitamente a compatibilidade entre entradas e previsões candidatas e, em seguida, reformulando problemas de previsão como otimização em relação a esse verificador. Especificamente, treinamos Transformers Baseados em Energia (EBTs) — uma nova classe de Modelos Baseados em Energia (EBMs) — para atribuir um valor de energia a cada par de entrada e previsão candidata, permitindo previsões por meio de minimização de energia baseada em gradiente até a convergência. Em ambas as modalidades discretas (texto) e contínuas (visuais), descobrimos que os EBTs escalam mais rapidamente do que a abordagem dominante Transformer++ durante o treinamento, alcançando uma taxa de escalonamento até 35% maior em relação a dados, tamanho do lote, parâmetros, FLOPs e profundidade. Durante a inferência, os EBTs melhoram o desempenho com Pensamento do Sistema 2 em 29% a mais do que o Transformer++ em tarefas de linguagem, e os EBTs superam os Transformers de Difusão na remoção de ruído de imagens enquanto usam menos passos diretos. Além disso, descobrimos que os EBTs alcançam melhores resultados do que os modelos existentes na maioria das tarefas subsequentes, dado o mesmo ou pior desempenho de pré-treinamento, sugerindo que os EBTs generalizam melhor do que as abordagens existentes. Consequentemente, os EBTs são um novo paradigma promissor para escalar tanto as capacidades de aprendizado quanto de pensamento dos 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.

Transformadores Baseados em Energia são Aprendizes e Pensadores Escaláveis

Energy-Based Transformers are Scalable Learners and Thinkers

Resumo

Support