Regression Language Models for Code
Authors: name: Yash Akhauri, name: Xingyou Song, name: Arissa Wongpanich, name: Bryan Lewandowski, name: Mohamed S. Abdelfattah
Organization: Google

We study 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) 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 initialized from T5Gemma, obtains > 0.9 Spearman-rank on
competitive programming submissions from APPS, and a single unified model
achieves > 0.5 average Spearman-rank across 17 separate 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.