Improved the code in IF and TRAK to support passing in dict by Suliang-Jin · Pull Request #225 · TRAIS-Lab/dattri

Suliang-Jin · 2025-11-26T23:51:30Z

Description

I have changed some code in IF, TRAK and TracIn so they now support passing in dict. If the user still passes in tuple or list, the behavior doesn't change.
I have also added some experiments on supporting IF and Huggingface transformers.

jiaqima · 2025-11-27T22:41:13Z

@Suliang-Jin given this is a relatively large PR, could you follow the PR template to provide more detailed context about this PR?

jiaqima

It seems that this PR contained changes from another PR #217. If the proposed changes in this PR depend on the other PR, please make this PR after merging the other one.

jiaqima · 2025-11-27T22:43:43Z

README.md


 The following is an example to use `IFAttributorCG` and `AttributionTask` to apply data attribution to a PyTorch model.

+Please reference [here](./docs/guide/README.md) for the guide on how to properly define train/test data for Attributor and loss/target function.


There is no such file for ./docs/guide/README.md?

I think I have created this README in this PR

jiaqima · 2025-11-27T22:45:12Z

.github/workflows/examples_test.yml

        python examples/brittleness/mnist_lr_brittleness.py --method cg --device cpu
        python examples/data_cleaning/influence_function_data_cleaning.py --device cpu --train_size 1000 --val_size 100 --test_size 100 --remove_number 10
        python examples/relatIF/influence_function_comparison.py --no_output
+        sed -i 's/range(1000)/range(100)/g' examples/lds_vs_gt/mnist.py


This change should belong to another PR?

Yes, sorry about it. I will fix it.

Suliang-Jin · 2025-11-28T19:44:55Z

Summary

I'm sorry about the confusion in this PR. Please refer to the PR created on Nov 26.
This PR is (1) fixing the support of dict type dataset passed in the attributor and (2) adding new experiments on the support of Huggingface transformers for Influence Function.

What’s Changed

Under directory algorithm/, I have added data type checking for tuple, list and dict, and the corresponding way of handling different data type in base.py, influence_function.py, trak.py and tracin.py. That is, if a user wants to pass in a dataset for attribution computation, the script will detect if the type of the dataset is tuple, list or dict, and then process it (e.g., by adding a dummy dimension on each data feature in influence_function.py).
Under directory experiments/gpt2_wikitext/, I have added experiment scripts score_IF.py to test the support of using Influence Function on Huggingface Transformers, the dataset passed in for attribution computation is still tuple. On the other hand, score_IF_dict.py and score_TRAK_dict.py are experiments to test if dict type dataset can be properly processed.
At the instruction from @TheaperDeng, I also added instruction document docs/guide/README.md to illustrate how to properly define the dataset/data loader and loss/target function.

Motivation

The original issue is raised from issue #165.
After discussion with @TheaperDeng, I aimed to provide both the support of Huggingface Transformers on Influence Function and the support of handling dict type data in attribution computation. The two issues are intertwined, as wikitext dataset is originally a dict and is commonly used for benchmarking LLMs. However, the two issues need to be tackled differently.

How It Works

On the support of Huggingface Transformers on Influence Function (see in experiments/gpt2_wikitext/), I simply defined special loss function tailored to the behavior of how influence_function.py and trak.py under algorithm/ handle data processing. That is, since influence_function.py explicitly added a dummy dimension to the dataset using unsqueeze(0), I "unsqueezed" the data in the loss function when calculating Influence Function. On the other hand, trak.py doesn't create a dummy dimension, so I didn't "unsqueezed" the data in the loss function when calculating TRAK.
On the other hand, the support of dict type dataset is simply adding an isinstance() check and processing the data in the corresonding way in base.py, influence_function.py, trak.py and tracin.py under algorithm/. I only added isinstance() in the places where the data processing was originally needed: for example, in base.py, I added the type checking when this happens:

train_batch_data = tuple(
    data.to(self.device).unsqueeze(0) for data in train_batch_data_
)

Testing

Under experiments/gpt2_wikitext/, I have tested all score_IF.py, score_IF_dict.py and score_TRAK_dict.py.
To ensure that the scripts still properly support tuple-type dataset, I have rerun trak_dropout_lds.py and influence_function_lds.py under examples/pretrained_benchmark/.

Related Issues

Fixes #165

TheaperDeng · 2025-12-16T21:38:14Z

dattri/algorithm/base.py

+                    )
+                elif isinstance(train_batch_data_, dict):
+                    train_batch_data = {
+                        k: v.unsqueeze(0) for k, v in train_batch_data_.items()


We also assume the value in dictionary to be tensor right?

I think we should put the data to self.device here.

TheaperDeng · 2025-12-16T21:43:51Z

dattri/algorithm/influence_function.py

+                            k: v.to(self.device) for k, v in full_data_.items()
+                        }
+                    else:
+                        raise Exception("We currently only support the train/test data to be tuple, list or dict.")


No need to fix IFAttributor API, I will delete it in another PR.

dattri/algorithm/base.py

TheaperDeng · 2025-12-16T21:57:28Z

dattri/model_util/retrain.py

    """
+    if seed is None:
+        seed = random.getrandbits(64)
+


This is coverred in another PR?

TheaperDeng · 2025-12-16T21:57:45Z

examples/lds_vs_gt/mnist.py

+    # Calculate and print LDS score
+    ##############################
+    lds_score = lds(score, ground_truth)[0]
+    print("lds:", torch.mean(lds_score[~torch.isnan(lds_score)]))


This is coverred in another PR?

experiments/gpt2_wikitext/requirements.txt

TheaperDeng · 2025-12-16T22:21:00Z

docs/guide/README.md

@@ -0,0 +1,120 @@
+# User Guide


The documentation is clear, a high-level summary table at the top would be beneficial. It should list the supported data types and callable types for all methods in the https://github.com/TRAIS-Lab/dattri?tab=readme-ov-file#supported-algorithms.

TheaperDeng · 2025-12-16T22:23:00Z

docs/guide/data_compatibility.md

+    )
+    logp = -outputs.loss
+    return logp - torch.log(1 - torch.exp(logp))
+```


Slightly different requirements should be applied to TRAK (Multi-class Margin) and TracIN (training loss and any target function).

TheaperDeng · 2025-12-16T22:23:34Z

experiments/gpt2_wikitext/score_IF_dict.py

@@ -0,0 +1,686 @@
+#!/usr/bin/env python


We only need one script for IF

TheaperDeng · 2025-12-16T22:23:48Z

experiments/gpt2_wikitext/score_TRAK_dict.py

@@ -0,0 +1,742 @@
+#!/usr/bin/env python


We only need one script for TRAK.

TheaperDeng · 2026-01-20T00:48:50Z

docs/guide/README.md

@@ -0,0 +1,142 @@
+# User Guide


Chage the title to "Data Type Compatibility for Loss and Target Functions". Rename the file to be data_compatibility.md

TheaperDeng · 2026-01-20T00:55:45Z

docs/guide/README.md

+|                                                   |              [EK-FAC](https://arxiv.org/abs/2308.03296)              | ✔️ | ✔️ | ❌ | [Code example](../../examples/brittleness/mnist_lr_brittleness.py) |
+|                                                   |             [RelatIF](https://arxiv.org/pdf/2003.11630)              | ✔️ | ✔️ | ❌ | [Code example](../../examples/brittleness/mnist_lr_brittleness.py) |
+|                                                   |              [LoGra](https://arxiv.org/pdf/2405.13954)               | ✔️ | ✔️ | ❌ | [Code example](../../examples/brittleness/mnist_lr_brittleness.py) |
+|                                                   |              [GraSS](https://arxiv.org/pdf/2505.18976)               | ✔️ | ✔️ | ❌ | [Code example](../../examples/brittleness/mnist_lr_brittleness.py) |


I think for GraSS, LoGra, RelateIF, EK-FAC, we don't have their examples in ../../examples/brittleness/mnist_lr_brittleness.py

TheaperDeng · 2026-01-20T00:57:13Z

experiments/gpt2_wikitext/score_TRAK.py

+        type=str,
+        default="tuple",
+        choices=["tuple", "list", "dict"]
+    )


We don't need to show-off what we have supported in the examples. Just choose the most convenient way and demonstrate it in the script.

TheaperDeng · 2026-01-20T00:57:40Z

experiments/gpt2_wikitext/score_IF.py

+        type=str,
+        default="tuple",
+        choices=["tuple", "list", "dict"]
+    )


TheaperDeng · 2026-01-20T00:58:44Z

experiments/gpt2_wikitext/score_IF.py

+        score = attributor.attribute(train_dataloader, eval_dataloader)
+
+    torch.save(score, "score_IF.pt")
+    logger.info("Attribution scores saved to score_IF.pt")


How does IF perform on GPT-2 + wikitext setting?

sx-liu · 2026-02-11T21:51:49Z

experiments/gpt2_wikitext/score_TRAK.py

+        type=str,
+        default="tuple",
+        help="What data structure to pass the training/test data for data attribution."
+    )


We could simply remove this argument, and assume the input structure is dict, as it is more natural for huggingface datasets.

sx-liu · 2026-02-11T21:53:01Z

experiments/gpt2_wikitext/score_TRAK.py

+    if args.data_structure == "dict":
+        train_dataset = [{k: torch.tensor(v, dtype=torch.long) for k, v in d.items()} for d in train_dataset]
+        eval_dataset = [{k: torch.tensor(v, dtype=torch.long) for k, v in d.items()} for d in eval_dataset]
+


Remove the conditional check

sx-liu · 2026-02-11T21:53:12Z

experiments/gpt2_wikitext/score_TRAK.py

+        train_dataloader = DataLoader(
+            train_dataset,
+            batch_size=args.per_device_train_batch_size,
+            sampler=train_sampler,


sx-liu · 2026-02-11T21:53:57Z

experiments/gpt2_wikitext/score_TRAK.py

+            train_dataset,
+            collate_fn=custom_collate_fn,
+            batch_size=args.per_device_train_batch_size,
+            sampler=train_sampler,


Just simply remove the list/tuple case

sx-liu · 2026-02-11T21:54:17Z

experiments/gpt2_wikitext/score_TRAK.py

+                input_ids,
+                kwargs={"attention_mask": attention_mask, "labels": labels},
+            )
+            return outputs.loss


sx-liu · 2026-02-11T22:01:26Z

Thanks @Suliang-Jin, I have walked through the changes and most of the parts LGTM. One additional thing is that we might also need a few new unit test cases with dict input.

Suliang-Jin · 2026-02-11T22:03:23Z

Thanks @sx-liu! I will make the update in these two days:)

jiaqima reviewed Nov 27, 2025

View reviewed changes

TheaperDeng reviewed Dec 16, 2025

View reviewed changes

Suliang Jin added 4 commits December 27, 2025 18:26

Add example on LDS (refs TRAIS-Lab#168)

ac5bcf1

Explicit train & target func def

b42b338

define a seed in retrain_lds for robustness

ada3d48

support passing in dict, experiments on score_IF

3ae6be8

Suliang-Jin force-pushed the score_IF branch from 9471f17 to 3ae6be8 Compare December 27, 2025 23:30

to_device fix, file from other PR removed

4cbb1d9

jiaqima closed this Jan 19, 2026

jiaqima reopened this Jan 19, 2026

TheaperDeng reviewed Jan 20, 2026

View reviewed changes

fixed linting and included minor changes

1d163b3

sx-liu reviewed Feb 11, 2026

View reviewed changes


		The following is an example to use `IFAttributorCG` and `AttributionTask` to apply data attribution to a PyTorch model.

		Please reference [here](./docs/guide/README.md) for the guide on how to properly define train/test data for Attributor and loss/target function.

Conversation

Suliang-Jin commented Nov 26, 2025

Description

Uh oh!

jiaqima commented Nov 27, 2025

Uh oh!

jiaqima left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Suliang-Jin commented Nov 28, 2025

Summary

What’s Changed

Motivation

How It Works

Testing

Related Issues

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

sx-liu commented Feb 11, 2026

Uh oh!

Suliang-Jin commented Feb 11, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants