Resumen Técnico: Modelos de Lenguaje de Regresión para Código

Declaración del Problema

El artículo aborda el desafío de la regresión código-métrica: predecir resultados numéricos de ejecuciones de código (como huella de memoria, latencia o precisión) directamente a partir del código fuente o representaciones de grafos computacionales. Históricamente, esta tarea ha sido obstaculizada por la naturaleza abierta de los lenguajes de programación y la dependencia de una ingeniería de características pesada y específica del dominio. Los métodos anteriores a menudo requerían convertir el código en representaciones de grafos rígidas (por ejemplo, árboles de sintaxis, matrices de adyacencia) o vectores tabulares mediante un conteo manual de comandos, lo que limitaba la transferibilidad a nuevas tareas, lenguajes o tipos de kernel. La dificultad central radica en aprender modelos de regresión precisos sobre datos altamente no tabulares y basados en grafos sin perder la flexibilidad para generalizar a través de diversos dominios.

Metodología

Los autores proponen un Modelo de Lenguaje de Regresión (RLM), un marco unificado que trata la regresión código-métrica como un problema estándar de predicción del siguiente token, alineándose con el paradigma moderno de los Modelos de Lenguaje Grande (LLM).

Arquitectura y Entrenamiento

• Codificador-Decodificador Unificado: El RLM utiliza un codificador congelado (específicamente T5Gemma-S) para procesar el texto de entrada y un decodificador entrenable para generar salidas numéricas. El codificador está preentrenado en datos de lenguaje estándar, mientras que solo el decodificador sufre retropropagación, reduciendo significativamente los costos de entrenamiento.
• Regresión Texto-a-Texto: A diferencia de los enfoques tradicionales que mapean entradas a vectores de longitud fija, el RLM acepta representaciones de texto crudo del código (por ejemplo, Python, C++, grafos ONNX) y genera valores numéricos como secuencias de tokens.
• Tokenización Numérica Especializada: Para manejar el amplio rango dinámico de las métricas (de $10^{-2}$ a $10^6$) sin normalización, el modelo utiliza una tokenización explícita dígito a dígito. Los números se representan utilizando tokens especiales para el signo, el exponente y la mantisa (por ejemplo, <+><-><1><7><2><5> para $72.5$). Esto evita la inestabilidad numérica y la necesidad de escalado específico del conjunto de datos.
• Modelado Multiobjetivo: Aprovechando la naturaleza autoregresiva del decodificador, el modelo puede predecir condicionalmente múltiples métricas secuencialmente (por ejemplo, $p(y_{latencia} | y_{precisión}, x)$), capturando correlaciones inherentes entre objetivos como la precisión y la latencia del hardware.

Fuentes de Datos

El modelo se entrena en un conjunto de datos diverso y unificado que comprende:

1. Programación de Alto Nivel: Código de APPS (problemas de LeetCode), CodeNet (14M de muestras en 37 lenguajes) y KernelBook (kernels GPU Triton).
2. Búsqueda de Arquitectura Neural (NAS): 520K arquitecturas únicas de puntos de referencia como NASBench-101/201, FBNet y Once-for-All, representadas como representaciones intermedias ONNX.

Contribuciones Clave

1. Modelo Unificado de Propósito General: El artículo demuestra que un único RLM puede predecir simultáneamente el uso de memoria en múltiples lenguajes (Python, C++, etc.), la latencia de kernels GPU y la precisión/latencia de redes neuronales, eliminando la necesidad de ingeniería de características específica de la tarea.
2. Superioridad sobre Métodos Basados en Grafos: A pesar de leer representaciones de texto densas de ONNX, el RLM supera a las líneas base de Redes Neuronales de Grafos (GNN) más avanzadas en puntos de referencia de NAS (por ejemplo, logrando un Kendall-Tau de 0.46 en espacios de diseño donde las GNN dominaban previamente).
3. Eficiencia y Escalabilidad: El uso de un codificador congelado permite un entrenamiento rentable. Los autores muestran que un modelo relativamente pequeño de 300M de parámetros logra un alto rendimiento, y escalar a 600M de parámetros produce ganancias adicionales.
4. Predicción Sin Normalización: La estrategia de tokenización dígito a dígito permite al modelo predecir valores a través de órdenes de magnitud sin normalización manual, una ventaja significativa sobre las cabezas de regresión basadas en MSE.

Resultados Experimentales

• Rendimiento del Código: En el conjunto de datos APPS, el RLM alcanza una correlación de rango de Spearman ($\rho$) de 0.93 para la predicción de memoria. En CodeNet, logra un $\rho$ promedio > 0.5 en 17 lenguajes, con C++ alcanzando 0.748.
• Rendimiento en NAS: El RLM alcanza un Kendall-Tau promedio de 0.46 en cinco espacios de diseño clásicos de NAS, superando a métodos SOTA anteriores como FLANT y codificadores basados en GNN. Predice con éxito tanto la precisión como las latencias específicas del hardware (por ejemplo, en Pixel3, Eyeriss, CPU Intel, GPU Nvidia) simultáneamente.
• Estudios de Ablación:
  • Codificadores Congelados: Los codificadores preentrenados mejoran significativamente la convergencia y el rendimiento final en comparación con la inicialización aleatoria.
  • Decodificador vs. Cabezas de Regresión: La salida basada en decodificador (0.800 Spearman $\rho$) supera sustancialmente a las cabezas de regresión explícitas (0.478 sin normalización, 0.717 con normalización).
  • Tokenización: Los tokenizadores personalizados y compactos para grafos ONNX mejoran el rendimiento (0.723 frente a 0.533 con tokenizadores predeterminados).
  • Aprendizaje en Contexto (ICL): Aunque el ICL con modelos grandes (por ejemplo, GPT-5) mejora con más ejemplos, permanece consistentemente por debajo del RLM entrenado en pesos y conlleva costos de inferencia significativamente mayores debido a las limitaciones de la ventana de contexto.

Significado y Afirmaciones

El artículo afirma allanar el camino para simplificar la regresión de grafos computacionales en un problema genérico de predicción del siguiente token. Al demostrar que los modelos de lenguaje pueden regredir efectivamente sobre representaciones de texto puro de código y grafos, el trabajo sugiere que el campo puede alejarse de la ingeniería de características rígida y manual hacia un enfoque más flexible y unificado.

Los autores posicionan esto como una herramienta valiosa para:

• Búsqueda de Programas: Acelerar la selección de soluciones de código eficientes.
• Co-diseño Hardware-Software: Predecir la latencia en diversas plataformas de hardware.
• Optimización de Compiladores: Estimar métricas de rendimiento para pasadas de compilador.

El artículo se mantiene modesto respecto a implicaciones más amplias, señalando que, aunque el método muestra promesa para predecir resultados numéricos a partir de código crudo, extender esto para predecir resultados de experimentos científicos completos sigue siendo una pregunta abierta para trabajos futuros. La contribución principal es establecer la capacidad general de los modelos de lenguaje para servir como regresores universales para código y grafos computacionales sin postprocesamiento ni extracción de características especializada.