OpenESM: Energy-Steered Models Are Scalable Foundation Learners

OpenESM. Energy-Steered Models Are Scalable Foundation Learners. By Yuan Pu and eight others, from the Shanghai Artificial Intelligence Laboratory. Published on the seventh of October. Its code repository is worked out: about forty source files, with a readme.

Energy-based models have shown promising properties on domain-specific tasks, but they remain hard to make work in general-purpose scenarios. The paper sets out to bridge that gap, and to bring energy-based ideas into foundational language modeling, where models are expected to handle almost anything. The question underneath is whether an energy-based approach can be trained at the scale of a foundation model, and whether it keeps improving as that scale grows.

Energy-steered models, or ESMs, integrate energy steering and token generation within a single architecture. The point of combining them is to let an energy-guided search refine what gets generated. Concretely, the model first optimizes a latent state over an energy landscape. Only then does it use that refined latent state to generate the token.

Pre-training is energy-guided as well, carried out on large-scale corpora. Through it, ESMs learn a generally effective energy landscape, one in which favorable tokens correspond to lower energy values. Generation then amounts to moving toward low energy before committing. To explore how this scales, the authors pre-train ESMs on three large-scale corpora, and evaluate them across twenty-six downstream tasks spanning a broad range of domains.

ESMs show promising scaling behavior as both model size and training data increase. They also outperform GPTs on most of the downstream tasks. The abstract gives these results without figures. The model weights and the training infrastructure are fully released, so the community can completely reproduce the results and extend them. That matters for a claim about scaling, since anyone who doubts the trend has what they need to rerun it. Both parts of the opening question get an answer here: the approach trains at scale, and it improves with scale.

The change sits at the moment of generation. Rather than producing a token straight from the latent state, an energy-steered model first searches the landscape for a better one, and only then commits. On the strength of the scaling results, the authors put ESMs forward as a promising architecture for large-scale pre-training and general-purpose problem solving.

OpenESM. Energy-Steered Models Are Scalable Foundation Learners. By Yuan Pu and eight others, from the Shanghai Artificial Intelligence Laboratory. Published on the seventh of October. This is the long version: the whole paper rather than the abstract.

Energy-based models score answers rather than write them. They learn a landscape in which desirable solutions sit low and undesirable ones sit high, and that idea has worked in pattern recognition, structured prediction, robot control and planning. But each of those models lives inside one task or one domain. Generative language models, by contrast, generalise across almost anything you put in front of them. The paper asks whether an energy-based model can be a foundation model at all. The obstacle is concrete: a model that outputs a single energy score is a critic, not a writer, and finding the lowest-energy next token by trying every token in the vocabulary is computationally infeasible. The question is how to make the critic's judgement improve the writing without enumerating anything.

The answer is a model with one Transformer backbone and two heads. The energy head is a single linear layer that outputs a scalar energy. The generative head is the familiar embedding layer and softmax that produces a distribution over the vocabulary. Both heads take the same two inputs: the context so far, and an extra latent vector the size of the model's hidden width. That latent vector is the hinge of the design. The paper describes it as a controllable representation that influences the subsequent tokens. The energy head judges how compatible that latent vector is with the context, and lower energy means a better latent for continuing this text.

Inference has two steps. First, energy steering. The latent starts as a draw from a standard Gaussian. The model computes the energy, takes the gradient of that energy with respect to the latent alone, and moves the latent a step downhill. The weights stay frozen during this. Repeated for a fixed number of steps, this is framed as Markov chain Monte Carlo sampling, but the authors drop the noise such samplers usually add, because deterministic updates proved more stable in large-scale pre-training. So in practice it is plain gradient descent on the latent. Second, token generation: the generative head reads the context and the steered latent and predicts the next token.

Training is where the engineering lives. Steering a separate latent for every position in a sequence, one at a time, would be far too slow. Instead the model appends a full row of latents after the text, one per position, and redesigns the attention mask so that each latent sees only the tokens before its position and itself. Each latent also receives the positional index of the token it predicts. One forward pass then yields every position's energy, and the steering gradient is taken on their sum, so all latents move at once while causality is preserved. The steered latents then feed the generative head, and the weights are trained with ordinary next-token cross-entropy. There is no separate energy loss: the energy head learns only through how steering changes the prediction. That means the training gradient passes through a gradient, which is second-order differentiation, and that fact returns in the limits. One detail sits only in the hyper-parameter table: the number of steering steps used in pre-training is one.

Models come at depths of twelve, twenty and twenty-six layers, with width scaling alongside depth, from two hundred and sixty-eight million to one point zero three billion parameters. Each is matched against a GPT of the same depth, trained on the same corpus with the same tokens, across Fineweb, ClimbMix and DCLM. On validation bits per byte, the energy-steered models sit below the GPTs throughout, and the gap widens as tokens and parameters grow. The authors read that widening as the central result: the advantage increases with scale rather than shrinking.

The nineteen zero-shot question-answering tasks test the same thesis downstream. At depth twenty-six, the average accuracy rises from thirty-six point zero six to forty point three two on ClimbMix, and from thirty-nine point seven nine to forty-four point seven one on DCLM. At depth twelve the margins are small, and on FineWeb the smallest energy-steered model is slightly behind its GPT. That pattern, little or no gain when small and a clear gain when larger, is the comparison that most directly supports the scaling claim. The remaining seven tasks are zero-shot generation after pre-training on OpenWebText. There the average gap to GPT is two hundredths of a bit per byte, with Lambada going the other way, and the masked diffusion baseline was trained for two epochs against one.

The paper does have a limitations section, framed as limits of the current ecosystem. The first is efficiency. Because training needs second-order gradients, existing sharding systems such as FSDP, Megatron-LM and DeepSpeed ZeRO do not readily support it. Each GPU must hold the whole model, which caps size at the memory of a single H200, and is why the largest model stops at one point zero three billion parameters. The future directions add that standard speed-ups such as Flash Attention and GPTQ are hard to apply to this architecture. The second limit is that generation is token by token, and the authors name block-wise steering over whole sentences as future work. Beyond what they say: results are single runs with no error bars, there is no wall-clock or compute-matched comparison with GPT, and the data-efficiency figures quoted for ClimbMix do not reconcile with each other in the text. The abstract counts twenty-six tasks, which is the nineteen question-answering sets plus seven generation sets.

The idea to carry away is a way to give a critic a voice. Rather than asking an energy model to search the output space, the model searches a continuous latent space it controls, by gradient, and hands the result to an ordinary generative head. Training makes the energy landscape useful only insofar as it improves prediction. At up to a billion parameters, that buys a margin over GPT that grows with scale. Whether it survives the compute accounting, and the jump past one GPU, is what the next paper has to show.