Riepilogo Tecnico: TRACE: Imparare a Computare su Grafi di Circuiti

1. Definizione del Problema

La sfida fondamentale affrontata da questo lavoro è l'apprendimento del computo (learning to compute), definito come la capacità di modellare accuratamente il comportamento funzionale dei grafi di circuiti. I grafi di circuiti (inclusi i grafi Register Transfer Level (RTL), gli And-Inverter Graphs (AIG) e i netlist post-mapping) catturano il flusso di computazione attraverso nodi diretti che rappresentano operazioni e variabili.

Gli autori sostengono che i paradigmi dominanti nell'apprendimento delle rappresentazioni su grafi — Message Passing Neural Networks (MPNNs) e Graph Transformers — siano architettonicamente incompatibili con questo specifico compito a causa di tre criticità fondamentali:

1. Le MPNN non riescono a modellare le interazioni tra gli input: Le MPNN si affidano ad aggregatori invarianti per permutazione (ad es. somma, media) che combinano i messaggi dei vicini in modo indipendente. Ciò impedisce al modello di catturare interazioni specifiche per l'operatore e sensibili alla posizione. Ad esempio, un multiplexer (MUX) dove $MUX(S, A, B) \neq MUX(A, S, B)$ non può essere modellato correttamente perché l'ordine degli input viene perso durante l'aggregazione.
2. I Transformer vanilla sono privi di gerarchia: I Transformer standard appiattiscono il grafo in una sequenza, distruggendo la struttura gerarchica esplicita e le dipendenze intermedie necessarie per modellare il computo passo dopo passo (ad es. fallendo nel catturare la dipendenza di $f(g(a,b), c)$ dal risultato intermedio $g(a,b)$).
3. I Transformer con maschera sugli archi ereditano i limiti delle MPNN: I Transformer che utilizzano maschere di attenzione per rispettare la topologia del grafo spesso si riducono a somme pesate (simili alle Graph Attention Networks), ereditando l'incapacità di modellare interazioni precise e non lineari tra molteplici input.

Inoltre, esiste una sfida più profonda nel catturare le funzioni globali che emergono da topologie complesse, in particolare le dipendenze riconvergenti. Nei circuiti, la funzione globale di un nodo può differire significativamente da una semplice approssimazione locale che assume l'indipendenza degli input. I metodi precedenti, che supervisionano direttamente sulla funzione globale, raggruppano implicitamente questi complessi effetti contestuali negli embedding, rendendo difficile per i modelli distinguere le vere dipendenze funzionali dalle correlazioni spurie.

2. Metodologia

Gli autori propongono TRACE (Transformer for Reasoning about Algebraic and Computational Expressions), un nuovo paradigma composto da due componenti core: un backbone architettonicamente solido e un obiettivo di apprendimento fondato su principi precisi.

2.1 Backbone Transformer Gerarchico

Per risolvere i difetti architettonici dei modelli precedenti, TRACE impiega un Transformer Gerarchico che rispecchia il flusso di computazione passo dopo passo.

• Computazione a livello logico: I nodi vengono elaborati livello per livello in base alla loro profondità logica topologica. Gli Input Primari (PI) si trovano al livello 0.
• Rappresentazione in Notazione Prefissa: Ispirandosi alla notazione prefissa, ogni passo computazionale a un determinato livello logico è rappresentato come una sequenza ordinata: [operatore, input_1, input_2, ...].
• Codifica Sensibile alla Posizione: Un encoder Transformer elabora queste sequenze con codifiche posizionali. Ciò consente al modello di catturare interazioni specifiche per l'operatore e la consapevolezza della posizione (ad es. distinguere il segnale di controllo dagli input di dati in un MUX).
• Elaborazione Ricorsiva: L'embedding di un nodo operatore viene aggiornato sulla base dell'output del Transformer per il primo token della sua sequenza. Questo embedding aggiornato funge poi da input per gli operatori successivi nel livello logico successivo.
• Efficienza: A differenza dei grafi generici con distribuzioni di grado a coda lunga, i grafi di circuiti hanno gradi di entrata piccoli e limitati, evitando l'overhead di memoria quadratico tipicamente associato ai Transformer.

2.2 Apprendimento dello Shift Funzionale (Function Shift Learning - FSL)

Per affrontare la difficoltà di prevedere direttamente funzioni globali complesse, TRACE introduce lo Function Shift Learning, un obiettivo di apprendimento innovativo che decoppia il problema.

• Funzione Globale ($y_{global}$): La vera probabilità che un nodo sia logicamente 1, che tiene conto di tutte le dipendenze (incluse le riconvergenze). Calcolarla direttamente è costoso ($O(2^k)$).
• Funzione Locale ($y_{local}$): Una semplice approssimazione che assume l'indipendenza degli input. È computazionalmente efficiente ($O(1)$) ma fallisce nel catturare il comportamento reale in presenza di riconvergenza.
• Shift Funzionale ($y_{FSL}$): Il modello viene addestrato a prevedere solo la discrepanza tra la funzione globale e quella locale:
  $$y_{FSL} = y_{global} - y_{local}$$
• Addestramento: Il modello viene addestrato per la regressione dello shift funzionale utilizzando una perdita L1. Lo shift reale (ground truth) è pre-calcolato dai dati di addestramento.
• Inferenza: Durante l'inferenza, la funzione globale viene ricostruita iterativamente. Per ogni nodo, il modello predice lo shift, che viene aggiunto alla funzione locale (calcolata usando le funzioni globali stimate dei suoi predecessori) per ottenere la stima globale finale.

3. Contributi Chiave

1. Allineamento Architetturale: L'introduzione di un Transformer Gerarchico che sostituisce l'aggregazione invariante per permutazione con l'elaborazione basata su sequenze e sensibile alla posizione, rispecchiando fedelmente il flusso computazionale dei grafi di circuiti.
2. Obiettivo di Apprendimento Decoppiato: La proposta dello Function Shift Learning, che isola gli effetti contestuali complessi della riconvergenza (lo "shift") dall'approssimazione locale semplice, semplificando il compito di apprendimento.
3. Prestazioni allo Stato dell'Arte: La dimostrazione che questo paradigma combinato supera significativamente le MPNN e i Graph Transformer esistenti in diverse modalità di circuiti.

4. Risultati Sperimentali

Gli autori hanno valutato TRACE su tre modalità di circuiti (RTL, AIG, PM Netlist) attraverso due tipi di task: apprendimento contrastivo e task predittivi.

4.1 Task Contrastivi

Nei compiti di retrieval (identificazione di circuiti funzionalmente equivalenti), TRACE ha superato costantemente tutti i baseline:

• RTL: Ha raggiunto il 94,45% di Recall@1 (rispetto a circa l'88% del secondo miglior modello).
• AIG: Ha raggiunto il 92,68% di Recall@1 (rispetto a circa l'88% del secondo miglior modello).
• PM Netlist: Ha raggiunto il 90,81% di Recall@1 (rispetto a circa l'86% del secondo miglior modello).
• Notevolmente, molti baseline di Graph Transformer (ad es. GraphGPS, SGFormer) hanno subito gravi degradazioni delle prestazioni o errori di Out-Of-Memory (OOM) su AIG e PM Netlist, mentre TRACE è rimasto stabile.

4.2 Task Predittivi

Nei compiti di previsione di proprietà a livello di nodo (probabilità Logic-1, similarità, probabilità di transizione), TRACE ha ottenuto il minor Mean Absolute Error (MAE) e i punteggi $R^2$ più elevati:

• AIG Combinatorio: MAE di 0,015 per la probabilità Logic-1 (rispetto a 0,027 per DeepGate4).
• AIG Sequenziale: MAE di 0,009 per la probabilità Logic-1 e 0,005 per la probabilità di transizione.
• PM Netlist: MAE di 0,013 per la probabilità Logic-1 (rispetto a 0,048 per GraphSAGE).

4.3 Ablation e Analisi

• Function Shift Learning: Rimuovere FSL ha degradato significativamente le prestazioni. Ad esempio, nella previsione della probabilità su PM Netlist, rimuovere FSL ha aumentato il MAE da 0,013 a 0,036 (un aumento relativo dell'errore del 63,9%).
• Accumulo di Errore: L'analisi attraverso i livelli logici ha mostrato che TRACE mantiene prestazioni stabili anche a livelli logici profondi, indicando robustezza contro l'accumulo di errori durante il processo di inferenza iterativa.

5. Significato

Il paper sostiene che TRACE stabilisca un nuovo stato dell'arte per l'analisi dei circuiti e, più ampiamente, fornisca un paradigma più robusto e architettonicamente solido per l'apprendimento su grafi di circuiti. Affrontando il disallineamento fondamentale tra le architetture standard di apprendimento su grafi e la natura del computo (specificamente la sensibilità alla posizione e la gerarchia), TRACE dimostra che allineare il bias induttivo del modello con la struttura del problema è critico per catturare il comportamento funzionale. Il lavoro serve come prova di principio che il decoupling delle approssimazioni locali dagli shift contestuali globali porta a una superiore generalizzazione attraverso diversi domini dell'automazione elettronica (EDA).