Technische Samenvatting: Regressie Taalmodellen voor Code

Probleemstelling

Het artikel behandelt de uitdaging van code-naar-maatstaf regressie: het voorspellen van numerieke uitkomsten van code-uitvoeringen (zoals geheugengebruik, latentie of nauwkeurigheid) rechtstreeks vanuit broncode of representaties van computationele grafieken. Historisch gezien werd deze taak gehinderd door de open-ended aard van programmeertalen en de afhankelijkheid van zware, domeinspecifieke feature engineering. Eerdere methoden vereisten vaak het converteren van code naar rigide grafische representaties (bijv. syntaxisbomen, burenmatrices) of tabulaire vectoren via handmatig tellen van commando's, wat de overdraagbaarheid naar nieuwe taken, talen of kernel-types beperkte. De kernmoeilijkheid ligt in het leren van accurate regressiemodellen over sterk niet-tabulaire, grafische data zonder de flexibiliteit te verliezen om te generaliseren over diverse domeinen.

Methodologie

De auteurs stellen een Regressie Taalmodel (RLM) voor, een unificerend kader dat code-naar-maatstaf regressie behandelt als een standaard next-token voorspellingprobleem, in lijn met het moderne Large Language Model (LLM) paradigma.

Architectuur en Training

• Gecombineerde Encoder-Decoder: Het RLM maakt gebruik van een bevroren encoder (specifiek T5Gemma-S) om invoer-tekst te verwerken en een trainbare decoder om numerieke uitkomsten te genereren. De encoder is voorgetraind op standaard taaldata, terwijl alleen de decoder backpropagation ondergaat, wat de trainingskosten aanzienlijk verlaagt.
• Tekst-naar-Tekst Regressie: In tegenstelling tot traditionele benaderingen die invoer afbeelden op vaste lengte vectoren, accepteert het RLM ruwe tekstrepresentaties van code (bijv. Python, C++, ONNX-grafieken) en geeft numerieke waarden uit als reeksen tokens.
• Gespecialiseerde Numerieke Tokenisatie: Om het brede dynamische bereik van maatstaven (van $10^{-2}$ tot $10^6$) zonder normalisatie te hanteren, gebruikt het model een expliciete tokenisatie per cijfer. Getallen worden weergegeven met speciale tokens voor teken, exponent en mantisse (bijv. <+><-><1><7><2><5> voor $72.5$). Dit voorkomt numerieke instabiliteit en de noodzaak van dataset-specifieke schaling.
• Multi-Doel Modellering: Door gebruik te maken van de autoregressieve aard van de decoder, kan het model conditioneel meerdere maatstaven sequentieel voorspellen (bijv. $p(y_{latentie} | y_{nauwkeurigheid}, x)$), waardoor inherente correlaties tussen doelen zoals nauwkeurigheid en hardware-latentie worden vastgelegd.

Data Bronnen

Het model is getraind op een diverse, gecombineerde dataset bestaande uit:

1. Hoog-niveau Programmeren: Code van APPS (LeetCode problemen), CodeNet (14M samples over 37 talen) en KernelBook (Triton GPU kernels).
2. Neurale Architectuur Zoeken (NAS): 520K unieke architecturen uit benchmarks zoals NASBench-101/201, FBNet en Once-for-All, weergegeven als ONNX intermediaire representaties.

Belangrijkste Bijdragen

1. Gecombineerd Algemeen Doel Model: Het artikel toont aan dat een enkel RLM tegelijkertijd geheugengebruik over meerdere talen (Python, C++, enz.), GPU kernel latentie en neurale netwerknauwkeurigheid/latentie kan voorspellen, waardoor de noodzaak voor taak-specifieke feature engineering wordt geëlimineerd.
2. Superioriteit ten opzichte van Grafische Methoden: Ondanks het lezen van dichte ONNX tekstrepresentaties, presteert het RLM beter dan state-of-the-art Graph Neural Network (GNN) baselines op NAS benchmarks (bijv. het behalen van een Kendall-Tau van 0,46 op ontwerpruimtes waar GNN's eerder domineerden).
3. Efficiëntie en Schaalbaarheid: Het gebruik van een bevroren encoder maakt kosteneffectieve training mogelijk. De auteurs tonen aan dat een relatief klein model van 300M parameters hoge prestaties bereikt, en dat schalen naar 600M parameters verdere winst oplevert.
4. Voorspelling zonder Normalisatie: De tokenisatiestrategie per cijfer stelt het model in staat waarden over ordes van grootte te voorspellen zonder handmatige normalisatie, een significant voordeel ten opzichte van op MSE gebaseerde regressiehoofden.

Experimentele Resultaten

• Code Prestaties: Op de APPS dataset bereikt het RLM een Spearman rang correlatie ($\rho$) van 0,93 voor geheugenvoorspelling. Op CodeNet bereikt het een gemiddelde $\rho > 0,5$ over 17 talen, waarbij C++ een score van 0,748 bereikt.
• NAS Prestaties: Het RLM bereikt een gemiddelde Kendall-Tau van 0,46 over vijf klassieke NAS ontwerpruimtes, en presteert beter dan eerdere SOTA methoden zoals FLANT en op GNN gebaseerde encoders. Het slaagt erin zowel nauwkeurigheid als hardware-specifieke latenties (bijv. op Pixel3, Eyeriss, Intel CPU, Nvidia GPU) tegelijkertijd te voorspellen.
• Ablatie Studies:
  • Bevroren Encoders: Voorgetrainde encoders verbeteren convergentie en uiteindelijke prestaties aanzienlijk ten opzichte van willekeurige initialisatie.
  • Decoder versus Regressiehoofden: De decoder-gebaseerde output (0,800 Spearman $\rho$) presteert aanzienlijk beter dan expliciete regressiehoofden (0,478 zonder normalisatie, 0,717 met normalisatie).
  • Tokenisatie: Aangepaste, compacte tokenizers voor ONNX-grafieken verbeteren de prestaties (0,723 versus 0,533 met standaard tokenizers).
  • In-Context Learning (ICL): Hoewel ICL met grote modellen (bijv. GPT-5) verbetert met meer voorbeelden, blijft het consequent onder de in-gewichten getrainde RLM en brengt het aanzienlijk hogere inferentiekosten met zich mee door beperkingen in het contextvenster.

Betekenis en Claims

Het artikel claimt de weg te effenen voor het vereenvoudigen van regressie van computationele grafieken tot een generiek next-token voorspellingprobleem. Door aan te tonen dat taalmodellen effectief kunnen regresseren over pure tekstrepresentaties van code en grafieken, suggereert het werk dat het vakgebied kan afstappen van rigide, handgemaakte feature engineering naar een flexibeler, unificerende aanpak.

De auteurs positioneren dit als een waardevol hulpmiddel voor:

• Programma Zoeken: Versnellen van de selectie van efficiënte code-oplossingen.
• Hardware-Software Co-ontwerp: Latentie voorspellen over diverse hardwareplatforms.
• Compiler Optimalisatie: Prestatiemaatstaven schatten voor compiler passes.

Het artikel blijft bescheiden wat betreft bredere implicaties, met de opmerking dat hoewel de methode veelbelovend is voor het voorspellen van numerieke uitkomsten uit ruwe code, het uitbreiden hiervan naar het voorspellen van uitkomsten van volledige wetenschappelijke experimenten een open vraag blijft voor toekomstig werk. De primaire bijdrage is het vaststellen van het algemene vermogen van taalmodellen om te fungeren als universele regressoren voor code en computationele grafieken zonder nabewerking of gespecialiseerde feature extractie.