💬 NLP

Regression Language Models for Code

This paper introduces a unified, lightweight Regression Language Model (RLM) that leverages a frozen LLM encoder to accurately predict diverse code execution metrics—including memory footprint, latency, and accuracy—across multiple programming languages and hardware platforms, outperforming prior feature-engineered and graph-based approaches without task-specific customization.

Original authors: Yash Akhauri, Xingyou Song, Arissa Wongpanich, Bryan Lewandowski, Mohamed S. Abdelfattah

Published 2026-05-28

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Yash Akhauri, Xingyou Song, Arissa Wongpanich, Bryan Lewandowski, Mohamed S. Abdelfattah

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you have a massive library of computer programs written in dozens of different languages (like Python, C++, or even specialized languages for graphics cards). Usually, if you want to know how much memory a program will eat up, how long it will take to run, or how accurate a neural network will be, you have to actually run the program on a computer to find out. This is slow, expensive, and sometimes impossible if you haven't built the hardware yet.

For a long time, experts tried to build special "crystal balls" to predict these numbers without running the code. But these crystal balls were like custom-made tools: to predict memory for Python, you needed one tool; to predict speed for a graphics card, you needed a completely different tool. They required humans to manually dig through the code, count specific commands, and turn the code into complex charts and graphs just to feed them into a calculator. If the code changed slightly, the tool often broke.

The New "Universal Translator" for Code

This paper introduces a new kind of AI called a Regression Language Model (RLM). Think of this not as a calculator, but as a super-translator that speaks both "Computer Code" and "Numbers."

Here is how it works, using simple analogies:

1. The "Frozen Brain" and the "Creative Writer"

The model is built like a two-part team:

The Frozen Brain (Encoder): This part is a pre-trained AI (based on a model called T5Gemma) that has already read millions of books and lines of code. It understands the structure of sentences and code perfectly. The researchers "froze" this part, meaning they didn't teach it anything new. It just acts as a smart reader that understands what the code is saying.
The Creative Writer (Decoder): This is the part they actually trained. Instead of writing a story, this writer's job is to look at the code the "Brain" read and write a number as the next word in a sentence.

The Analogy: Imagine you are reading a recipe (the code). A normal AI might try to guess the ingredients list. This new AI looks at the recipe and says, "Based on this text, the next word is '500' (calories)." It treats the prediction of a number (like memory usage or speed) exactly like finishing a sentence.

2. One Tool for Everything

The biggest breakthrough is that this single "Universal Translator" can handle everything at once.

It can read Python code and guess how much memory it uses.
It can read C++ code and guess how long it takes to run.
It can read a Neural Network design (written in a format called ONNX) and guess how accurate it will be or how fast it runs on a specific chip.

Previously, you needed a different "expert" for each of these tasks. This model is like a Swiss Army Knife that can switch between being a memory expert, a speed expert, and an accuracy expert just by looking at the text.

3. How It "Counts" Without Math

One of the clever tricks is how the model outputs numbers. Instead of trying to do complex math to output "72.5," it breaks the number down into tiny pieces, like a digital lock.

It doesn't say "72.5".
It says: "Plus sign... 7... 2... decimal point... 5."
The Analogy: Imagine you are teaching a child to count. Instead of asking them to solve "50 + 25," you ask them to write the answer digit by digit: "5... 0... plus... 2... 5... equals... 7... 2... 5." This makes it much easier for the AI to handle huge numbers (like millions of bytes) or tiny numbers without getting confused.

4. Why It's Better Than the Old Way

No Manual Labor: Old methods required humans to manually design features (like "count the number of loops"). This model reads the raw text directly, just like a human reads a book.
Better at Ranking: The paper shows that if you give this model 10 different versions of the same code, it is very good at sorting them from "fastest" to "slowest" or "smallest memory" to "largest memory." It's like a judge who can look at 10 runners and instantly know who will win, even without a stopwatch.
Cost: The researchers found that using this model is much cheaper than using giant, general-purpose AI chatbots (like GPT-5) to guess the numbers. The chatbots are like using a sledgehammer to crack a nut; they are expensive and often wrong. This model is a specialized, lightweight tool that gets the job done for pennies.

What They Actually Proved

The paper tested this on:

CodeNet: A huge dataset with millions of code snippets in 37 different languages. The model successfully predicted memory usage across all of them.
APPS: A set of coding challenges. The model could look at a solution and predict its memory usage with very high accuracy (better than 90% correlation with reality).
Neural Networks: It predicted the speed and accuracy of AI designs better than the previous state-of-the-art methods that used complex graph diagrams.

In Summary:
The paper claims that we don't need to build complex, custom tools to predict how code performs. Instead, we can treat code like a story and numbers like the next word in that story. By training a single, unified AI to do this, we can predict memory, speed, and accuracy for almost any programming language or AI design, simply by reading the text. It turns a difficult math problem into a simple "what comes next?" game.

Technical Summary: Regression Language Models for Code

Problem Statement

The paper addresses the challenge of code-to-metric regression: predicting numeric outcomes of code executions (such as memory footprint, latency, or accuracy) directly from source code or computational graph representations. Historically, this task has been hindered by the open-ended nature of programming languages and the reliance on heavy, domain-specific feature engineering. Prior methods often required converting code into rigid graph representations (e.g., syntax trees, adjacency matrices) or tabular vectors via manual counting of commands, which limited transferability to new tasks, languages, or kernel types. The core difficulty lies in learning accurate regression models over highly non-tabular, graph-based data without losing the flexibility to generalize across diverse domains.

Methodology

The authors propose a Regression Language Model (RLM), a unified framework that treats code-to-metric regression as a standard next-token prediction problem, aligning with the modern Large Language Model (LLM) paradigm.

Architecture and Training

Unified Encoder-Decoder: The RLM utilizes a frozen encoder (specifically T5Gemma-S) to process input text and a trainable decoder to generate numeric outputs. The encoder is pretrained on standard language data, while only the decoder undergoes backpropagation, significantly reducing training costs.
Text-to-Text Regression: Unlike traditional approaches that map inputs to fixed-length vectors, the RLM accepts raw text representations of code (e.g., Python, C++, ONNX graphs) and outputs numeric values as sequences of tokens.
Specialized Numeric Tokenization: To handle the wide dynamic range of metrics (from $10^{-2}$ to $10^6$ ) without normalization, the model uses an explicit digit-by-digit tokenization. Numbers are represented using special tokens for sign, exponent, and mantissa (e.g., <+><-><1><7><2><5> for $72.5$). This avoids numeric instability and the need for dataset-specific scaling.
Multi-Objective Modeling: Leveraging the autoregressive nature of the decoder, the model can conditionally predict multiple metrics sequentially (e.g., $p(y_{latency} | y_{accuracy}, x)$ ), capturing inherent correlations between objectives like accuracy and hardware latency.

Data Sources

The model is trained on a diverse, unified dataset comprising:

High-Level Programming: Code from APPS (LeetCode problems), CodeNet (14M samples across 37 languages), and KernelBook (Triton GPU kernels).
Neural Architecture Search (NAS): 520K unique architectures from benchmarks like NASBench-101/201, FBNet, and Once-for-All, represented as ONNX intermediate representations.

Key Contributions

Unified General-Purpose Model: The paper demonstrates that a single RLM can simultaneously predict memory usage across multiple languages (Python, C++, etc.), GPU kernel latency, and neural network accuracy/latency, eliminating the need for task-specific feature engineering.
Superiority over Graph-Based Methods: Despite reading dense ONNX text representations, the RLM outperforms state-of-the-art Graph Neural Network (GNN) baselines on NAS benchmarks (e.g., achieving a Kendall-Tau of 0.46 on design spaces where GNNs previously dominated).
Efficiency and Scalability: The use of a frozen encoder allows for cost-effective training. The authors show that a relatively small 300M parameter model achieves high performance, and scaling to 600M parameters yields further gains.
Normalization-Free Prediction: The digit-by-digit tokenization strategy enables the model to predict values across orders of magnitude without manual normalization, a significant advantage over MSE-based regression heads.

Experimental Results

Code Performance: On the APPS dataset, the RLM achieves a Spearman rank correlation ( $\rho$ ) of 0.93 for memory prediction. On CodeNet, it achieves an average $\rho > 0.5$ across 17 languages, with C++ reaching 0.748.
NAS Performance: The RLM achieves an average Kendall-Tau of 0.46 across five classic NAS design spaces, outperforming previous SOTA methods like FLANT and GNN-based encoders. It successfully predicts both accuracy and hardware-specific latencies (e.g., on Pixel3, Eyeriss, Intel CPU, Nvidia GPU) simultaneously.
Ablation Studies:
- Frozen Encoders: Pretrained encoders significantly improve convergence and final performance compared to random initialization.
- Decoder vs. Regression Heads: The decoder-based output (0.800 Spearman $\rho$ ) substantially outperforms explicit regression heads (0.478 without normalization, 0.717 with normalization).
- Tokenization: Custom, compact tokenizers for ONNX graphs improve performance (0.723 vs. 0.533 with default tokenizers).
- In-Context Learning (ICL): While ICL with large models (e.g., GPT-5) improves with more shots, it remains consistently below the in-weight trained RLM and incurs significantly higher inference costs due to context window limitations.

Significance and Claims

The paper claims to pave the way for simplifying computational graph regression into a generic next-token prediction problem. By demonstrating that language models can effectively regress over pure text representations of code and graphs, the work suggests that the field can move away from rigid, hand-crafted feature engineering toward a more flexible, unified approach.

The authors position this as a valuable tool for:

Program Search: Accelerating the selection of efficient code solutions.
Hardware-Software Co-design: Predicting latency across diverse hardware platforms.
Compiler Optimization: Estimating performance metrics for compiler passes.

The paper remains modest regarding broader implications, noting that while the method shows promise for predicting numeric outcomes from raw code, extending this to predict outcomes of entire scientific experiments remains an open question for future work. The primary contribution is establishing the general capability of language models to serve as universal regressors for code and computational graphs without post-processing or specialized feature extraction.