Resumen Técnico: TRACE: Aprendiendo a Computar en Grafos de Circuitos

1. Planteamiento del Problema

El desafío fundamental abordado por este trabajo es el aprendizaje para computar (learning to compute), definido como la capacidad de modelar con precisión el comportamiento funcional de los grafos de circuitos. Los grafos de circuitos (incluyendo grafos de Transferencia de Registro a Nivel de RTL, Grafos de Compuertas y Negación [AIG] y listas de conexión post-mapeo) capturan el flujo de la computación a través de nodos dirigidos que representan operaciones y variables.

Los autores argumentan que los paradigmas dominantes en el aprendizaje de representaciones de grafos —Redes Neuronales de Paso de Mensajes (MPNN) y Transformers de Grafos— están arquitectónicamente desajustados para esta tarea específica debido a tres fallas críticas:

1. Las MPNN no logran modelar las interacciones de entrada: Las MPNN dependen de agregadores invariantes a la permutación (por ejemplo, suma, media) que combinan los mensajes de los vecinos de forma independiente. Esto impide que el modelo capture interacciones dependientes de la posición y específicas del operador. Por ejemplo, un multiplexor (MUX) donde $MUX(S, A, B) \neq MUX(A, S, B)$ no puede modelarse correctamente porque el orden de las entradas se pierde durante la agregación.
2. Los Transformers Vanilla carecen de conciencia jerárquica: Los Transformers estándar aplanan el grafo en una secuencia, destruyendo la estructura jerárquica explícita y las dependencias intermedias necesarias para modelar la computación paso a paso (por ejemplo, al no capturar la dependencia de $f(g(a,b), c)$ sobre el resultado intermedio $g(a,b)$).
3. Los Transformers con Máscara de Aristas heredan las limitaciones de las MPNN: Los Transformers que utilizan máscaras de atención para respetar la topología del grafo a menudo se reducen a sumas ponderadas (similares a las Redes de Atención de Grafos), heredando la incapacidad de modelar interacciones precisas y no lineales entre múltiples entradas.

Además, existe un desafío más profundo al capturar funciones globales que emergen de topologías complejas, particularmente las dependencias reconvergentes. En los circuitos, la función global de un nodo puede diferir significamente de una simple aproximación local que asume la independencia de las entradas. Los métodos anteriores que supervisan directamente la función global agrupan implícitamente estos efectos contextuales complejos en los embeddings, lo que dificulta que los modelos distingan las dependencias funcionales reales de las correlaciones espurias.

2. Metodología

Los autores proponen TRACE (Transformer for Reasoning about Algebraic and Computational Expressions), un nuevo paradigma que consta de dos componentes principales: un esqueleto arquitectónicamente sólido y un objetivo de aprendizaje principista.

2.1 Esqueleto de Transformer Jerárquico

Para resolver las fallas arquitectónicas de los modelos anteriores, TRACE emplea un Transformer Jerárquico que refleja el flujo paso a paso de la computación.

• Computación a Nivel Lógico: Los nodos se procesan nivel por nivel basándose en su profundidad lógica topológica. Las Entradas Primarias (PIs) se encuentran en el nivel 0.
• Representación en Notación Prefija: Inspirado en la notación prefija, cada paso de computación en un nivel lógico específico se representa como una secuencia ordenada: [operador, entrada_1, entrada_2, ...].
• Codificación Sensible a la Posición: Un codificador Transformer procesa estas secuencias con codificaciones posicionales. Esto permite al modelo capturar interacciones específicas del operador y la conciencia de posición (por ejemplo, distinguir la señal de control de las entradas de datos en un MUX).
• Procesamiento Recursivo: El embedding de un nodo operador se actualiza basándose en la salida del Transformer para el primer token de su secuencia. Este embedding actualizado sirve entonces como entrada para los operadores subsiguientes en el siguiente nivel lógico.
• Eficiencia: A diferencia de los grafos generales con distribuciones de grado de cola larga, los grafos de circuitos tienen grados de entrada pequeños y acotados, lo que evita la sobrecarga de memoria cuadrática típicamente asociada con los Transformers.

2.2 Aprendizaje de Desplazamiento de Función (FSL)

Para abordar la dificultad de predecir funciones globales complejas directamente, TRACE introduce el Aprendizaje de Desplazamiento de Función (Function Shift Learning), un objetivo novedoso que desacopla el problema de aprendizaje.

• Función Global ($y_{global}$): La verdadera probabilidad de que un nodo sea lógica-1, la cual tiene en cuenta todas las dependencias (incluyendo la reconvergencia). Calcular esto directamente es costoso ($O(2^k)$).
• Función Local ($y_{local}$): Una aproximación simple que asume la independencia de las entradas. Es computacionalmente eficiente ($O(1)$) pero falla al capturar el comportamiento real en presencia de reconvergencia.
• Desplazamiento de Función ($y_{FSL}$): El modelo es entrenado para predecir únicamente la discrepancia entre la función global y la local:
  $$y_{FSL} = y_{global} - y_{local}$$
• Entrenamiento: El modelo se entrena para realizar la regresión del desplazamiento de la función utilizando una pérdida L1. El desplazamiento real (ground truth) se precalcula a partir de los datos de entrenamiento.
• Inferencia: Durante la inferencia, la función global se reconstruye iterativamente. Para cada nodo, el modelo predice el desplazamiento, el cual se suma a la función local (calculada usando las funciones globales estimadas de sus predecesores) para obtener la estimación global final.

3. Contribuciones Clave

1. Alineación Arquitectónica: La introducción de un Transformer Jerárquico que reemplaza la agregación invariante a la permutación con un procesamiento basado en secuencias y sensible a la posición, reflejando fielmente el flujo computacional de los grafos de circuitos.
2. Objetivo de Aprendizaje Desacoplado: La propuesta de Aprendizaje de Desplazamiento de Función, que aísla los efectos contextuales complejos de la reconvergencia (el "desplazamiento") de la aproximación local simple, simplificando la tarea de aprendizaje.
3. Rendimiento de Vanguardia (State-of-the-Art): La demostración de que este paradigma combinado supera sustancialmente a las MPNN y a los Transformers de Grafos existentes en diversas modalidades de circuitos.

la 4. Resultados Experimentales

Los autores evaluaron TRACE en tres modalidades de circuitos (RTL, AIG, PM Netlist) en dos tipos de tareas: aprendizaje contrastivo y tareas predictivas.

4.1 Tareas Contrastivas

En tareas de recuperación (identificación de circuitos funcionalmente equivalentes), TRACE superó consistentemente a todos los modelos de referencia:

• RTL: Logró un 94.45% de Recall@1 (frente al ~88% del siguiente mejor).
• AIG: Logró un 92.68% de Recall@1 (frente al ~88% del siguiente mejor).
• PM Netlist: Logró un 90.81% de Recall@1 (frente al ~86% del siguiente mejor).
• Notablemente, muchos modelos base de Transformers de Grafos (por ejemplo, GraphGPS, SGFormer) sufrieron una degradación severa del rendimiento o errores de falta de memoria (OOM) en AIGs y PM Netlists, mientras que TRACE se mantuvo estable.

4.2 Tareas Predictivas

En tareas de predicción de propiedades a nivel de nodo (probabilidad de Lógica-1, similitud, probabilidad de transición), TRACE logró el menor Error Absoluto Medio (MAE) y los puntajes $R^2$ más altos:

• AIG Combinacional: MAE de 0.015 para la probabilidad de Lógica-1 (frente a 0.027 para DeepGate4).
• AIG Secuencial: MAE de 0.009 para la probabilidad de Lógica-1 y 0.005 para la probabilidad de transición.
• PM Netlists: MAE de 0.013 para la probabilidad de Lógica-1 (frente a 0.048 para GraphSAGE).

4.3 Ablación y Análisis

• Aprendizaje de Desplazamiento de Función: Eliminar FSL degradó significativamente el rendimiento. Por ejemplo, en la predicción de probabilidad de PM Netlist, eliminar FSL aumentó el MAE de 0.013 a 0.036 (un incremento relativo de error del 63.9%).
• Acumulación de Error: El análisis a través de los niveles lógicos mostró que TRACE mantiene un rendimiento estable incluso en niveles lógicos profundos, lo que indica robustez contra la acumulación de errores durante el proceso de inferencia iterativa.

5. Significado

El artículo afirma que TRACE establece un nuevo estado del arte para el análisis de circuitos y, más ampliamente, proporciona un paradigma más robusto y arquitectónicamente sólido para el aprendizaje en grafos de circuitos. Al abordar el desajuste fundamental entre las arquitecturas estándar de aprendizaje de grafos y la naturaleza de la computación (específicamente la sensibilidad a la posición y la jerarquía), TRACE demuestra que alinear el sesgo inductivo del modelo con la estructura del problema es crítico para capturar el comportamiento funcional. El trabajo sirve como una prueba de concepto de que el desacoplamiento de las aproximaciones locales de los desplazamientos contextuales globales conduce a una generalización superior a través de diversos dominios de la automatización de diseño electrónico (EDA).