🤖 AI

TRACE: Learning to Compute on Circuit Graphs

The paper introduces TRACE, a new paradigm for learning circuit graph functionality that employs a Hierarchical Transformer to capture computation flow and a function shift learning objective to decouple global prediction, significantly outperforming existing architectures across diverse circuit modalities.

Original authors: Ziyang Zheng, Jiaying Zhu, Jingyi Zhou, Qiang Xu

Published 2026-06-17

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Ziyang Zheng, Jiaying Zhu, Jingyi Zhou, Qiang Xu

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Problem: The "Blind" Calculator

Imagine you are trying to teach a computer to understand how a complex machine (like a digital circuit) works. This machine is made of many small parts (gates) connected together, passing information from one to the next.

For a long time, the best way to teach computers this was using a method called MPNNs (Message Passing Neural Networks). The paper argues that this method is fundamentally broken for this specific job.

The Analogy: The Blindfolded Chef
Think of a circuit as a recipe where ingredients (inputs) are mixed by a chef (an operator, like a MUX or a logic gate).

The Old Way (MPNNs): Imagine the chef is blindfolded. They can smell the ingredients one by one, but they can't see the order. If you hand them "Salt" then "Pepper," they treat it the same as "Pepper" then "Salt." But in a circuit, order matters! A switch (MUX) behaves completely differently depending on which wire is the control and which is the data. The old method gets confused because it treats the ingredients as a messy pile rather than a specific sequence.
The Other Old Way (Transformers): Imagine a chef who sees all the ingredients at once on a table but doesn't understand the steps of the recipe. They miss the fact that Step 2 depends on the result of Step 1. They try to guess the final dish without following the cooking process.

The Solution: TRACE

The authors created a new system called TRACE. It fixes the problems by changing two main things: the "kitchen layout" (Architecture) and the "teaching method" (Learning Objective).

1. The New Kitchen Layout: The "Step-by-Step" Chef

Instead of the blindfolded chef or the confused one, TRACE uses a Hierarchical Transformer.

How it works: It treats the circuit like a sentence written in a specific order (like a math equation: Add(Multiply(A, B), C)).
The Analogy: Imagine the chef now has a clear instruction card. They don't just look at the ingredients; they look at the order they are handed.
- First, they see the "Operator" (the action, like "Mix").
- Then, they see "Input 1" (the first ingredient).
- Then, "Input 2" (the second ingredient).
Why it helps: Because the chef sees the order, they know exactly how to mix them. If the recipe says "Mix Salt into Pepper," they do it differently than "Mix Pepper into Salt." This allows the computer to understand the exact logic of the circuit, respecting the hierarchy (who depends on whom) and the position of every wire.

2. The New Teaching Method: "Function Shift Learning"

Even with a good chef, predicting the final taste of a giant, complex dish is hard. Sometimes, ingredients interact in surprising ways because they were used in multiple steps earlier in the recipe (this is called "reconvergence" in circuits).

The Problem: If you ask the chef, "What is the final taste?" they might get overwhelmed by the complexity.
The TRACE Fix: Instead of asking the chef to predict the entire final taste from scratch, the authors teach them to predict the difference (the "shift") between a simple guess and the real answer.
- Step A (The Simple Guess): The chef makes a quick guess assuming all ingredients are independent (e.g., "If I have 50% salt and 50% pepper, the mix is 50/50"). This is easy and fast.
- Step B (The Real Answer): The real answer is more complex because the salt and pepper might have interacted earlier in the process.
- The Task: The model is trained to find the gap (the "shift") between the simple guess and the real answer.
The Analogy: It's like a student taking a math test. Instead of trying to solve a massive, scary equation from zero, they are taught: "Here is a simple, easy answer. Now, just tell me how much you need to adjust that answer to get it right." This makes the learning much easier and more accurate.

The Results: A Clear Win

The authors tested TRACE on three different types of circuit "languages" (RTL, AIGs, and Netlists) and compared it to all the previous best methods.

The Outcome: TRACE won every time. It was much better at:
1. Finding twins: Identifying two circuits that do the exact same thing (even if they look different).
2. Predicting behavior: Accurately guessing how a circuit will behave under different conditions.
Why it matters: The paper shows that by fixing the "blindness" of the old methods and simplifying the learning task, we can build much smarter tools for understanding electronic circuits.

Summary

TRACE is a new way for computers to learn how circuits work.

It stops treating circuit inputs as a messy pile and starts treating them as an ordered sequence (like a sentence), so it understands the logic flow.
It stops trying to guess the whole complex answer at once and instead learns to predict the small difference between a simple guess and the truth.

This combination makes it the most accurate tool yet for understanding the "functional behavior" of electronic circuits.

Technical Summary: TRACE: Learning to Compute on Circuit Graphs

1. Problem Statement

The fundamental challenge addressed by this work is learning to compute, defined as the ability to accurately model the functional behavior of circuit graphs. Circuit graphs (including Register Transfer Level (RTL) graphs, And-Inverter Graphs (AIGs), and post-mapping netlists) capture the flow of computation through directed nodes representing operations and variables.

The authors argue that the dominant paradigms in graph representation learning—Message Passing Neural Networks (MPNNs) and Graph Transformers—are architecturally mismatched for this specific task due to three critical flaws:

MPNNs Fail to Model Input Interactions: MPNNs rely on permutation-invariant aggregators (e.g., sum, mean) that combine messages from neighbors independently. This prevents the model from capturing position-aware, operator-specific interactions. For instance, a multiplexer (MUX) where $MUX(S, A, B) \neq MUX(A, S, B)$ cannot be correctly modeled because the order of inputs is lost during aggregation.
Vanilla Transformers are Hierarchy-Agnostic: Standard Transformers flatten the graph into a sequence, destroying the explicit hierarchical structure and intermediate dependencies required to model step-by-step computation (e.g., failing to capture the dependency of $f(g(a,b), c)$ on the intermediate result $g(a,b)$ ).
Edge-Masked Transformers Inherit MPNN Limitations: Transformers that use attention masks to respect graph topology often reduce to weighted sums (similar to Graph Attention Networks), inheriting the inability to model precise, non-linear interactions between multiple inputs.

Furthermore, a deeper challenge exists in capturing global functions that emerge from complex topologies, particularly reconvergent dependencies. In circuits, a node's global function may differ significantly from a simple local approximation that assumes input independence. Previous methods that supervise directly on the global function implicitly bundle these complex contextual effects into embeddings, making it difficult for models to distinguish true functional dependencies from spurious correlations.

2. Methodology

The authors propose TRACE (Transformer for Reasoning about Algebraic and Computational Expressions), a new paradigm consisting of two core components: an architecturally sound backbone and a principled learning objective.

2.1 Hierarchical Transformer Backbone

To resolve the architectural flaws of prior models, TRACE employs a Hierarchical Transformer that mirrors the step-by-step flow of computation.

Logic Level Computation: Nodes are processed level-by-level based on their topological logic depth. Primary Inputs (PIs) are at level 0.
Prefix Notation Representation: Inspired by prefix notation, each computation step at a specific logic level is represented as an ordered sequence: [operator, input_1, input_2, ...].
Position-Aware Encoding: A Transformer encoder processes these sequences with positional encodings. This allows the model to capture operator-specific interactions and position-awareness (e.g., distinguishing the control signal from data inputs in a MUX).
Recursive Processing: The embedding of an operator node is updated based on the Transformer's output for the first token of its sequence. This updated embedding then serves as the input for subsequent operators in the next logic level.
Efficiency: Unlike general graphs with long-tailed degree distributions, circuit graphs have small, bounded in-degrees, preventing the quadratic memory overhead typically associated with Transformers.

2.2 Function Shift Learning (FSL)

To address the difficulty of predicting complex global functions directly, TRACE introduces Function Shift Learning, a novel objective that decouples the learning problem.

Global Function ( $y_{global}$ ): The true probability of a node being logic-1, which accounts for all dependencies (including reconvergence). Computing this directly is expensive ( $O(2^k)$ ).
Local Function ( $y_{local}$ ): A simple approximation assuming input independence. It is computationally efficient ( $O(1)$ ) but fails to capture the true behavior in the presence of reconvergence.
Function Shift ( $y_{FSL}$ ): The model is trained to predict only the discrepancy between the global and local functions:
$y_{FSL} = y_{global} - y_{local}$
Training: The model is trained to regress the function shift using an L1 loss. The ground truth shift is pre-computed from training data.
Inference: During inference, the global function is reconstructed iteratively. For each node, the model predicts the shift, which is added to the local function (computed using the estimated global functions of its predecessors) to yield the final global estimate.

3. Key Contributions

Architectural Alignment: The introduction of a Hierarchical Transformer that replaces permutation-invariant aggregation with position-aware, sequence-based processing, faithfully mirroring the computational flow of circuit graphs.
Decoupled Learning Objective: The proposal of Function Shift Learning, which isolates the complex contextual effects of reconvergence (the "shift") from the simple local approximation, simplifying the learning task.
State-of-the-Art Performance: The demonstration that this combined paradigm substantially outperforms existing MPNNs and Graph Transformers across diverse circuit modalities.

4. Experimental Results

The authors evaluated TRACE on three circuit modalities (RTL, AIG, PM Netlists) across two task types: contrastive learning and predictive tasks.

4.1 Contrastive Tasks

In retrieval tasks (identifying functionally equivalent circuits), TRACE consistently outperformed all baselines:

RTL: Achieved 94.45% Recall@1 (vs. ~88% for the next best).
AIG: Achieved 92.68% Recall@1 (vs. ~88% for the next best).
PM Netlist: Achieved 90.81% Recall@1 (vs. ~86% for the next best).
Notably, many Graph Transformer baselines (e.g., GraphGPS, SGFormer) suffered severe performance degradation or Out-Of-Memory (OOM) errors on AIGs and PM Netlists, whereas TRACE remained stable.

4.2 Predictive Tasks

In tasks predicting node-level properties (Logic-1 probability, similarity, transition probability), TRACE achieved the lowest Mean Absolute Error (MAE) and highest $R^2$ scores:

Combinational AIG: MAE of 0.015 for Logic-1 probability (vs. 0.027 for DeepGate4).
Sequential AIG: MAE of 0.009 for Logic-1 probability and 0.005 for transition probability.
PM Netlists: MAE of 0.013 for Logic-1 probability (vs. 0.048 for GraphSAGE).

4.3 Ablation and Analysis

Function Shift Learning: Removing FSL significantly degraded performance. For example, on PM Netlist probability prediction, removing FSL increased MAE from 0.013 to 0.036 (a 63.9% relative error increase).
Error Accumulation: Analysis across logic levels showed that TRACE maintains stable performance even at deep logic levels, indicating robustness against error accumulation during the iterative inference process.

5. Significance

The paper claims that TRACE establishes a new state of the art for circuit analysis and, more broadly, provides a more robust and architecturally sound paradigm for learning on circuit graphs. By addressing the fundamental mismatch between standard graph learning architectures and the nature of computation (specifically position-awareness and hierarchy), TRACE demonstrates that aligning the model's inductive bias with the problem structure is critical for capturing functional behavior. The work serves as a proof of principle that decoupling local approximations from global contextual shifts leads to superior generalization across diverse electronic design automation (EDA) domains.