Generación de Código Tipado con LLMs: Más Allá de la Inferencia Post-Hoc

Los Large Language Models (LLMs) han demostrado una capacidad sorprendente para generar código, incluso en lenguajes con sistemas de tipos avanzados como Idris o Agda. Sin embargo, la mayoría de los LLMs actuales están entrenados para predecir la siguiente secuencia de tokens sin una comprensión inherente de la estructura de tipos del lenguaje objetivo. Esto conduce a un proceso de generación ineficiente y propenso a errores, donde la verificación de tipos se realiza como un paso posterior a la inferencia, a menudo mediante bucles de reintento o decodificación restringida.

El problema fundamental que aborda este artículo es cómo integrar la noción de tipos y estructura directamente en el proceso de entrenamiento de los LLMs, en lugar de tratarla como una restricción externa. La tesis central es que, al permitir que los modelos aprendan a diferenciar 'con respecto a la estructura' (es decir, aprender las decisiones de tipo discretas como parte del entrenamiento), se puede lograr una generación de código más eficiente, precisa y semánticamente coherente. Esto contrasta con los enfoques actuales que diferencian 'a través de la estructura' (donde las decisiones de tipo son andamiajes fijos) o que simplemente aplican la verificación de tipos en la inferencia, lo que no permite al modelo aprender de sus errores de tipo durante el entrenamiento.

La relevancia de este problema es creciente a medida que los LLMs se utilizan para tareas de programación más complejas y críticas. La ineficiencia de los métodos actuales de post-procesamiento se vuelve una barrera para la escalabilidad y la fiabilidad. La propuesta de integrar los tipos en el entrenamiento busca emular cómo los sistemas como AlphaZero logran un rendimiento superior al incorporar las reglas del dominio (en este caso, las reglas de un juego) directamente en su proceso de aprendizaje, en lugar de aplicarlas solo en la inferencia.

Arquitectura del Sistema

Los LLMs tradicionales operan con una arquitectura de transformador que consume una secuencia de tokens de entrada y produce una distribución sobre el siguiente token. El tipo de la función es típicamente LLM : List Token -> D(List Token). Durante la inferencia, se muestrea un token de esta distribución y se realimenta al modelo. Los enfoques actuales para la generación de código tipado se centran en modificar este proceso de inferencia.

Dos métodos principales se discuten: 1) 'Try; Compile; If Error Repeat': El modelo genera una secuencia completa de tokens hasta un token STOP. Luego, un typechecker externo evalúa la validez del código. Si hay un error, el mensaje de error se retroalimenta al modelo para un nuevo intento. Este enfoque tiene baja granularidad (la verificación es al final) y alto ancho de banda (mensajes de error estructurados). Sin embargo, es ineficiente ya que los errores tempranos no se detectan hasta el final de la generación, y el modelo no aprende de los errores, ya que los pesos no se actualizan. 2) 'Constrained Decoding': Antes de muestrear cada token, un typechecker consulta el estado actual de la generación y enmascara (establece la probabilidad a cero) cualquier token que llevaría a un resultado mal tipado. Este enfoque tiene máxima granularidad (cada token) y mínimo ancho de banda (una señal binaria de aceptar/rechazar). Garantiza la corrección de tipos, pero no cambia lo que el modelo 'quiere decir', forzándolo a muestrear tokens de baja probabilidad y potencialmente produciendo código sintácticamente correcto pero semánticamente sin sentido. Tampoco hay flujo de gradiente a través de la máscara, por lo que el modelo no aprende a asignar mayor probabilidad a ramas bien tipadas.

La propuesta de 'diferenciar con respecto a la estructura' introduce una arquitectura donde el modelo aprende las decisiones de tipo discretas. En lugar de una única red, se utilizan tres mapas diferenciables: f_c : X × P → D(Bool) para decidir la rama de tipo (por ejemplo, Left o Right en un coproducto Either A B), f_A : X × P → A para generar el valor si se elige la rama A, y f_B : X × P → B para la rama B. La salida final es una distribución sobre los tipos y sus valores, f : X → D(A + B). Durante el entrenamiento supervisado, solo se diferencia a través de la rama tomada, utilizando una función de pérdida de entropía cruzada. Esto permite que la red aprenda la partición del espacio de entrada y las ramas de salida simultáneamente, produciendo código bien tipado por construcción y permitiendo que el modelo aprenda la estructura de salida. Este enfoque se generaliza a estructuras más complejas utilizando la teoría de contenedores y lentes dependientes.

Fundamentos Teóricos

La discusión sobre la diferenciación a través de estructuras discretas y la necesidad de aprender particiones del espacio de entrada se conecta con principios fundamentales de la teoría de categorías y la optimización. La descomposición de funciones en coproductos (Either A B) en Haskell, por ejemplo, refleja propiedades de la categoría Set, específicamente su naturaleza como categoría extensiva. Esto implica que cualquier función a un coproducto debe descomponerse en un predicado que particiona el dominio y funciones separadas para cada rama. Trabajos como CHAD (Categorical Higher-Order Automatic Differentiation) y Higher Order AD de Higher Order Functions exploran cómo diferenciar mapas en coproductos, pero a menudo asumen una partición fija, lo que limita la capacidad de aprendizaje del modelo.

El concepto de 'diferenciar con respecto a la estructura' se relaciona con la idea de aprender distribuciones sobre elecciones discretas, un problema común en el aprendizaje por refuerzo y los modelos generativos. Aunque el artículo no cita un paper específico para esta formulación exacta, se menciona que el 'teacher forcing' es una variante no tipada, y las Abstract Syntax Networks (2017) describen un enfoque similar pero sin tipos dependientes. La integración de reglas de dominio en el entrenamiento, como se ve en AlphaZero para juegos como el ajedrez, es un precedente fuerte para la propuesta de integrar sistemas de tipos en el entrenamiento de LLMs. AlphaZero (Silver et al., 2017) demostró cómo un modelo puede alcanzar un rendimiento sobrehumano al aprender las reglas del juego y usarlas para guiar la búsqueda y el entrenamiento, en contraste con modelos que solo aprenden de la superficie de los datos.