Regression Language Models for Code (2509.26476v1)

Abstract: 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 LLM (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.

• The paper presents a novel unified model that predicts numeric code metrics directly from raw code using an autoregressive numeric output, bypassing manual feature engineering.
• It employs a pretrained T5Gemma encoder combined with a decoder that outputs digitized numeric tokens, achieving high rank correlations on diverse datasets across multiple languages.
• The approach supports multi-task predictions including memory, latency, and accuracy, offering robust and transferable insights for code optimization and hardware-software co-design.

Regression LLMs for Code: Unified Code-to-Metric Prediction

Motivation and Problem Statement

The paper addresses the challenge of code-to-metric regression, i.e., predicting numeric outcomes (such as memory usage, latency, or accuracy) directly from code or computational graph representations. Traditional approaches rely heavily on domain-specific feature engineering, which is brittle and non-transferable across languages, compilers, or hardware platforms. The authors propose a unified Regression LLM (RLM) that leverages pretrained LLM encoders (specifically T5Gemma) and a decoder-based numeric output scheme to predict metrics from raw code or graph text, spanning multiple programming languages and representation levels.

Methodology

Model Architecture

The RLM is structured as an encoder-decoder transformer, initialized from T5Gemma. The encoder ingests code or graph representations as plain text, while the decoder autoregressively generates numeric tokens representing the target metric(s). Numeric outputs are tokenized digit-by-digit, including sign and exponent tokens, enabling normalization-free regression and robust handling of wide-ranging target values. Constrained decoding ensures valid numeric outputs and supports both pointwise prediction and density estimation for uncertainty quantification.

Multi-Task and Multi-Objective Regression

The model is trained on a mixture of regression tasks, including:

• Memory footprint prediction for code in Python, C++, and other languages
• Latency estimation for Triton GPU kernels
• Accuracy and speed prediction for neural networks represented in ONNX IR

The autoregressive decoder enables conditional modeling of multiple objectives (e.g., predicting latency and then accuracy), capturing inter-metric dependencies that parallel regression heads cannot.

Data Sources

The RLM is trained and evaluated on diverse datasets:

• CodeNet: 7.3M code samples across 17 languages, with memory usage labels
• APPS Leetcode: 98.9K Python problems, with peak memory usage
• KernelBook: 12.6K Triton GPU kernels, with latency measurements
• NASBench, FBNet, OFA, SNAS, etc.: 520K neural architectures in ONNX IR, with accuracy, FLOPs, parameter count, and hardware-specific latencies

Experimental Results

High-Level Code Regression

• On APPS Leetcode, the RLM achieves Spearman rank correlation $p > 0.9$ for memory prediction.
• On CodeNet, a single model achieves $p > 0.5$ average across 17 languages, with best results on C++ ($p = 0.748$) and competitive performance on niche languages (Lua, Haskell, etc.).
• The RLM can rank candidate solutions within a problem, outperforming random selection in identifying the most memory-efficient code.

Kernel and Graph Regression

• For Triton kernel latency, the RLM achieves $p = 0.516$.
• On NAS benchmarks, the RLM matches or exceeds state-of-the-art GNN-based methods (e.g., FLAN, Arch2Vec, CATE) in Kendall-Tau rank correlation, with average $\tau = 0.46$ across five NAS design spaces.
• The RLM supports multi-metric prediction, accurately modeling Pareto frontiers of accuracy vs. latency across multiple hardware platforms.

Ablation Studies

• Pretraining: Initializing from a language-pretrained encoder (T5Gemma) accelerates convergence and improves final regression metrics. Regression pretraining on synthetic tasks (e.g., FLOPs) further boosts performance.
• Decoder vs. Regression Head: Decoder-based numeric output (cross-entropy loss) outperforms explicit regression heads (MSE loss), especially in normalization-free settings and with wide-ranging target values.
• Scaling: Larger encoder models (600M vs. 300M parameters) yield higher rank correlations, but require careful hyperparameter tuning and more compute.
• Tokenization and Sequence Length: Custom tokenizers (e.g., ONNX-aware) and longer input sequences improve performance, especially for large graphs.
• Fine-Tuning: Fine-tuning on specific languages does not consistently improve performance if the pretraining corpus is sufficiently rich, but is beneficial for low-resource NAS tasks.

Implementation Considerations

• Training: The RLM is trained with Adafactor, using a linear warmup and cosine decay schedule. Input sequences are cropped to a maximum length (2048 tokens), with custom tokenization for code and ONNX graphs.
• Inference: Median aggregation over multiple decoder samples (default: 64) is used for pointwise prediction. Increasing sample size can improve accuracy.
• Resource Requirements: The default model (300M parameters) is tractable for modern hardware, but scaling to larger models or longer sequences may require distributed training and memory optimization.
• Deployment: The normalization-free decoder and text-based input format enable straightforward deployment across diverse codebases and hardware platforms, without the need for feature engineering or dataset-specific normalization.

Implications and Future Directions

The RLM paradigm demonstrates that unified, text-based regression models can subsume traditional feature engineering and specialized graph models for code-to-metric prediction. This approach aligns with the LLM paradigm, treating regression as next-token prediction and leveraging transfer learning from language data. The ability to predict multiple metrics from raw code or graph text has direct applications in program search, hardware-software co-design, compiler optimization, and automated machine learning.

Open questions remain regarding the generalization of RLMs to entirely new code domains, the limits of scaling, and the potential for predicting experimental outcomes from raw code. Further research may explore larger models, more diverse pretraining corpora, and integration with program synthesis and optimization pipelines.

Conclusion

The paper establishes Regression LLMs as effective, normalization-free, and transferable predictors for code-to-metric regression across languages, compilers, and hardware. By leveraging pretrained LLM encoders and autoregressive numeric decoders, RLMs outperform or match specialized GNNs and regression heads, while simplifying the modeling pipeline and enabling multi-task, multi-objective prediction. This work provides a foundation for future research in unified code analysis, performance prediction, and automated optimization.