Energy-Entropy Regularization: The True Power of Minimal Looped Transformers

TL;DR

This paper introduces a novel training framework using Tsallis entropy and Hamiltonian dynamics to improve the training of looped Transformers, enabling them to solve complex reasoning tasks by transforming the loss landscape.

Contribution

It presents a new method that leverages physical flow concepts to train single-head looped Transformers effectively, revealing their internal reasoning mechanisms.

Findings

Successfully trained a single-head looped Transformer on a 1000-token induction head task.

Demonstrated the effectiveness of the proposed framework in transforming the loss landscape.

Revealed the internal mechanisms behind the superior reasoning capabilities of looped Transformers.

Abstract

Recent research suggests that looped Transformers have superior reasoning capabilities compared to standard deep architectures. Current approaches to training single-head looped architectures on benchmark tasks frequently fail or yield suboptimal performance due to a highly non-convex and irregular loss landscape. In these settings, optimization often stagnates in poor local minima and saddle points of the loss landscape, preventing the model from discovering the global minimum point. The internal mechanisms of these single-head looped transformer models remain poorly understood, and training them from scratch remains a significant challenge. In this paper, we propose a novel training framework that leverages Tsallis entropy and Hamiltonian dynamics to transform the geometry of the loss landscape. By treating the parameter updates as a physical flow, we successfully trained a single-head looped Transformer with model dimension $d = 8$ to solve induction head task with input sequence length of 1000 tokens. This success reveals the internal mechanism behind the superior reasoning capability.