Source-linked AI summary
Mechanistic Circuit Identification for Controllable Data Generation
Nakyung Lee, Sangwoo Hong, Jungwoo Lee
TL;DR
Existing synthetic-data pipelines rely heavily on heuristic, black-box control and provide limited insight into sample-level learning dynamics. The paper discovers and validates utility-associated circuits, uses them to steer generation, and introduces SAMS for stage-aware scheduling. Circuit-steered data consistently improves downstream performance and calibration over prompt-based baselines, while the authors identify broader architectures and tasks as future scope.
Problem
Heuristic, black-box generation offers limited explanation of how samples affect learning and can reduce multidimensional utility to a single score.
Method
The framework discovers and causally validates circuits associated with AUM, EL2N, and GradAlign, steers them to generate targeted samples, and schedules those datasets across training stages with SAMS.
Results
Circuit-steered data consistently outperforms prompt-based baselines, with SAMS reaching 85.8% peak accuracy and 0.055 ECE on SciQ.
Takeaways & Limitations
The work demonstrates that mechanistic interpretability can function as a controllable interface for utility-targeted data generation rather than only as a descriptive analysis tool.
Takeaways & Limitations
The intervention study uses mean ablation for GradAlign because zero ablation can collapse its gradient-alignment signal and confound causal measurement.
Abstract
from arXiv · showhide
While recent advances in data synthesis aim to curate high-quality datasets, most generation pipelines still rely on heuristic prompt-based control. This black-box paradigm provides limited insight into how individual samples interact with a model's underlying learning dynamics. To bridge this gap, we propose a circuit-grounded framework that connects training-dynamics-based data valuation with mechanistic interpretability (MI). Specifically, we conceptualize data quality along three complementary utility axes, learnability, challenge, and alignment. First, we uncover specialized model-internal circuits that causally govern these utility signals. Then, moving beyond heuristic prompting toward mechanistic control, we leverage these circuits as controllable interfaces, actively steering generation to produce utility-targeted data. Building on this capability, we introduce SAMS (Stage-Aware Mechanistic Scheduling), which schedules circuit-steered data according to the model's evolving optimization needs. Experiments on multiple-choice QA tasks demonstrate that our approach yields precisely controlled data with greater diversity than prompt-based baselines, consistently improving downstream performance and calibration. Ultimately, this work establishes a principled white-box paradigm for interpretable data generation, pioneering the use of MI not just as an analytical tool, but as a practical, controllable interface.
1 Introduction
The paper reframes synthetic-data quality as a multi-axis, training-dynamics-based problem and connects it to causal mechanistic circuits. It uses those circuits to steer data generation and schedule utility-targeted data according to evolving training needs.
- Motivation: Existing synthesis pipelines are largely heuristic and black-box, offering limited explanation of how individual samples affect learning.Surface-level proxies and single scalar scores can obscure the internal mechanisms and multidimensional nature of sample utility.
- Training Utility: Data quality is represented through complementary axes of Learnability, Challenge, and Alignment.These correspond to AUM, EL2N, and GradAlign, respectively.
- Framework: The proposed pipeline discovers and causally validates utility-associated circuits, steers them to generate targeted samples, and dynamically schedules resulting datasets during fine-tuning.SAMS adjusts the mixture over training time to match stage-specific learning needs.
- Research Questions: The framework asks whether specific internal circuits govern training-utility signals and can serve as causal interfaces for synthetic-data generation.The motivation is that training-dynamics metrics describe utility, whereas mechanistic interpretability can reveal how utility is realized internally.
2 Related Works
Related work positions mechanistic interpretability as a way to reverse-engineer internal circuits and data synthesis as an area still dominated by surface-level generation control.
- Mechanistic Interpretability: Mechanistic interpretability reverse-engineers neural networks by locating granular circuits through which information is processed and transformed.Causal intervention methods such as attribution patching, EAP, and EAP-IG are used to identify these pathways.
- Data Synthesis: Data-synthesis research has explored mixture optimization, model-aware selection, diversity, and pruning, while generation remains largely driven by prompts and in-context examples.These approaches generally treat the generator as a black box.
3 Training-Utility-Aware Circuit Discovery
The paper discovers and validates model-internal circuits associated with learnability, challenge, and alignment, then uses their structure to support utility-aware data filtering and controllable generation.
- Circuit Discovery Setup: The method measures learnability, challenge, and alignment with AUM, EL2N, and GradAlign, then ranks samples into contrasting high- and low-utility subsets.This contrastive setup aims to isolate pathways responsible for differences in training utility rather than general task solving.
- Circuit Discovery Setup: EAP-IG estimates edge-level causal influence using noisy sample-embedding counterfactuals and a differentiable maximum logit-margin target.The clean embedding is perturbed with Gaussian noise, preserving sequence length for circuit analysis.
- Circuit Validation: Discovered circuits achieve substantially higher Abs-CPR than random edge subsets across all three utility metrics, with performance generally rising before saturation or decline as K increases.The strongest recovery is reported for GradAlign-High, EL2N-High, and AUM-Low.
- Circuit Validation: Ablation causes substantial score drops and reduces rank correlations below 0.1, versus above 0.89 for random baselines; AUM and EL2N show GapRed above 98%.GradAlign also exhibits large score drops and near-zero correlations, although its gap fluctuates structurally.
- Circuit-based Data Filtering: Circuit-based filtering consistently outperforms random selection and original scalar-score filtering, particularly for GradAlign, where raw filtering severely degrades downstream accuracy.The circuit representation preserves utility-relevant structure not fully captured by scalar scores.
- Circuit Structure and Functional Roles: All six circuits share early- and late-layer backbone structure but diverge in middle layers, where utility-specific routing redistributes computational mass.The circuits are organized into Canonical Anchor, Informative Rerouting, and Target-Aligned Control modules.
- Circuit Structure and Functional Roles: Learnability primarily uses the Canonical Anchor route, driven by AUM-High (+1.03) and EL2N-Low (+1.18) to extract clear label evidence.The pathway is described as a stable default route for highly learnable, low-pressure samples.
- Circuit Structure and Functional Roles: Challenge uses Informative Rerouting, dominated by EL2N-High (+1.81) and suppressed by AUM-High (−1.33), redirecting computation through alternative mid-layer pathways.The route engages deeper MLPs to resolve ambiguity near the decision boundary.
4 Circuit-Steered Data Synthesis
The framework uses discovered utility circuits as causal interfaces for steering generation toward targeted profiles, followed by selection to improve profile consistency. Ablations and comparisons indicate that the axes jointly support coverage, reasoning signals, alignment, targeted utility, and gradient-space diversity.
- Circuit-Steered Data Generation: Circuit-steered generation combines activation addition and targeted attention steering to control utility-related representations and information routing.Activation addition shifts hidden states, while attention steering modulates circuit-relevant token routing.
- Circuit-Steered Data Generation: Steering vectors are computed from mean activation differences between high- and low-utility groups, with λ_c controlling circuit amplification or suppression.
- Circuit-Steered Data Generation: Post-generation selection removes noisy or weakly matched candidates, improving consistency without changing the steering mechanism.
- Utility-Axis Ablation: Removing AUM preserves fluency but collapses domain coverage and gradient-space diversity, while removing EL2N reduces rationale-related performance gains.
- Utility-Axis Ablation: Removing GradAlign lowers solving accuracy and confidence margin, supporting its role in preserving target-aligned utility.
- Generation Fidelity: Circuit steering produces targeted learnability, challenge, and alignment profiles and higher diversity, particularly under G-Vendi.D_challenging has the highest EL2N with negative AUM, while D_aligned has higher GradAlign than its prompt-based counterpart.
5 SAMS: Stage-Aware Mechanistic Scheduling
SAMS schedules circuit-steered data according to the model's changing optimization needs rather than treating utility-specific pools as a flat mixture. Across source-data ratios and benchmarks, circuit-steered training outperforms prompt-based baselines and improves calibration.
- Motivation: SAMS addresses the mismatch between flat data mixtures and phase-dependent fine-tuning needs by scheduling utility-specific pools across training stages.
- Three-Phase Schedule: Warm-up emphasizes D_learnable, Transition increases D_challenging while retaining an anchor, and Challenge emphasizes D_aligned and D_challenging.
- Downstream Evaluation: Evaluation fine-tunes Qwen2.5-0.5B-Instruct on source-and-synthetic mixtures with 50%–70% source-data ratios using Accuracy, ECE, and NLL.
- Three-Phase Schedule: The schedule uses pool ratios of (0.60, 0.15, 0.25), then (0.25, 0.45, 0.30), and finally (0.10, 0.60, 0.30) for learnable, challenging, and aligned data.
- Downstream Evaluation: 85.8% peak accuracy and 0.055 ECE are reported for SAMS on in-domain SciQ, while circuit-steering configurations also outperform prompt-tuned models on ARC-Easy.
6 Conclusion
The paper presents mechanistic interpretability as a controllable interface for data synthesis by converting training-data utility into circuit-level variables. SAMS aligns these utility profiles with evolving optimization needs, while broader architectures and tasks remain future scope.
- Conclusion: Training-data utility is reverse-engineered into controllable circuit-level variables for generating samples with learnability, informative challenge, and target alignment.
- Conclusion: SAMS dynamically aligns circuit-steered datasets with the model's evolving optimization needs.
- Conclusion: Future work will extend localized interventions to broader architectures and tasks.
A Detailed Experimental Setup for Circuit Discovery
Circuit discovery uses utility scores, contrastive high/low buckets, paired clean-corrupted inputs, and EAP-IG attribution adapted to Qwen2.5's grouped-query attention. The extracted graph is analyzed at node and edge granularity across a large computational graph.
- Utility Scores: AUM, EL2N, and GradAlign provide sample-level measures of learnability, challenge, and alignment for circuit discovery.
- Dataset and Model: The setup uses Qwen2.5-1.5B-Instruct with 11,679 SciQ training examples for computationally tractable circuit discovery and validation.
- Contrastive Curation: Contrastive curation selects the top and bottom 15% of samples, creating 1,751-example high- and low-utility buckets for attribution patching.
- Contrastive Curation: Each bucket example pairs a clean chat-template input with a corrupted counterpart formed by injecting Gaussian embedding noise with σ = 0.1.
- EAP-IG Circuit Discovery: EAP-IG is extended to trace edge-level information flow under Qwen2.5's grouped-query attention, with attribution computed at node and edge granularity.
- EAP-IG Circuit Discovery: The full Qwen2.5-1.5B graph contains 365 nodes and 84,715 potential valid edges before Top-K sparsification.
B Detailed Experimental Setup for Circuit Validation
The validation experiments test whether discovered utility circuits causally determine utility scores and apply them to mechanistic data filtering. Circuit-based subsets are compared with full, random, and original-formulation selection baselines.
- Circuit Validation: Top-250 EAP-IG edges for each utility metric are zero-ablated during forward passes to test causal effects.AUM and EL2N use persistent ablation throughout fine-tuning, whereas GradAlign is measured from an ablated forward-backward pass at a fixed checkpoint.
- Circuit Validation: Ablation impact is measured through full-versus-ablated ranking correlation and high/low bucket score-gap collapse.Substantial decreases in both measures indicate that the ablated circuit generates the corresponding utility signal.
- Data Filtering: SciQ samples are ranked by circuit-based selection scores, retaining the top q = 30% for high- and low-utility filtering.The procedure computes circuit scores for each sample relative to the discovered circuits of each metric before selection.
- Data Filtering: Filtering compares mechanistic subsets with the full dataset, random q% subsets, and top-q% subsets selected using original utility formulations.Downstream validation fine-tunes Qwen2.5-1.5B-Instruct for 3 epochs at 2 × 10^-5.
C Detailed Topology of Training-Utility Circuits
The six utility circuits share early and late structural scaffolds but diverge in middle layers, yielding metric-specific computational routes. Attention-interface analysis finds a largely conserved Q/K/V composition dominated by value pathways.
- Shared Topology: All six circuits share early input loading and late pre-output bottlenecks, while mid-to-late layers contain metric-specific attention, MLP, and suppressive pathways.Recurring components include layer-0 value heads, routing through m1 and m2, and a late m23 → logits bottleneck.
- AUM Circuits: AUM-high forms a stable canonical backbone, whereas AUM-low is a weaker, heterogeneous ambiguity-sensitive branch.The low circuit resembles a failure-adjacent route rather than a clean canonical path.
- EL2N Circuits: EL2N-low follows a canonical low-pressure route, while EL2N-high reroutes mass toward alternative mid-layer paths and deeper MLP processing.The EL2N-high circuit is characterized as a rerouting-pressure branch.
- GradAlign Circuits: GradAlign-high introduces late control through m13 and suppresses shortcut-like routes, while GradAlign-low uses a short-path evidence mode with strong suppression near m3 and late outputs.Both circuits retain aspects of the standard ingress and readout scaffold.
- Functional Modules: The structural pathways are abstracted into three overarching functional modules used in the main text.The modules summarize fine-grained routes across utility metrics.
- Attention-Interface Decomposition: Value pathways dominate attention-interface attribution mass across the six circuits, followed by key pathways, with query pathways contributing less.The Q/K/V composition is largely conserved across circuits, indicating a common attention-interface organization.
D.1 Detailed Circuit Steering Implementation
Circuit-steered generation combines activation interventions with sparse, targeted attention-logit steering, then filters candidates by internal compatibility with the intended diagnostic route. The procedure uses common settings and a neutral prompt across utility profiles.
- Attention Steering: Targeted attention steering perturbs pre-softmax logits only in attention heads selected from discovered circuit presets.The intervention complements activation addition by controlling token-level information routing.
- Attention Steering: A structural routing mask prioritizes semantic regions such as the prefix and local context around the current decoding step.A bounded temperature function adjusts selected-head sharpness according to steering scale.
- Attention Steering: Steering remains sparse, localized, and tied to circuit roles rather than broadly disrupting global attention dynamics.Profile-dependent strengths determine the intervention applied to selected heads.
- Internal Compatibility: Candidate selection removes formatting-invalid outputs and rescored valid candidates using teacher-forced likelihood under the same target circuit configurations.The support score compares each route’s candidate likelihood with the unsteered model.
- Internal Compatibility: The selector shifts generated-data distributions toward desired utility profiles while retaining outputs that are surface-valid and internally aligned.Figure 8 presents the resulting comparison of mean difficulty proxies between raw and selected datasets.
- Generation Configurations: Three targeted pools are constructed: Dlearnable, Dchallenging, and Daligned.Circuit presets and steering scales provide the profile-specific intervention targets.
E Generated Dataset Statistics
The generated datasets retain valid MCQA structure while differing in complexity, lexical overlap, and internal proxy distributions. Circuit-steered profiles show lower copying and shifts consistent with their intended utility targets.
- Basic Statistics: 20,000 generation attempts per profile are reduced to the top 10,000 valid samples through internal compatibility selection.Table 10 reports post-selection dataset statistics, including JSON parsing and lexical-overlap measures.
- Basic Statistics: Circuit-steered and prompt-based generations both achieve consistently high JSON parsing rates, while Dchallenging has slightly longer questions and rationales than Dlearnable.Prompt-based baselines produce significantly longer rationales.
- Lexical Diversity: Circuit-steered profiles have lower BLEU-1 and Jaccard similarity to seed questions than prompt-based baselines.Lower overlap indicates less lexical copying and more diverse generation trajectories.
- Internal Proxy Alignment: The support score measures how much a target circuit increases a candidate’s likelihood relative to the unsteered model.Positive support means the target circuit intrinsically prefers the generated completion.
- Internal Proxy Alignment: Dlearnable shifts toward lower difficulty proxies, Dchallenging toward higher proxies, and Daligned toward a balanced moderate profile.A high Difficulty Proxy indicates greater emphasis on challenging reasoning than canonical anchors.
F Generation Fidelity Experiments Details.
The experiments evaluate generated data through complexity, diversity, and quality metrics, then compare controlled mixtures and stage-aware scheduling under standardized downstream fine-tuning conditions. SAMS shifts sampling from learnable to challenging data while keeping aligned data active throughout training.
- Evaluation Metrics: Complexity uses PPL and Margin to measure generated-sample difficulty and answer discriminability.Higher PPL indicates greater structural difficulty and lexical unpredictability, while lower Margin indicates harder, more ambiguous questions.
- Evaluation Metrics: Diversity is measured with Vendi for lexical variety and G-Vendi for differences in how samples perturb model parameters.Vendi uses sentence embeddings, whereas G-Vendi uses proxy-model input-gradient fingerprints.
- Evaluation Metrics: Quality is assessed with Ans-Acc and R-Gain, capturing answerability and the informativeness of provided rationales.Higher R-Gain indicates questions that are well-structured yet genuinely supported by their accompanying rationale.
- Ablation and Data Mixtures: The ablation removes each utility axis by setting its activation and attention steering scales to zero while holding other conditions identical to the Full configuration.This leave-one-out design isolates the contribution of AUM, EL2N, and GradAlign circuit axes to overall data quality.
- Ablation and Data Mixtures: Training batches mix source and generated data at source ratios of 50%, 60%, or 70%, keeping total training budget constant across methods.UNIFORM-MIX samples equally from learnable, challenging, and aligned synthetic pools, while SAMS uses a stage-specific curriculum.
- Stage-aware Mechanistic Scheduling: SAMS divides fine-tuning into warm-up, transition, and challenge stages, retaining aligned data while shifting mass from learnable to challenging pools.The total training steps are divided into three equal intervals unless otherwise specified.
- Downstream Validation: Downstream evaluation uses SciQ as the in-domain benchmark and ARC-Easy as the out-of-domain benchmark, reporting accuracy, ECE, and NLL.All methods, schedules, and source ratios use the same fine-tuning configuration, including six AdamW epochs at a learning rate of 2 × 10^-5.