Fix/per module accumulate iterations

DOCUMENTATION.md

@@ -136,7 +136,7 @@ You can organize all factors and scores for the specific model with `factors_nam
 
 **What should I do if my model does not have any nn.Linear or nn.Conv2d modules?**
 Currently, the implementation does not support influence computations for modules other than `nn.Linear` or `nn.Conv2d`.
-Try rewriting the model so that it uses supported modules (as done for the `conv1d` module in the [GPT-2 example](https://github.com/pomonam/kronfluence/tree/documentation/examples/wikitext)).
+Try rewriting the model so that it uses supported modules (as done for the `conv1d` module in the [GPT-2 example](https://github.com/pomonam/kronfluence/tree/main/examples/wikitext)).
 Alternatively, you can create a subclass of `TrackedModule` to compute influence scores for your custom module.
 If there are specific modules you would like to see supported, please submit an issue.
 
@@ -150,7 +150,7 @@ inspect `model.named_modules()` to determine what modules to use. You can specif
 
 > [!NOTE]
 > If the embedding layer for transformers are defined with `nn.Linear`, you must write your own
-> `task.tracked_modules` to avoid influence computations embedding matrices.
+> `task.get_influence_tracked_modules` to avoid influence computations on embedding matrices.
 
 **How should I implement Task.compute_train_loss?**
 Implement the loss function used to train the model. Note that the function should return 
@@ -378,6 +378,7 @@ all modules, this will keep track of intermediate module-wise scores.
 - `aggregate_train_gradients`: Whether to use the summed training gradient instead of per-sample training gradients.
 - `use_measurement_for_self_influence`: Whether to use the measurement (instead of the loss) when computing self-influence scores.
 - `query_gradient_low_rank`: The rank for the query batching (low-rank approximation to the preconditioned query gradient; see **Section 3.2.2**). If `None`, no query batching will be used.
+- `use_full_svd`: Whether to use the full SVD (`torch.linalg.svd`) instead of the lower-precision but faster `torch.svd_lowrank` when computing the low-rank query gradient. Only relevant when `query_gradient_low_rank` is set.
 - `query_gradient_svd_dtype`: `dtype` for performing singular value decomposition (SVD) for query batch. You can also use `torch.float64`.
 - `query_gradient_accumulation_steps`: Number of query gradients to accumulate over. For example, when `query_gradient_accumulation_steps=2` with 
 `query_batch_size=16`, a total of 32 query gradients will be stored in memory when computing dot products with training gradients.
@@ -391,7 +392,14 @@ To compute pairwise influence scores (**Equation 5** in the paper), you can run:
 
 ```python
 # Computing pairwise influence scores.
-analyzer.compute_pairwise_scores(scores_name="pairwise", factors_name="ekfac", score_args=score_args)
+analyzer.compute_pairwise_scores(
+    scores_name="pairwise",
+    factors_name="ekfac",
+    query_dataset=query_dataset,
+    train_dataset=train_dataset,
+    per_device_query_batch_size=64,
+    score_args=score_args,
+)
 # Loading pairwise influence scores.
 scores = analyzer.load_pairwise_scores(scores_name="pairwise")
 ```
@@ -400,7 +408,12 @@ To compute self-influence scores (see **Section 5.4** from [this paper](https://
 
 ```python
 # Computing self-influence scores.
-analyzer.compute_self_scores(scores_name="self", factors_name="ekfac", score_args=score_args)
+analyzer.compute_self_scores(
+    scores_name="self",
+    factors_name="ekfac",
+    train_dataset=train_dataset,
+    score_args=score_args,
+)
 # Loading self-influence scores.
 scores = analyzer.load_self_scores(scores_name="self")
 ```

README.md

@@ -66,7 +66,7 @@ from torch import nn
 from kronfluence.analyzer import Analyzer, prepare_model
 
 # Define the model and load the trained model weights.
-model = torch.nn.Sequential(
+model = nn.Sequential(
     nn.Flatten(),
     nn.Linear(784, 1024, bias=True),
     nn.ReLU(),
@@ -76,7 +76,7 @@ model = torch.nn.Sequential(
     nn.ReLU(),
     nn.Linear(1024, 10, bias=True),
 )
-model.load_state_dict(torch.load("model_path.pth"))
+model.load_state_dict(torch.load("model_path.pth", weights_only=True))
 
 # Load the dataset.
 train_dataset = torchvision.datasets.MNIST(
@@ -87,10 +87,10 @@ train_dataset = torchvision.datasets.MNIST(
 eval_dataset = torchvision.datasets.MNIST(
     root="./data",
     download=True,
-    train=True,
+    train=False,
 )
 
-# Define the task. See the Technical Documentation page for details.
+# Define the task by subclassing `kronfluence.task.Task` — see DOCUMENTATION.md for details.
 task = MnistTask()
 
 # Prepare the model for influence computation.
@@ -110,6 +110,7 @@ analyzer.compute_pairwise_scores(
 )
 
 # Load the scores with dimension `len(eval_dataset) x len(train_dataset)`.
+# `scores["all_modules"][i, j]` is the influence of train example `j` on eval example `i`.
 scores = analyzer.load_pairwise_scores(scores_name="my_scores")
 ```
 

kronfluence/computer/computer.py

@@ -273,7 +273,7 @@ def _get_data_partition(
             target_data_partitions = [target_data_partitions]
 
         for data_partition in target_data_partitions:
-            if data_partition < 0 or data_partition > data_partitions:
+            if data_partition < 0 or data_partition >= data_partitions:
                 error_msg = f"Invalid data partition {data_partition}. Must be in range [0, {data_partitions})."
                 self.logger.error(error_msg)
                 raise ValueError(error_msg)
@@ -308,7 +308,7 @@ def _get_module_partition(
             target_module_partitions = [target_module_partitions]
 
         for module_partition in target_module_partitions:
-            if module_partition < 0 or module_partition > module_partitions:
+            if module_partition < 0 or module_partition >= module_partitions:
                 error_msg = f"Invalid module partition {module_partition}. Must be in range [0, {module_partitions})."
                 self.logger.error(error_msg)
                 raise ValueError(error_msg)

kronfluence/factor/eigen.py

@@ -210,6 +210,9 @@ def perform_eigendecomposition(
                         eigenvalues, eigenvectors = torch.linalg.eigh(covariance_matrix)
                     else:
                         raise
+                # Covariance matrices are PSD; clamp tiny negative roundoff to zero so
+                # downstream damping/reciprocal can't divide by a negative value.
+                eigenvalues.clamp_(min=0)
                 del covariance_matrix
                 eigen_factors[eigenvalues_name][module_name] = eigenvalues.contiguous().to(
                     dtype=original_dtype, device="cpu"

kronfluence/module/conv2d.py

@@ -179,6 +179,14 @@ def compute_per_sample_gradient(
     def compute_pairwise_score(
         self, preconditioned_gradient: torch.Tensor, input_activation: torch.Tensor, output_gradient: torch.Tensor
     ) -> torch.Tensor:
+        if self.score_args.compute_per_token_scores:
+            raise UnsupportableModuleError(
+                f"`compute_per_token_scores=True` is not supported for `nn.Conv2d` modules "
+                f"(module name: {self.name}). Per-token scores are only supported for "
+                f"`nn.Linear` modules whose input activation has a sequence (token) axis. "
+                f"Either set `score_args.compute_per_token_scores=False`, or exclude "
+                f"`nn.Conv2d` modules from the tracked modules."
+            )
         input_activation = self._flatten_input_activation(input_activation=input_activation)
         input_activation = input_activation.view(output_gradient.size(0), -1, input_activation.size(-1))
         output_gradient = rearrange(tensor=output_gradient, pattern="b o i1 i2 -> b (i1 i2) o")

kronfluence/score/dot_product.py

@@ -116,7 +116,7 @@ def compute_dot_products_with_loader(
                             raise
                     pairwise_scores = pairwise_scores.cpu()
                     score_chunks[ALL_MODULE_NAME].append(pairwise_scores)
-                    accumulate_iterations(model=model, tracked_module_names=tracked_module_names)
+                accumulate_iterations(model=model, tracked_module_names=tracked_module_names)
 
             if state.use_distributed and total_steps % DISTRIBUTED_SYNC_INTERVAL == 0:
                 state.wait_for_everyone()

kronfluence/score/pairwise.py

@@ -202,7 +202,7 @@ def compute_pairwise_scores_with_loaders(
     query_iter = iter(query_loader)
     num_accumulations = 0
     enable_amp = score_args.amp_dtype is not None
-    enable_grad_scaler = enable_amp and factor_args.amp_dtype == torch.float16
+    enable_grad_scaler = enable_amp and score_args.amp_dtype == torch.float16
     scaler = GradScaler(init_scale=factor_args.amp_scale, enabled=enable_grad_scaler)
     if enable_grad_scaler:
         gradient_scale = 1.0 / scaler.get_scale()
@@ -325,7 +325,7 @@ def compute_pairwise_query_aggregated_scores_with_loaders(
     prepare_modules(model=model, tracked_module_names=tracked_module_names, device=state.device)
 
     enable_amp = score_args.amp_dtype is not None
-    enable_grad_scaler = enable_amp and factor_args.amp_dtype == torch.float16
+    enable_grad_scaler = enable_amp and score_args.amp_dtype == torch.float16
     scaler = GradScaler(init_scale=factor_args.amp_scale, enabled=enable_grad_scaler)
     if enable_grad_scaler:
         gradient_scale = 1.0 / scaler.get_scale()

kronfluence/score/self.py

@@ -201,7 +201,7 @@ def compute_self_scores_with_loaders(
 
     total_steps = 0
     enable_amp = score_args.amp_dtype is not None
-    enable_grad_scaler = enable_amp and factor_args.amp_dtype == torch.float16
+    enable_grad_scaler = enable_amp and score_args.amp_dtype == torch.float16
     scaler = GradScaler(init_scale=factor_args.amp_scale, enabled=enable_grad_scaler)
     if enable_grad_scaler:
         gradient_scale = 1.0 / scaler.get_scale()
@@ -332,7 +332,7 @@ def compute_self_measurement_scores_with_loaders(
 
     total_steps = 0
     enable_amp = score_args.amp_dtype is not None
-    enable_grad_scaler = enable_amp and factor_args.amp_dtype == torch.float16
+    enable_grad_scaler = enable_amp and score_args.amp_dtype == torch.float16
     scaler = GradScaler(init_scale=factor_args.amp_scale, enabled=enable_grad_scaler)
     if enable_grad_scaler:
         gradient_scale = 1.0 / scaler.get_scale()

kronfluence/task.py

@@ -99,7 +99,7 @@ def get_attention_mask(self, batch: Any) -> Optional[Union[Dict[str, torch.Tenso
     def post_process_per_sample_gradient(self, module_name: str, gradient: torch.Tensor) -> torch.Tensor:
         """Post-processes the per-sample gradient of a specific module.
 
-        This method is called only if `do_post_process_per_sample_gradient` is set to `True`.
+        This method is called only if `enable_post_process_per_sample_gradient` is set to `True`.
         Override this method in subclasses to implement custom gradient post-processing.
 
         Args:

tests/scores/test_pairwise_scores.py

@@ -316,6 +316,9 @@ def test_pairwise_scores_partition_equivalence(
         overwrite_output_dir=True,
     )
 
+    if test_name == "conv" and compute_per_token_scores:
+        pytest.skip("per-token scores are not supported for `nn.Conv2d` modules.")
+
     score_args = pytest_score_arguments()
     score_args.compute_per_module_scores = compute_per_module_scores
     score_args.compute_per_token_scores = compute_per_token_scores
@@ -434,7 +437,7 @@ def test_per_module_scores_equivalence(
     assert torch.allclose(total_scores, scores[ALL_MODULE_NAME], atol=ATOL, rtol=RTOL)
 
 
-@pytest.mark.parametrize("test_name", ["mlp", "conv", "gpt"])
+@pytest.mark.parametrize("test_name", ["mlp", "gpt"])
 @pytest.mark.parametrize("compute_per_module_scores", [True, False])
 @pytest.mark.parametrize("query_size", [12])
 @pytest.mark.parametrize("train_size", [64])
@@ -504,6 +507,49 @@ def test_per_token_scores_equivalence(
             assert torch.allclose(per_token_scores[module_name], scores[module_name], atol=ATOL, rtol=RTOL)
 
 
+@pytest.mark.parametrize("query_size", [4])
+@pytest.mark.parametrize("train_size", [8])
+@pytest.mark.parametrize("seed", [5])
+def test_per_token_scores_rejected_for_conv2d(
+    query_size: int,
+    train_size: int,
+    seed: int,
+) -> None:
+    # Conv2d does not support per-token scores; we expect a clear error rather
+    # than silently producing inconsistent shapes across modules.
+    from kronfluence.utils.exceptions import UnsupportableModuleError
+
+    model, train_dataset, test_dataset, data_collator, task = prepare_test(
+        test_name="conv",
+        query_size=query_size,
+        train_size=train_size,
+        seed=seed,
+    )
+    kwargs = DataLoaderKwargs(collate_fn=data_collator)
+    model, analyzer = prepare_model_and_analyzer(model=model, task=task)
+    analyzer.fit_all_factors(
+        factors_name=DEFAULT_FACTORS_NAME,
+        dataset=train_dataset,
+        dataloader_kwargs=kwargs,
+        per_device_batch_size=8,
+        overwrite_output_dir=True,
+    )
+    score_args = pytest_score_arguments()
+    score_args.compute_per_token_scores = True
+    with pytest.raises(UnsupportableModuleError, match="Conv2d"):
+        analyzer.compute_pairwise_scores(
+            scores_name=DEFAULT_SCORES_NAME,
+            factors_name=DEFAULT_FACTORS_NAME,
+            query_dataset=test_dataset,
+            per_device_query_batch_size=4,
+            train_dataset=train_dataset,
+            per_device_train_batch_size=8,
+            dataloader_kwargs=kwargs,
+            score_args=score_args,
+            overwrite_output_dir=True,
+        )
+
+
 @pytest.mark.parametrize(
     "test_name",
     [