Regression Language Models for Code

1. Introduction

• Predicting useful metrics from code could inform decisions and optimization, but accurate regression over graph-based representations requires transferable feature engineering.Earlier text-to-text regression had not established whether it could predict metrics for compiler-produced programs and machine-learning architectures.Source: §1 · p. 1
• A unified RLM with a frozen T5Gemma-S encoder predicts metrics across GPU kernels, neural-network architectures, and numerous programming languages.Its decoder is trained on a diverse combination of regression data.Source: §1 · p. 1
• The RLM remains competitive with and can outperform state-of-the-art GNN-based methods on standard NAS benchmarks, despite reading dense ONNX representations.It also supports latency prediction across different hardware.Source: §1 · p. 1
• Comprehensive ablations report faster convergence with a frozen pretrained encoder, better performance from decoder-based numeric outputs than MSE heads, gains with larger encoders, and sensitivity to tokenization and sequence length.Source: §1 · p. 1

2. Related Work and Motivation

• Many computational-graph methods require substantial feature engineering, and graph-specific features may not transfer even to similar tasks.Source: §2 · p. 2
• Code regression has used command counts and program metrics or syntax-tree and graph features, but adding a new command or kernel may force the feature pipeline to be rebuilt.Source: §2 · p. 2
• NAS regression spans path encodings, graph statistics, zero-cost proxies, activation information, graph kernels, GNN embeddings, and hardware embeddings for latency or cost.Source: §2 · p. 2
• The paper establishes the ability to regress over many code variants from pure text, a capability it says had not been investigated and is valuable across computing fields.Earlier LLM regression work included regression heads for NAS, tabular data, and recommender systems.Source: §2 · pp. 2–3

3. Method

• The encoder-decoder RLM predicts y-values as next tokens from text inputs, using a frozen T5Gemma encoder to reduce training costs while retaining pretrained code knowledge.Source: §3 · p. 3
• Digit-by-digit sign, exponent, and mantissa tokens encode numeric outputs without normalization, while constrained decoding ensures valid numbers for point predictions or density estimates.Source: §3 · p. 3

3.1. Multi-task Regression

• Universal text representations let one model train on diverse regression tasks and support pretraining followed by fine-tuning, but may require more data than specialized GNNs with graph-specific inductive bias.Source: §3.1 · p. 3

3.2. Multi-Objective Modeling

• Autoregressive decoding models multiple metrics conditionally, including latency followed by accuracy or latencies across hardware, rather than treating outputs as conditionally independent.Source: §3.2 · p. 3

4.1. High-Level Programming Datasets

• APPS Leetcode provides Python solutions and input-output pairs; 99K solutions ran in a minimal sandbox, with peak memory as the primary metric.The dataset contains 10K problems, 232.4K ground-truth solutions, and 131.7K test cases.Source: §4.1 · p. 3
• 12,652 of KernelBook’s 18.2K Triton kernels ran successfully and were profiled on one NVIDIA A6000.KernelBook pairs PyTorch programs with Triton kernels produced by TorchInductor.Source: §4.1 · p. 3
• Table 3 evaluates Spearman ρ on 1,024 programs per language and includes 24 CodeNet languages after filtering those with insufficient test examples.The high-level programming datasets predict memory or execution latency on fixed hardware.Source: §4.1 · p. 3, Table 3 · p. 4
• CodeNet’s 7.3M accepted entries span several languages and provide memory values, but lack program inputs needed for zero-shot prediction on new questions and submissions.Because train and test splits share questions, the setup permits limited-information evaluation using same-question training submissions.Source: §4.1 · p. 4

4.2. NAS Datasets

• NAS predicts trained-network accuracy, and the study represents architectures as ONNX graphs containing operations and connectivity information.The authors describe ONNX as universal for neural-network or computation graphs and transferable to new networks.Source: §4.2 · p. 4
• The collected NAS spaces comprise 520K unique architectures in unified text-based ONNX, with accuracy, FLOPs, parameter counts, and latencies collated.The authors also create SNAS, an 85.5K-architecture CIFAR-10 space trained for 32 steps as a pretraining space.Source: §4.2 · p. 4

5. Experiments

• The study jointly trains on the datasets' training splits; its ICL prompts concatenate labeled (x, y) examples before the test input and report Spearman ρ.The number of in-context examples varies across k ∈ {0, 3, 5, 20, MAX}.Source: §5 · p. 4, §5.1 · p. 5

5.1. In-context Regression with LLMs

• GPT-5-Thinking ICL improves as examples increase but remains below the in-weight RLM across domains, with APPS performance saturating below it.The comparison uses k ∈ {0, 3, 5, 20, MAX} examples; MAX-shot is approximately 50 examples.Source: §5.1 · p. 5, Figure 3 · p. 5
• MAX-shot ICL uses about 2.7 × 10^5 input tokens per query, versus typically 10^3–10^4 raw-input tokens for RLM inference, raising ICL's marginal cost at scale.Long code and ONNX strings also limit the number of labeled examples that fit in the context window.Source: §5.1 · pp. 5–6

5.2. High-Level Programming Languages

• ρ > 0.9 on APPS Leetcode peak-memory; over half of APPS problems exceed 0.54 in-problem ρ, with predictions spanning 10^-2–10^6.The RLM also distinguishes memory consumption between substantially different solutions to the same problem.Source: §5.2 · p. 6

5.3. NAS Results

• The ONNX-input RLM remains competitive with FLAN, outperforms Arch2Vec and CATE, and predicts accuracy alongside hardware-specific latencies.The NASBench-201 evaluation includes latency objectives for Pixel3, Eyeriss, Intel CPU, and Nvidia GPU.Source: §5.3 · p. 7, Figure 6 · p. 7

6.1. Comparing with Regression Heads

• Spearman’s ρ is 0.800 for the decoder head, versus 0.717 for the normalized and 0.478 for the unnormalized regression head.The decoder head also avoids per-dataset normalization.Source: §6.1 · p. 7, Table 5 · p. 8

6.2. Scaling Regression Language Models

• The 600M-parameter T5Gemma model performs better than the smaller pretrained models, but requires extensive hyperparameter tuning and could not be run under limited compute.The authors leave further scaling analysis for future work.Source: §6.2 · p. 8, Table 6 · p. 8

6.3. Encoder-Decoder Settings

• The learned compact ONNX-aware tokenizer improves rank performance over T5’s default tokenizer by merging frequent operator strings and reducing token counts.Fewer tokens allow longer graphs per sequence.Source: §6.3 · p. 8, Table 7 · p. 8
• Spearman-ρ rises from 0.819 at 1K sequence length to 0.838 at 4K with the learned tokenizer and same training procedure.The authors attribute the improvement to reading more graph information.Source: §6.3 · p. 8, Table 8 · p. 8
• Digit-by-digit numeric tokenization outperforms T5Gemma’s regular numeric tokenizer; decoder pretraining helps only the regular tokenizer, while digit-tokenizer performance is unchanged.The authors conclude that the pretrained encoder suffices when using digit-by-digit tokenization.Source: §6.3 · p. 8, Table 9 · p. 8

7. Conclusion

• RLMs regress metrics across programming languages and code representations directly from text, without post-processing or feature engineering.The paper identifies program search, hardware-software co-design, and compiler optimization as applications.Source: §7 · p. 8
• Predicting numeric outcomes of entire experiments from raw code remains an open question for future work.The authors leave this broader use of code-based RLMs for future work.Source: §7 · p. 8

A.1. Decoder-only Regressors: Qwen and Llama

• On APPS memory prediction, the T5Gemma-based RLM reaches Spearman ρ ≈0.90, versus ≈0.17 for Llama-3.2-1B and ≈0.15 for Qwen3-0.6B.The decoder-only baselines used 128 sampled numeric generations per input and reported the median parsed value.Source: §A.1 · p. 13
• Without constrained decoding and a digit-centric vocabulary, decoder-only baselines often emit non-numeric text, making precise regression sample-inefficient and brittle to prompt formatting.The paper reports the encoder–decoder with specialized numeric output as substantially more sample-efficient for long, structured code or graph inputs.Source: §A.1 · p. 13

• Language-pretrained models achieve lower validation losses and better regression metrics on Triton kernels; freezing the encoder does not affect the run and is significantly cheaper.Figure 7 compares training from the T5Gemma checkpoint with random initialization.Source: §A.2 · p. 13, Figure 7 · p. 13
• Pretraining on 381K NASBench-101 examples nearly perfectly predicts FLOPS, and that checkpoint accelerates convergence and raises final Spearman-ρ when predicting accuracy from the same examples.Figure 9 shows FLOPS predictions; Figure 8 shows the accuracy-training curves.Source: §A.2 · p. 13, Figure 8 · p. 13, Figure 9 · p. 14

A.3. Training on Code and NAS

• Adding NAS graph data to code training does not significantly harm coding-benchmark ranking, and KernelBook validation Spearman ρ remains stable throughout training.Table 10 compares code-only and code-plus-NAS training on coding benchmarks.Source: §A.3 · p. 14, Table 10 · p. 14, Figure 10 · p. 14

A.4. Pretraining Diversity and Impact on Language

• Increasing pretraining language diversity does not clearly improve results on Go, Haskell, and Rust, and sometimes worsens results on zero-shot languages.The authors hypothesize that problem-content variety may help more than syntactic diversity, given CodeNet’s structure and T5Gemma’s code-language representations.Source: §A.4 · pp. 14–15, Table 11 · p. 15

A.5. Fine-tuning

• Fine-tuning on 1,024 language-specific examples does not necessarily help when that task was already well represented in pretraining, whereas 1K NAS examples yield +0.35 Spearman ρ on Amoeba and ENAS.The authors hypothesize that over-focusing on a language may harm general reasoning and regression abilities; Amoeba and ENAS together were only 0.08% of pretraining data.Source: §A.5 · p. 15, Table 13 · p. 15, Table 12 · p. 15

B.1. Limited Information Scenario

• 0.08 ρ drop: omitting the problem statement does not significantly harm APPS code-latency predictions.Table 14 compares models trained with problem and code against models observing code only, across 1024 programs per language.Source: §B.1 · p. 15, Table 14 · p. 16

B.2. Ranking

• Figure 11 varies the top percentile to assess whether the RLM identifies fast submissions when several solutions have similar or identical implementations.The figure plots the fraction of problems with more than eight solutions where the predicted best solution lies within the true top-p% of solutions, against a random-pick baseline.Source: §B.2 · p. 16, Figure 11 · p. 16

C. Experimental Settings

• Training uses Adafactor with linear warmup and cosine decay; inference takes the median of 64 decoder samples, and encoder inputs are capped at 2048 tokens.Sequences beyond the limit were truncated only for ONNX NAS graphs, without significant performance harm according to the cited table.Source: C. Experimental Settings · p. 16

D.1. y-Value Distributions

• Target values span widely varying ranges, motivating per-example cross-entropy loss and a decoder head that requires no y-normalization.The appendix describes the ranges as extending from 10−1 to 105 orders of magnitude and notes that MSE would make loss magnitudes highly variable.Source: §D.1 · p. 17, Figure 12 · p. 17

D.2. Extra Leetcode Data for APPS

• APPS includes 600 additional EffiBench examples, timed by repeatedly running solutions on generated test cases and averaging across iterations and trials.The examples are LeetCode problems and submissions.Source: §D.2 · p. 17

D.3. Generating the SNAS Dataset

• SNAS examples are sampled CNN architectures trained on CIFAR-10 under a fixed budget, with top-1 accuracy as the primary label and architecture and training metadata logged.Sampling varies macro configurations and micro-cell DAG operations; training uses SGD with momentum and a cosine learning-rate schedule.Source: §D.3 · p. 17

D.4.1. APPS

• APPS profiling targets Python-heap peak allocations from tracemalloc; peak RSS is a noisier secondary metric that also captures native allocations and allocator effects.The wrapper profiles APPS train-split solutions and records timing, allocation, and execution counters.Source: §D.4.1 · p. 18

D.4.2. KERNELBOOK

• KernelBook Triton kernels are profiled on an NVIDIA A6000 after JIT warmup, using five trials with at least a one-second timing window and median latency as the reported metric.Inputs come from dataset-provided constructors and activations; failed runs such as OOM or shape mismatches are skipped.Source: §D.4.2 · pp. 18–19

D.5. Triton Code Sample (KernelBook)

• Triton Kernel Example #4949 is listed with a latency of 0.0152 ms.The appendix pairs the kernel code with a DiceLossNew wrapper.Source: §D.5 · pp. 19–20

D.6. ONNX Graph Code Sample

• SNAS Architecture #10 is shown as an ONNX graph with an accuracy of 60.93%.The appendix includes portions of the graph representation.Source: §D.6 · p. 21

D.7. Example Code Submissions

• 5608 bytes ground truth: the one-pass solution uses O(1) extra space, versus 7136 bytes for the three-array version, and the RLM predicts 5549 and 6228 bytes.The larger implementation also materializes list(enumerate(arr)) and creates large temporary strings in debug prints.Source: Figure 13 · p. 22
• 5508 bytes ground truth: the O(1)-space circular-subarray solution uses less memory than the O(n)-space version, and the RLM estimates 5634 versus 6430 bytes.Both solutions run in O(n) time; the larger version allocates several length-n arrays and reverses the array in place.Source: Figure 14 · p. 23
• 6518 bytes estimated: reusing one frequency map uses less memory than materializing and sorting lists with repeated Counters, estimated at 10197 bytes.Ground-truth memory is 6178 and 10588 bytes, respectively.Source: Figure 15 · p. 24