Regression Language Models for Code

We study \textbf{code-to-metric regression}: predicting numeric outcomes of code executions, a challenging task due to the open-ended nature of programming languages. While prior methods have resorted to heavy and domain-specific feature engineering, we show that a single unified Regression Language Model (RLM) using a frozen LLM encoder can simultaneously predict directly from text, (i) the memory footprint of code across multiple high-level languages such as Python and C++, (ii) the latency of Triton GPU kernels, and (iii) the accuracy and speed of trained neural networks represented in ONNX. In particular, a relatively small 300M parameter RLM based on T5Gemma, obtains $>$0.9 Spearman-rank on competitive programming submissions from APPS, and a single unified model achieves $>$0.5 average Spearman-rank across 24 different programming languages from CodeNet. Furthermore, the RLM can obtain the highest average Kendall-Tau of 0.46 on five classic NAS design spaces previously dominated by graph neural networks, and simultaneously predict architecture latencies on numerous hardware platforms.

Predicting how a program or computation graph will run -- such as its memory footprint, execution speed, or accuracy, is incredibly valuable but notoriously difficult. Traditionally, these predictions required researchers to perform feature engineering, by representing the code as abstract syntax trees or complicated graph features. To simplify this, we developed a unified "Regression Language Model" (RLM) using a frozen T5Gemma encoder that reads raw code as text and directly predicts numerical performance metrics. Using a specialized numeric decoder, it simply performs supervised next-token prediction on numeric outputs represented as tokens, and can make precise numeric predictions across vastly different scales. Our open-source package can be found at https://github.com/google-deepmind/regress-lm, and paves the way for research in LLM-based regression.