Fast Data Attribution for Text-to-Image Models

Thanks for the helpful suggestions and comments. Below are our responses.

Code release. As mentioned in the abstract and in the review, we will release code, data, evaluation script, and models to ensure reproducibility upon publication.

Writing improvement. Thank you for suggesting writing improvement for our method. While other reviewers didn’t raise issues with the writing, we agree that it can be polished, and we are happy to provide any further clarifications during the discussion period.

Symbols in Equation 1 (AbU+). In Eqn 1, $\theta_0$ and $\hat z$ are defined in L138, and $\theta_{-\hat z}$, $\alpha$, and $F$ are defined in L139. $N$ is the size of the training set. $\mathcal L$ is the training loss of text-to-image diffusion model, as defined in Section A.1 in the supplemental. In the revision, we will more clearly refer to Section A.1 and define all the symbols.

"Clear and readable algorithm." For reference, we provide a rigorous algorithm block below, which will be added to the revision.

(We use mathjax in this response. Please refresh the page if it's not properly rendered, thanks!)

Algorithm 1 Data-curation (Sec. 3.1).

For distillation, we curate a teacher set with AbU+. For each synthetic query, we first retrieve the top-$K$ neighbors to narrow the search space, then compute AbU+ scores only on that subset (optionally subsampling to size $M$) to improve efficiency. Finally, we rank all scored neighbors for that query and store the ranks. We normalize the rank to $[0, 1]$.

Input
– $\mathcal D=\{(x_i,c_i)\}_{i=1}^{N}$:   Training set (image–caption pairs)
– $\phi$:   Frozen feature encoder for retrieval
– $\tau(\hat z,z)$:   Attribution scores from slow teacher AbU+
– $\mathcal C$:   Caption set to generate query images
– $K$:   Retrieval cut-off
– $M$:   Subsample size

Algorithm

1. Initialize teacher attribution set $\mathcal U \leftarrow \varnothing$.
2. For each caption $c\in\mathcal C$:
3.     Generate synthetic image $\hat x$ and form query $\hat z=(\hat x,c)$.
4.     Retrieve top-$K$ neighbors $\mathcal D_{\hat z}\subset\mathcal D$ of $\hat z$ using $\phi$.
5.     Uniformly subsample $\tilde{\mathcal D_{\hat z}}\subset\mathcal D_{\hat z}$ with $|\tilde{\mathcal D_{\hat z}}| = M$.
6.     Initialize list $\mathcal S_{\hat z}\leftarrow\varnothing$.
7.     For each $z_i\in\tilde{\mathcal D_{\hat z}}$:
8.         Compute score $s_i=\tau(\hat z,z_i)$ and append $s_i$ to $\mathcal S_{\hat z}$.
9.     Sort $\mathcal S_{\hat z}$ in descending order by $s_i$ to assign ranks $\pi^i_{\hat z}=1,2,\dots,|\mathcal S_{\hat z}|$ to each training point $z\in\tilde{\mathcal D_{\hat z}}$.
10.     For each $z_i\in\tilde{\mathcal D_{\hat z}}$ and its rank $\pi^i_{\hat z}$:
11.         Normalize and update $\pi^i_{\hat z} \leftarrow \dfrac{\pi^i_{\hat z}}{|\mathcal S_{\hat z}|}$.
12.         Append $(\hat z,z_i, \pi^i_{\hat z},\mathcal D_{\hat z})$ to $\mathcal U$.

Output Curated teacher set $\mathcal U$ containing tuples $(\hat z,z_i, \pi^i_{\hat z},\mathcal D_{\hat z})$.

(We introduce a new symbol $\mathcal U$ for clarity, and other symbols follows the same convention in the submission.)

Algorithm 2 Training the fast attribution embedder (Sec. 3.2).

Using the teacher set $\mathcal U$, we finetune features so that the feature similarity between the training and synthesized image predicts the ranks.

Input
– $\mathcal D$:   Training set (from Algorithm 1)
– $\mathcal U$:   Teacher set (from Algorithm 1)
– $\phi$:   Frozen encoder
– $g_\psi$:   Trainable MLP
– $(\alpha,\beta)$:   Trainable logit scale
– $p_n$:   Probability of sampling non-neighbor images (for learning to assign lower rank to unrelated image)
– $\eta$:   Learning rate
– $T$ :   Training iteration

Algorithm

1. Repeat for $T$ iterations:
2.     Sample $(\hat z, z_i, \pi^i_{\hat z},\mathcal D_{\hat z})$ from $\mathcal U$.
3.      Sample non-neighbor With probability $p_n$:
            (a) replace data with a non-neighbor $z_i \leftarrow z \sim \mathcal D /\mathcal D_{\hat z}$
            (b) set the lowest rank to the non-neighbor $\pi^i_{\hat z} \leftarrow 1$
4.     Get features $f_\psi(\hat z) = g_\psi\bigl(\phi(\hat z)\bigr),\quad f_\psi(z_i) = g_\psi\bigl(\phi(z_i)\bigr)$.
5.     Similarity $r_\psi(\hat z, z_i) = \cos\bigl(f_\psi(\hat z), f_\psi(z_i)\bigr)$.
6.     Predicted rank $\sigma_{\alpha,\beta}(r_\psi(\hat z, z_i)) = \operatorname{sigmoid}(\alpha\odot r_\psi(\hat z,z_i) + \beta)$.
7.     BCE Loss $\mathcal L(\psi, \alpha, \beta) = \ell_\text{BCE}\bigl(\pi^i_{\hat z}\ ,\ \sigma_{\alpha,\beta}(r_\psi(\hat z, z_i))\bigr)$.
8.     Update $\psi,\alpha,\beta$ using $\nabla\mathcal L$ with AdamW optimizer with learning rate $\eta$.
9. Store the optimized parameters $\psi^\* \leftarrow \psi$, $\alpha^\* \leftarrow \alpha$, $\beta^\* \leftarrow \beta$.

Output Trained $\psi^\*, \alpha^\*, \beta^\*$; pre-compute and index $f_{\psi^\*}(z_i)$ for all training images $z_i \in \mathcal D$.

(Besides $\mathcal U$, we introduce $p_n$, $T$ for clarity, and other symbols follow the same convention in the submission.)

Thank you for acknowledging that our paper addresses a timely problem, provides a systematic study of design choices, demonstrates a scalable method, and presents clear methodology. We address your comments below.

Comparison to previous works. Thanks for the suggestion! We will discuss them in the revision. There are several differences between their approaches and ours:

• The task is different. They focus on data selection for a specified downstream task, where they assign a single influence score to each training point and select ones with high influence. We focus on training point influence for any given synthesized image. The training point influence is different per query synthesized image.

• The distillation method is different. They take the training point feature as input and regress the single influence score. We use feature similarity between the training and synthesized image to predict influence, where we finetune features using learning-to-rank objectives to encourage similarities to have the correct rank.

• Modality is different. They focus on language models, and we focus on text-to-image models.

"The LTR approach distills the relative ranking." Yes, we agree and have discussed this in the limitation (L330). Following Reviewer PCEG’s suggestion, we conducted a preliminary experiment on regressing the absolute attribution scores (normalized by mean and standard deviation). This regression leads to worse ranking performance than our learning-to-rank method. Below, we report mAP (1000) for comparison (higher is better):

Predicting accurate absolute attribution scores remains challenging, and it will be interesting to explore improvements in this in the future. We will add this discussion in the revision.

Linear Datamodeling Scores (LDS). LDS was introduced by Park et al. [1], with the assumption that attribution methods are additive, where the influence of a group of training points is the sum of the individual training point influences. As noted by AbU [2], while the additive assumption sometimes holds true (e.g., influence functions), it does not hold true for many methods (e.g., AbU, feature similarity).

Meanwhile, we did compare previous methods with the counterfactual forgetting metric, as used by many prior works [1,2,3,4]. We will clarify this in the revision.

[1] Park et al. TRAK: Attributing Model Behavior at Scale.

[2] Wang et al. Data Attribution for Text-to-Image Models by Unlearning Synthesized Images.

[3] Ilyas et al. Datamodels: Predicting Predictions from Training Data.

[4] Georgiev et al. The Journey, Not the Destination: How Data Guides Diffusion Models.

Thank you for acknowledging that our paper proposes a novel idea, offers a practical attribution application, and conducts comprehensive experiments. We address your comments below.

"Any failure case in AbU+ would be learned." Yes, we agree and have discussed this in the limitation (L333). This would be true for all distillation methods. Like all distillation methods, our method can be generally applied to different teacher methods, enabling future improvements in (slow) teachers to be distilled into fast students, using our procedure.

Predicting relative instead of absolute attribution scores. Yes, we agree and have discussed this in the limitation (L330). Following Reviewer PCEG’s suggestion, we conducted a preliminary experiment on regressing the absolute attribution scores (normalized by mean and standard deviation). This regression leads to worse ranking performance than our learning-to-rank method. Below, we report mAP (1000) for comparison (higher is better):

Predicting accurate absolute attribution scores remains challenging, and it will be interesting to explore improvements in this in the future. We will add this discussion in the revision.

Computational cost of the training set construction. Yes, the cost of data curation remains high, and we will discuss this in the limitations. We will clarify that our work focuses on achieving efficient attribution, with a low storage cost, at inference time. Currently, most attribution methods cannot handle user requests in real time, making it difficult to deploy in real-world use cases. In scenarios where an attribution system handles many daily user requests, our method is more favorable, since it provides a cost-efficient and fast solution at inference time.

"Do these metrics correlate well with real attribution quality?" Yes, the counterfactual forgetting metric directly measures the real attribution quality, which is widely adopted by AbU and other recent works [1,2,3,4]. As described in AbU [1], the metric evaluates attribution in the following way:

If an attribution algorithm accurately identifies influential training images for the generated image $\hat{**z**}$, then a model trained without those images would be incapable of generating or representing $\hat{**z**}$.

Meanwhile, as discussed in L227-231, we use cheap rank prediction metrics (e.g., mAP (1000)) to evaluate design choices for ranking models. To reduce computation cost, we only perform expensive counterfactual metrics on the best-ranking models to confirm that our method indeed yields good attribution quality. We will clarify this in the revision. We are happy to consider other metrics suggested by the reviewer.

[1] Wang et al. Data Attribution for Text-to-Image Models by Unlearning Synthesized Images.

[2] Ilyas et al. Datamodels: Predicting Predictions from Training Data.

[3] Park et al. TRAK: Attributing Model Behavior at Scale.

[4] Georgiev et al. The Journey, Not the Destination: How Data Guides Diffusion Models.

Thank you for acknowledging that our paper addresses a critical challenge, proposes an appropriate learning-to-rank framework, offers a thorough evaluation, and achieves low-latency attribution with minimal storage requirements. We address your comments below.

"The approach is inherently limited by the teacher method's quality." Yes, we agree and have discussed this in the limitation (L333). This would be true for all distillation methods. Like all distillation methods, our method can be generally applied to different teacher methods, enabling future improvements in (slow) teachers to be distilled into fast students, using our procedure.

"Ensembling multiple different teacher (attribution) models may achieve better results?" Thank you for your suggestion. We also believe that if an ensemble leads to better attribution rankings, we expect our method to perform better by distilling from those rankings. We plan to explore it in future work.

Predicting relative instead of absolute attribution scores. Yes, we agree and have discussed this in the limitation (L330). Following your suggestion, we conducted a preliminary experiment on regressing the absolute attribution scores (normalized by mean and standard deviation). This regression leads to worse ranking performance than our learning-to-rank method. Below, we report mAP (1000) for comparison (higher is better):

Predicting accurate absolute attribution scores remains challenging, and it will be interesting to explore improvements in this in the future. We will add this discussion in the revision.

Effectiveness of K-NN retrieval using off-the-shelf features. Following the reviewer’s suggestion, we study the effect of applying K-NN for data collection. As in Table 1 in the main text, we report counterfactual metrics for AbU+ ran on (1) the full training set (AbU+) and (2) top neighbors after K-NN (AbU+ (K-NN)). We also copied numbers from D-TRAK (2nd best teacher) and random baselines as reference:

We find that applying K-NN retrieval does not introduce a significant performance drop. We are happy to include this study in the revision.

Access to the pretrained T2I model. Yes, our method assumes full access to the model, dataset, and training loss. This setup directly applies to a company running attribution on its in-house model. We also note that recent works on data attribution all follow the same setup [1,2,3].

[1] Georgiev et al. The Journey, Not the Destination: How Data Guides Diffusion Models.

[2] Zheng et al. Intriguing Properties of Data Attribution on Diffusion Models.

[3] Wang et al. Data Attribution for Text-to-Image Models by Unlearning Synthesized Images.

Computational cost of the data curation phase. Yes, the cost of data curation remains high, and we will discuss this in the limitations. We will clarify that our work focuses on achieving efficient attribution, with a low storage cost, at inference time. Currently, most attribution methods cannot handle user requests in real time, making it difficult to deploy in real-world use cases. In scenarios where an attribution system handles many daily user requests, our method is more favorable, since it provides a cost-efficient and fast solution at inference time.

Hello, a gentle reminder for discussions. Authors have provided an extensive rebuttal, and would be great if you can acknowledge it and comment on it!