Technische Samenvatting: TRACE: Leren rekenen op circuitgrafen

1. Probleemstelling

De fundamentele uitdaging die in dit werk wordt aangepakt, is leren rekenen (learning to compute), gedefinieerd als het vermogen om het functionele gedrag van circuitgrafen accuraat te modelleren. Circuitgrafen (inclusief Register Transfer Level (RTL)-grafen, And-Inverter Graphs (AIG's) en post-mapping netlijsten) vatten de stroom van berekeningen samen via gerichte knopen die operaties en variabelen vertegenwoordigen.

De auteurs stellen dat de dominante paradigma's voor grafische representatie-leren — Message Passing Neural Networks (MPNN's) en Graph Transformers — architecturaal niet overeenkomen met deze specifieke taak vanwege drie kritieke gebreken:

1. MPNN's falen in het modelleren van input-interacties: MPNN's vertrouwen op permutatie-invariante aggregatoren (bijv. som, gemiddelde) die berichten van buren onafhankelijk combineren. Dit voorkomt dat het model positie-bewuste, operator-specifieke interacties kan vastleggen. Bijvoorbeeld, een multiplexer (MUX) waarbij $MUX(S, A, B) \neq MUX(A, S, B)$ kan niet correct worden gemodelleerd omdat de volgorde van de inputs verloren gaat tijdens de aggregatie.
2. Vanilla Transformers zijn hiërarchie-agnostisch: Standaard Transformers platten de graaf af naar een sequentie, waardoor de expliciete hiërarchische structuur en intermediaire afhankelijkheden die nodig zijn om stap-voor-stap berekeningen te modelleren, worden vernietigd (bijv. het niet kunnen vastleggen van de afhankelijkheid van $f(g(a,b), c)$ van het intermediaire resultaat $g(a,b)$).
3. Edge-gemaskeerde Transformers erven MPNN-beperkingen: Transformers die aandachtsmaskers gebruiken om de grafiektopologie te respecteren, reduceren vaak tot gewogen gemiddelden (vergelijkbaar met Graph Attention Networks), waardoor ze de onmogelijkheid erven om precieze, niet-lineaire interacties tussen meerdere inputs te modelleren.

Verder bestaat er een diepere uitdaging in het vastleggen van globale functies die voortkomen uit complexe topologieën, met name reconvergerende afhankelijkheden. In circuits kan de globale functie van een knoop aanzienlijk verschillen van een eenvoudige lokale benadering die uitgaat van onafhankelijke inputs. Eerdere methoden die direct op de globale functie superviseren, bundelen deze complexe contextuele effecten impliciet in embeddings, wat het voor modellen moeilijk maakt om ware functionele afhankelijkheden te onderscheiden van spuriële correlaties.

2. Methodologie

De auteurs stellen TRACE voor (Transformer for Reasoning about Algebraic and Computational Expressions), een nieuw paradigma bestaande uit twee kerncomponenten: een architecturaal solide backbone en een principieel leerdoel.

2.1 Hiërarchische Transformer Backbone

Om de architecturale gebreken van eerdere modellen op te lossen, gebruikt TRACE een Hiërarchische Transformer die de stap-voor-stap stroom van berekeningen weerspiegelt.

• Logisch niveau berekening: Knopen worden niveau voor niveau verwerkt op basis van hun topologische logische diepte. Primaire Inputs (PI's) bevinden zich op niveau 0.
• Prefix-notatie Representatie: Geïnspireerd door prefix-notatie, wordt elke berekeningsstap op een specif kind niveau vertegenwoordigd als een geordende sequentie: [operator, input_1, input_2, ...].
• Positie-bewuste Encoding: Een Transformer encoder verwerkt deze sequenties met positionele encodings. Dit stelt het model in staat om operator-specifieke interacties en positie-bewustzijn te vangen (bijv. het onderscheiden van het besturingssignaal van datainputs in een MUX).
• Recursieve Verwerking: De embedding van een operator-knoop wordt bijgewerkt op basis van de Transformer-output voor het eerste token van zijn sequentie. Deze bijgewerkte embedding dient vervolgens als de input voor opeenvolgende operators in het volgende logische niveau.
• Efficiëntie: In tegenstelling tot algemene grafen met long-tailed degree distributies, hebben circuitgrafen een kleine, begrensde in-degree, wat de kwadratische geheugenoverhead die gewoonlijk met Transformers gepaard gaat, voorkomt.

2.2 Function Shift Learning (FSL)

Om de moeilijkheid van het direct voorspellen van complexe globale functies aan te pakken, introduceert TRACE Function Shift Learning, een nieuw leerdoel dat het leerprobleem ontkoppelt.

• Globale Functie ($y_{global}$): De ware waarschijnlijkheid dat een knoop logisch '1' is, die alle afhankelijkheden (inclus kind reconvergentie) meeneemt. Het direct berekenen hiervan is duur ($O(2^k)$).
• Lokale Functie ($y_{local}$): Een eenvoudige benadering die uitgaat van onafhanklijke inputs. Dit is computationeel efficiënt ($O(1)$), maar faalt in het vastleggen van het ware gedrag in aanwezigheid van reconvergentie.
• Function Shift ($y_{FSL}$): Het model wordt getraind om alleen de discrepantie tussen de globale en de lokale functie te voorspellen:
  $$y_{FSL} = y_{global} - y_{local}$$
• Training: Het model wordt getraind om de function shift te regresseren met een L1-loss. De grondwaarheid van de shift wordt vooraf berekend vanuit de trainingsdata.
• Inference: Tijdens de inference wordt de globale functie iteratief gereconstrueerd. Voor elke knoop voorspelt het model de shift, die wordt toegevoegd aan de lokale functie (berekend met de geschatte globale functies van de voorgangers) om de uiteindelijke globale schatting te verkrijgen.

3. Belangrijkste Bijdragen

1. Architecturale Afstemming: De introductie van een Hiërarchische Transformer die permutatie-invariante aggregatie vervangt door positie-bewuste, sequentie-gebaseerde verwerking, wat de computationele stroom van circuitgrafen getrouw weerspiegelt.
2. Ontkoppeld Leerdoel: De voorstel van Function Shift Learning, dat de complexe contextuele effecten van reconvergentie (de "shift") isoleert van de eenvoudige lokale benadering, wat de leeropgave vereenvoudigt.
3. State-of-the-Art Prestaties: Het aantonen dat dit gecombineerde paradigma bestaande MPNN's en Graph Transformers aanzienlijk overtreft over diverse circuit-modaliteiten.

4. Experimentele Resultaten

De auteurs evalueerden TRACE op drie circuit-modaliteiten (RTL, AIG, PM Netlists) over twee typen taken: contrastief leren en predictieve taken.

4.1 Contrastieve Taken

In retrieval-taken (het identificeren van functioneel equivalente circuits) presteerde TRACE consistent beter dan alle baselines:

• RTL: Behaalde 94,45% Recall@1 (tegenover ~88% voor de op één na beste).
• AIG: Behaalde 92,68% Recall@1 (tegenover ~88% voor de op één na beste).
• PM Netlist: Behaalde 90,81% Recall@1 (tegenover ~86% voor de op één na beste).
• Opvallend genoeg leden veel Graph Transformer-baselines (bijv. GraphGPS, SGFormer) onder ernstige prestatievermindering of Out-Of-Memory (OOM) fouten op AIG's en PM Netlists, terwijl TRACE stabiel bleef.

4.2 Predictieve Taken

In taken voor het voorspellen van knoop-eigenschappen (Logic-1 waarschijnlijkheid, gelijkenis, transitie-waarschijnlijkheid), behaalde TRACE de laagste Mean Absolute Error (MAE) en de hoogste $R^2$-scores:

• Combinational AIG: MAE van 0,015 voor Logic-1 waarschijnlijkheid (tegenover 0,027 voor DeepGate4).
• Sequential AIG: MAE van 0,009 voor Logic-1 waarschijnlijkheid en 0,005 voor transitie-waarschijnlijkheid.
• PM Netlists: MAE van 0,013 voor Logic-1 waarschijnlijkheid (tegenover 0,048 voor GraphSAGE).

4.3 Ablatie en Analyse

• Function Shift Learning: Het verwijderen van FSL verslechterde de prestaties aanzienlijk. Bijvoorbeeld, bij de PM Netlist waarschijnlijkheidsvoorspelling verhoogde het verwijderen van FSL de MAE van 0,013 naar 0,036 (een relatieve fouttoename van 63,9%).
• Foutaccumulatie: Analyse over logische niveaus toonde aan dat TRACE een stabiele prestatie behoudt, zelfs bij diepe logische niveaus, wat wijst op robuustheid tegen foutaccumulatie tijdens het iteratieve inference-proces.

5. Betekenis

Het artikel beweert dat TRACE een nieuwe state-of-the-art vestigt voor circuitanalyse en, breder genomen, een robuuster en architecturaal sterker paradigma biedt voor het leren op circuitgrafen. Door de fundamentele mismatch tussen standaard grafisch leerarchitecturen en de aard van berekening (met name positie-bewustzijn en hiërarchie) aan te pakken, demonstreert TRACE dat het afstemmen van de inductieve bias van het model met de probleemstructuur cruciaal is voor het vastleggen van functioneel gedrag. Het werk dient als een bewijs van principe dat het ontkoppelen van lokale benaderingen van globale contextuele shifts leidt tot superieure generalisatie over diverse Electronic Design Automation (EDA) domeinen.