random projection in influence functions by haochend413 · Pull Request #231 · TRAIS-Lab/dattri

haochend413 · 2025-12-31T01:35:33Z

Description

Add random projection for some IF attributors

IF_Explicit
IF_EKFAC
IF_DataInf

Closes #227

haochend413 · 2026-01-01T08:27:27Z

I ran the IF attributors on the influence_functions_lds example after adding random projection.

Explicit on mnist_lr: ? -> 0.4851
EKFAC on mnist_mlp: 0.1151 -> 0.1140
DataInf on mnist_lr: 0.3173->0.2978

sleepymalc · 2026-01-09T08:21:27Z

dattri/algorithm/influence_function.py

+        Returns:
+            torch.Tensor: Transformed train representations with projected dimension.
+        """
+        from dattri.func.projection import random_project


This seems redundant:

https://github.com/TRAIS-Lab/dattri/pull/231/changes#diff-e756f81a0441b8319db94cc2a72041edee5cee3936e2150a424ada0afb6803bbR21

TonyZhou05 · 2026-01-25T09:01:56Z

Could we check the attribute function call to make sure that the projection has some effect on larger models such as mobilenet_v2? Currently calculating the influence score on any IF attributor is really slow (No progress on A100 GPU for mobilenet_v2 in 10 minutes) I can create this as a separate bug if needed.

Code snippets

train_subset = Subset(full_train_noisy, range(0, 5000))
train_loader = DataLoader(train_subset, batch_size=32, shuffle=False)
val_loader = DataLoader(val_subset, batch_size=32, shuffle=False)
with torch.no_grad():
    # This steps won't finish
    influence_scores = if_attributor.attribute(train_loader, val_loader)

haochend413 · 2026-02-01T04:45:21Z

Could we check the attribute function call to make sure that the projection has some effect on larger models such as mobilenet_v2? Currently calculating the influence score on any IF attributor is really slow (No progress on A100 GPU for mobilenet_v2 in 10 minutes) I can create this as a separate bug if needed.

Code snippets
train_subset = Subset(full_train_noisy, range(0, 5000))
train_loader = DataLoader(train_subset, batch_size=32, shuffle=False)
val_loader = DataLoader(val_subset, batch_size=32, shuffle=False)
with torch.no_grad():
    # This steps won't finish
    influence_scores = if_attributor.attribute(train_loader, val_loader)

Hi! Would you mind sharing your code / repo so that I can try running your script locally?

TonyZhou05 · 2026-02-02T07:22:09Z

Yes, I'm using the LiSSA Attributor, with the CIFAR10 datasets and mobilev2 model. Here is the link to a running file. You can use these checkpoints and datasets attached. Let me know if you need anything else. Thank you!
notebook:
https://colab.research.google.com/drive/1ccdlWwcp6WhKe2SRIdK8MIRR7MGs_95m?usp=sharing

checkpoints and datasets:
datasets
checkpoints

haochend413 · 2026-02-02T07:27:03Z

Yes, I'm using the LiSSA Attributor, with the CIFAR10 datasets and mobilev2 model. Here is the link to a running file. You can use these checkpoints and datasets attached. Let me know if you need anything else. Thank you! notebook: https://colab.research.google.com/drive/1ccdlWwcp6WhKe2SRIdK8MIRR7MGs_95m?usp=sharing

checkpoints and datasets: datasets checkpoints

Thank you for the information! Though I don't think projection is available for LiSSA since hvp for LiSSA is calculated using the torch vjp directly?

TheaperDeng

Good job. I think we also need unit test to make sure the change works for all three attributors.

TheaperDeng · 2026-01-28T22:33:49Z

dattri/algorithm/influence_function.py

        self,
        task: AttributionTask,
        layer_name: Optional[Union[str, List[str]]] = None,
+        projector_kwargs: Optional[Dict[str, Any]] = None,


projector_kwargs may not be a very good API for external users since they may need to check many document or even unit tests to learn the structure. We may flatten some key parameters (like proj_dim and proj_seed) here in the init

same for projector_kwargs in other attributors.

Proposing either using a pydantic class to wrap around the kwargs to keep it transparent or flattening works too.

Something like this, so that when user wants to configure their projector arguments, it'll have autocomplete:

class BaseProjectorConfig(BaseModel): config_1: str config_2: int class IFAttributorCGProjectorConfig(BaseProjectorConfig): config_3: str class IFAttributorLiSSA(BaseProjectorConfig): config_4: float # IFAttributorCG will have config_1, config_2 and config_3.

Easier to maintain as well

This sounds great to me

TheaperDeng · 2026-02-03T06:46:46Z

dattri/algorithm/influence_function.py

+                    sample_features,
+                    feature_batch_size=1,
+                    **self.projector_kwargs,
+                )


There is no need to change this code block, but I think the random_project may need some change to avoid such kind of cumbersome projector creation. TODO: make this an issue

TheaperDeng · 2026-02-03T07:07:38Z

dattri/algorithm/influence_function.py

+                        blksz_out,
+                        feature_batch_size=1,
+                        **self.projector_kwargs,
+                    )


Again, the cumbersome happens many times. @haochend413 Do you have any suggestion to the API design of the random_project?

I think we can definitely improve attributors' APIs, flattening the projection related inputs for straight-forwardness and adding default values. Since we're currently using null pointer to indicate projection usage, I think we probably should also add a use_projection_or_not bool indicator.

I'm not entirely sure which part of random_project API is potentially redundant? It's using a sample feature to infer sizes since it supports per-layer projection where sizes can vary between layers, but everything else looks fine to me?

def random_project( feature: Union[Dict[str, Tensor], Tensor], feature_batch_size: int, proj_dim: int, proj_max_batch_size: int, proj_seed: int = 0, proj_type: str = "normal", *, device: Union[str, torch.device] = "cpu", ) -> Callable:

TheaperDeng · 2026-02-03T07:15:24Z

dattri/func/fisher.py

+        input_projectors (Optional[Dict[str, Callable]]): A dict of projector functions
+            for projecting input activations. Keys are layer names.
+        output_projectors (Optional[Dict[str, Callable]]): A dict of projector functions
+            for projecting output gradients. Keys are layer names.


Is it possible that we have only one of input_projectors and output_projectors? If not, we may add a check here to report the error.

TheaperDeng · 2026-02-03T07:16:56Z

dattri/func/hessian.py

        A function that takes a tuple of Tensor `x` and a vector `v` and returns
        the IHVP of the Hessian of `func` and `v`.
    """
+    from dattri.func.projection import random_project


Is it possible to put this import to the header of the file?

It will cause circular import since projection module is using hvp_at_x function. I think by original design hessian module is at lower level ? Maybe a solution is to define a higher level module combining projection and hessian, instead of adding projection to hessian directly.

TheaperDeng · 2026-02-03T07:22:06Z

dattri/func/hessian.py

+            projector = random_project(
+                sample_features,
+                1,
+                **projector_kwargs,


Should we control the seed of the projection here? If not, what will happen if we have different projection matrix among H and g?

Seems like this is also a bug here that random_project() call is missing a required parameter proj_max_batch_size

This is a required setting for random projection, you may set (normally to 32) when you give the projection_kwargs to the attributor. We may set some default value for this key as well.

I looked into trak's projection_kwargs handling. It allows partially defined projection_kwargs by defining a static default param dict and uses update() method to merge new params provided by the user. We can do it this way if we want to keep projection_kwargs. However for simplicity I think it's better to flatten the attributor inputs.

# TRAKAttributor(...) DEFAULT_PROJECTOR_KWARGS = { "proj_dim": 512, "proj_max_batch_size": 32, "proj_seed": 0, "device": "cpu", } # ... self.projector_kwargs = DEFAULT_PROJECTOR_KWARGS if projector_kwargs is not None: self.projector_kwargs.update(projector_kwargs) # Usage in unit tests projector_kwargs = { "device": "cpu", } attributor = TRAKAttributor( task=task, correct_probability_func=m, device=torch.device("cpu"), projector_kwargs=projector_kwargs, )

I rechecked the code. For one single attribution, the seed used in all projectors should be the same, which is specified by the user in projector_kwargs. So I think the seeds used are fixed.

TheaperDeng · 2026-02-03T07:29:08Z

BTW, I wonder what may happen if we directly use layer_name to select the parameters we want to include in the gradient (train/test representations). Will it has a lower correaltion than projection? This argument may also be an alternative for algorithms like CG and LiSSA where projection is difficult to be applied. @TonyZhou05

TonyZhou05 · 2026-02-03T07:34:05Z

BTW, I wonder what may happen if we directly use layer_name to select the parameters we want to include in the gradient (train/test representations). Will it has a lower correaltion than projection? This argument may also be an alternative for algorithms like CG and LiSSA where projection is difficult to be applied. @TonyZhou05

I can try this to see how the numbers look like for this specific model. Should probably experiment on other smaller models as well.

haochend413 · 2026-02-03T18:59:07Z

dattri/func/fisher.py

        # The t here is the sequence length or time steps for sequential input
        # t = 1 if the given input is not sequential
        if a_prev_raw.ndim == 2:  # noqa: PLR2004
            a_prev = a_prev_raw.unsqueeze(1)


There is a potential bug here and line 364, where a_prev is not defined on all logic paths. It could cause errors when I run some examples. Probably we should make another PR for this since it's the original code?

haochend413 · 2026-02-04T20:00:03Z

I added a simple unit test for projection. I noticed that EK-FAC will not pass the self-attribute test after projection with 0.03 max diff. I'll look into whether it's the projection's problem (maybe it's related to seed control?).

haochend413 and others added 2 commits December 30, 2025 19:33

random projection in IF-EKFAC

8792296

Merge branch 'main' into IF_rand_proj

93ee6fe

haochend413 mentioned this pull request Dec 31, 2025

Influence Function with Projection #227

Open

add projector for DataInf

77dd8ca

haochend413 added 6 commits January 1, 2026 02:52

function argument

7013d1c

add random projection for IF_Explicit

f61ff88

ruff & darglint

a75aabd

ruff

2c6d212

minor fix

fc5e1d7

make projection optional

7356423

sleepymalc self-requested a review January 9, 2026 08:18

sleepymalc reviewed Jan 9, 2026

View reviewed changes

haochend413 added 2 commits January 9, 2026 11:49

minor fix

97cfc34

ruff

abdca1f

TheaperDeng reviewed Feb 3, 2026

View reviewed changes

haochend413 commented Feb 3, 2026

View reviewed changes

add simple unit test

1155e39

ruff

10a8e0a

Conversation

haochend413 commented Dec 31, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Uh oh!

haochend413 commented Jan 1, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

TonyZhou05 commented Jan 25, 2026

Uh oh!

haochend413 commented Feb 1, 2026

Uh oh!

TonyZhou05 commented Feb 2, 2026

Uh oh!

haochend413 commented Feb 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

TheaperDeng 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!

TonyZhou05 Feb 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

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!

haochend413 Feb 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

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!

haochend413 Feb 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

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!

haochend413 Feb 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

TheaperDeng commented Feb 3, 2026

Uh oh!

TonyZhou05 commented Feb 3, 2026

Uh oh!

haochend413 Feb 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

haochend413 commented Feb 4, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

haochend413 commented Dec 31, 2025 •

edited

Loading

haochend413 commented Jan 1, 2026 •

edited

Loading

haochend413 commented Feb 2, 2026 •

edited

Loading

TonyZhou05 Feb 3, 2026 •

edited

Loading

haochend413 Feb 3, 2026 •

edited

Loading

haochend413 Feb 3, 2026 •

edited

Loading

haochend413 Feb 3, 2026 •

edited

Loading

haochend413 Feb 3, 2026 •

edited

Loading