Hyper-SET: Designing Transformers via Hyperspherical Energy Minimization

TL;DR

Hyper-SET introduces a hyperspherical energy minimization framework for designing Transformers, leading to a theoretically grounded, scalable, and interpretable model that performs well across various tasks.

Contribution

The paper proposes a novel energy-based, top-down approach for Transformer design, resulting in Hyper-SET, a scalable, parameter-efficient, and interpretable Transformer variant.

Findings

Hyper-SET achieves competitive performance on diverse tasks.

It scales effectively with depth using shared parameters.

The model demonstrates improved interpretability and principled design.

Abstract

Transformer-based models have achieved remarkable success, but their core components, Transformer layers, are largely heuristics-driven and engineered from the bottom up, calling for a prototypical model with high interpretability and practical competence. To this end, we conceptualize a principled, top-down approach grounded in energy-based interpretation. Specifically, we formalize token dynamics as a joint maximum likelihood estimation on the hypersphere, featuring two properties: semantic alignment in the high-dimensional space and distributional uniformity in the low-dimensional space. By quantifying them with extended Hopfield energy functions, we instantiate this idea as a constrained energy minimization problem, which enables designs of symmetric attention and feedforward modules with RMS normalization. We further present \textit{Hyper-Spherical Energy Transformer} (Hyper-SET), a recurrent-depth alternative to vanilla Transformers naturally emerging from iterative energy optimization on the hypersphere. With shared parameters across layers, Hyper-SET can scale to arbitrary depth with fewer parameters. Theoretically grounded and compact, it achieves competitive or superior performance across diverse tasks, including Sudoku solving, image classification, and masked image modeling. We also design novel variations under the proposed general principle, such as linear attention and gated feedforward layer. Moreover, we showcase its scalability with depth-wise LoRA. Our results highlight Hyper-SET as a step toward interpretable and principled Transformer design.