Scalable Attribution and Control of Model Behavior During Training

Sleem Abdelghafar

cs.LGSep 27, 2026

Scalable Attribution and Control of Model Behavior During Training

Authors: Sleem Abdelghafar

Organizations: Rice University

Abstract

Attributing and controlling model behavior during training requires identifying each example's contribution quickly enough to act before the next update. However, examples in the same training batch can produce similar behavioral changes, making their individual contributions difficult to distinguish. We address this ambiguity through mutual information, accounting for interference within the batch by quantifying how much the combined behavioral change reveals about each example's contribution. We show that this mutual information is a logarithmic function of Behavioral Gradient Uniqueness (BGU). BGU gives the information measure its geometric interpretation. Our Batch-Space Ghost (BS-Ghost) algorithm makes these scores practical inside the training loop through shared computation in batch space, without storing model-sized example gradients. On a complete 1,000-example Qwen2.5-7B-Instruct workload, our BS-Ghost implementation adds 27 seconds (8.0%) to 5.5 minutes of ordinary training. Removal and retraining demonstrate that BGU identifies data that causally shapes final behavior. At each training step, signed information identifies which examples strengthen or weaken the target behavior, explaining how behavior develops during training. Signed information also enables cheap intervention during training: it predicts how changing example weights will affect behavior in the next update. We then use these predictions to choose weights that steer behavior toward a desired target. This makes our framework a practical foundation for scalable oversight and verification of training pipelines and processes, helping evaluators assess model alignment, understand how it develops during training, and guide interventions that shape ongoing learning.

Figures & tables

Panel of Figure 1 — Figure 1: Similar behavioral effects can make individual contributions hard to distinguish. (A) Optimizer-aware response directions from a measured batch; cosine similarities use the full behavioral space (B) We vary example weights and train a decoder to recover which were increased or decreased from answer-loss changes. Lower recovery error on independent test interventions means an example’s contribution is easier to distinguish. (C) Mean Spearman correlations between attribution scores and negative test recovery error. BGU predicts recovery rankings best. Its improvement over our optimizer-aware response shows the value of accounting for batch interference (Appendix ).

Figure 2: Information-based data selection changes final behavior. BGU circles and bands show three-seed means and Student 95% intervals after removal and retraining. Lower is better.

Run or method	Time	Timing Scope
Ordinary training	331.90 s	same 1,000 examples and optimizer
Training with BS-Ghost	358.51 s	every example scored during training
Added by BS-Ghost	26.61 s	8.02% increase; additional training time
Vector Filter ( Kowal et al., 2026 )	57 s	post-training pass
Projection Difference ( Chen et al., 2025 )	142 s	post-training pass
Concept Influence † ( Kowal et al., 2026 )	1,170 s	post-training pass

Table 1: BS-Ghost scores every example during training with 8.02% overhead. Qwen2.5-7B-Instruct, 1,000 examples. External rows give published post-training times;

\dagger

excludes inverse-Hessian preparation (Appendix ).

Method	Execution	ToxicChat	XSTest	JailbreakBench	Post-training scoring (s)
		AUPRC $\uparrow$
LESS Xia et al. ( 2024 )	Published	0.388	0.724	1.000	$\sim$ 5,400
	Reference V100	0.2710	0.7105	1.0000	4,704–4,739
	BS-Ghost post-training	0.2710	0.7105	1.0000	2,032–2,062
Grad-Dot Pruthi et al. ( 2020 )	Published	0.084	0.483	0.999	$\sim$ 1,800
	Reference V100	0.0372	0.1752	0.9089	2,110–2,287

Table 2: BS-Ghost accelerates post-training attribution while preserving reference scores. Published times use A6000 ( Jiao et al., 2025 ) . Ranges cover the three datasets; denoted by

\dagger

Panel of Figure 3 — Figure 3: BGU traces how behavior develops during training. Rows align data order, raw internal persona, signed information, and mean judged persona. Each dot is one occurrence. The signed information curves captures strengthening misaligned contributions (red) and weakening benign contributions (blue). Answer error bars are standard errors over 200 generations.

Panel of Figure 5 — Figure 5: BGU feedback suppresses persona while misaligned training continues. Top: our feedback control system. Bottom left: internal persona. Bottom right: mean judged persona, with standard errors over 200 answers. Shading marks training steps when the feedback control is active.

Appendix figures & tables2 assets

Supplementary material from the paper’s appendix.

Appendix

Dataset	Removed	Evil run 1	Evil run 2	Evil run 3	Coherence mean	$\Delta$ NLL
Opinions	0%	27.15	20.37	21.97	88.24	0
	10%	17.45	18.61	22.46	89.35	.0089
	25%	9.40	12.97	11.34	92.60	.0276
	50%	1.96	4.45	3.99	95.93	.0711
	75%	.85	1.23	2.37	96.06	.1302
	90%	1.47	2.13	1.79	96.00	.2362

Appendix

Table 3: Behavioral change and utility after information-selected removal. Each run contributes 200 generations. Lower evil and higher coherence are better.

\Delta

NLL is filtered minus matched unfiltered negative log-likelihood; lower is better. These runs give Figure ’s means and uncertainty.

Training	Final persona	Boundaries above 30	Exposure above 30
Ordinary training	43.03	555	22,923.5
Response feedback	10.29	154	6,100.1
Signed-information feedback	10.63	160	6,107.1
Matched random	35.67	555	21,404.6

Appendix

Table 4: Both targeted controllers suppress internal behavior. One Early trajectory per arm, with 564 observed boundaries through update 563. Exposure integrates

\max(p_{t}-30,0)

by the trapezoidal rule, in persona-point updates. Random uses the signed-information donor’s weights. These are internal readouts, not judged-answer scores.

Explore similar work

Data Attribution at Scale via Influence Matrix Estimation
Sep 14, 2026Yuxi Chen, Hamza Golubovic, Han Tong +2Influence Function Influence
Carnegie Mellon University 1Department of Statistics & Data Science · Columbia University 2Department of Statistics · Columbia University 3Electrical Engineering Department +2
Interactive Training 2: Auditable Control Plane for Live Model Training
Jul 17, 2026Wentao Zhang, Xuanhe Pan, Han Zhou +2Model Training
University of Waterloo · University of Wisconsin-Madison
Anatomy of Post-Training: Using Interpretability to Characterize Data and Shape the Learning Signal
Jun 10, 2026Leon Bergen, Usha Bhalla, Sidharth Baskaran +14Post-Training Interpretability
Core contributors

Sep 14, 2026, stat.MLJ/K move · Enter open · S save

Data Attribution at Scale via Influence Matrix Estimation

Yuxi Chen, Hamza Golubovic, Han Tong, Arian Maleki+1

Carnegie Mellon University 1Department of Statistics & Data Science · Columbia University 2Department of Statistics · Columbia University 3Electrical Engineering Department+2

Influence Function Influence Interpretability Matrix

Open paper AskPDF arXiv

Data attribution seeks to quantify how individual training examples shape a model's predictions and underpins problems including data valuation, machine unlearning, and model interpretability. Despite having a long line of work, computationally scalable methods often struggle to predict the effect of removing training data in neural networks due to their non-convex nature. To overcome this challenge, metagradient-based methods such as MAGIC (Ilyas and Engstrom, 2025) differentiate each prediction through the entire training run and compute its exact influence with respect to the training data, but require a separate run for every prediction. To reduce this cost, we cast budgeted attribution as estimating a large influence matrix from a small number of measurements. We show that the measurements most appropriate for recovering this matrix differ from those best suited for attribution itself. We then present two algorithms, MAGE and SPELL, suited for reconstruction and attribution respectively, that run on existing metagradient machinery at no extra cost. Empirical studies demonstrate strong performance over existing baselines across training scales and measurement budgets.

Abstract

Explore similar work

Data Attribution at Scale via Influence Matrix Estimation

Interactive Training 2: Auditable Control Plane for Live Model Training

Anatomy of Post-Training: Using Interpretability to Characterize Data and Shape the Learning Signal