Regression Language Models for Code

 

This AI Doesn't Write Your Code—It Predicts Its Performance

Beyond Writing Code

We've become accustomed to AI models that can generate code, but a new frontier is emerging: AI that can predict code's performance simply by reading it. This is a fundamentally challenging task, known as code-to-metric regression, that has traditionally required intensive, domain-specific feature engineering. A new approach using Regression Language Models, however, promises to unify this fragmented landscape.

One Model to Predict Diverse Performance Metrics

One Model to Predict It All

The core innovation is the "Regression Language Model" (RLM), a single, unified model designed for "code-to-metric regression." Its function is to predict a wide range of numeric outcomes directly from the source code text. This is a significant departure from previous methods, which required separate, specialized tools for different languages and performance metrics.

...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.

This capability is important because it dramatically simplifies performance prediction. This unification works not just across performance domains (memory, latency, network accuracy) but also across programming languages, allowing one versatile model to replace a complex ecosystem of specialized tools.

It's Surprisingly Effective on Competitive Code

Excels on Competitive Code: A Stress Test for Performance Prediction

Despite being a relatively small 300M parameter model initialized from T5Gemma, the RLM demonstrates remarkable accuracy on highly complex code. The model achieves a Spearman-rank correlation greater than 0.9 on competitive programming submissions from APPS.

This result is particularly impactful because competitive programming code is intentionally complex and engineered to push the limits of performance. The model’s high score is a powerful validation of its deep comprehension of intricate logic, and its initialization from a foundational model highlights a powerful trend in AI: fine-tuning general-purpose models to achieve state-of-the-art performance on highly specialized tasks.

It Beats Specialized AI in Its Own Arena

Outperforming Specialized GNNs in Neural Architecture Search

In a counter-intuitive finding, the RLM excels in an area that has been the traditional domain of a different type of AI. The model achieved the highest average Kendall-Tau of 0.46 on five classic Neural Architecture Search (NAS) design spaces previously dominated by graph neural networks.

In simple terms, this means a language model proved more effective at predicting the performance of different neural network architectures than the specialized Graph Neural Networks (GNNs) built for this task. Furthermore, the RLM can simultaneously predict architecture latencies on numerous hardware platforms, adding another layer of versatility. This is a significant development, suggesting that text-based language models can outperform graph-based models even on tasks that seem inherently structural.

A Single Model Understands 17 Programming Languages

Unified Prediction Across 17 Languages

The model's multilingual capabilities are one of its most impressive features. Research shows that a single unified RLM achieves an average Spearman-rank greater than 0.5 across 17 separate languages from CodeNet.

This multi-language proficiency represents a crucial step toward universal code analysis tools, removing the typical silos imposed by language-specific profilers. Creating a single performance prediction tool that works reliably across such a vast and diverse set of programming languages is a major breakthrough.

The Future of Code Optimization

The emergence of Regression Language Models signals a pivotal shift—from generating code to a deep, multi-faceted comprehension that unifies performance prediction across languages, hardware, and even model architectures. This moves us closer to a future where performance bottlenecks can be identified and optimized before a single line of code is ever executed.

If AI can accurately predict the performance of any code before it runs, what will that mean for the future of software development and optimization?