feat: Switch Class Detection and Subclass Detection Metrics to Recall at K formulation

config/awa2_resnet50.yaml

@@ -36,6 +36,7 @@ log_dir: hydra_logs
 repo_id: "quanda-bench-test"
 cache_dir: "tmp"
 top_k: 5
+detection_s: 5
 model_id: "0"
 correlation_fn: "spearman"
 offline: true

config/bert_qnli.yaml

@@ -23,6 +23,7 @@ log_dir: hydra_logs
 repo_id: "quanda-bench-test"
 cache_dir: "tmp"
 top_k: 5
+detection_s: 5
 model_id: "0"
 correlation_fn: "spearman"
 offline: true

config/cifar_resnet9.yaml

@@ -38,6 +38,7 @@ log_dir: hydra_logs
 repo_id: "quanda-bench-test"
 cache_dir: "tmp"
 top_k: 5
+detection_s: 5
 model_id: "0"
 correlation_fn: "spearman"
 offline: true

config/eval/explainer/arnoldi.yaml

@@ -7,6 +7,5 @@ kwargs:
   projection_dim: 50
   arnoldi_dim: 100
   batch_size: 256
-  precompute_data_ratio: 1.0
   layers:
     - fc_3

config/mnist_lenet.yaml

@@ -36,6 +36,7 @@ log_dir: hydra_logs
 repo_id: "quanda-bench-test"
 cache_dir: "tmp"
 top_k: 5
+detection_s: 5
 model_id: "0"
 correlation_fn: "spearman"
 offline: true

quanda/benchmarks/downstream_eval/class_detection.py

@@ -37,10 +37,12 @@ class ClassDetection(Benchmark):
     eval_args = ["test_data", "test_targets", "explanations"]
     default_use_predictions: bool = True
     default_filter_by_prediction: bool = False
+    default_s: int = 1
 
     def __init__(
         self,
         *args,
+        s: int = 1,
         filter_by_prediction: bool = False,
         **kwargs,
     ):
@@ -50,6 +52,9 @@ def __init__(
         ----------
         *args
             Positional arguments passed to the base class.
+        s : int, optional
+            Number of top-attributed training points to consider in the
+            same-class fraction computation, by default 1.
         filter_by_prediction : bool, optional
             Whether to filter the test samples to only calculate the metric on
             those samples, where the correct class is predicted, by
@@ -59,6 +64,7 @@ def __init__(
 
         """
         super().__init__(*args, **kwargs)
+        self.s = s
         self.filter_by_prediction = filter_by_prediction
 
     @classmethod
@@ -72,6 +78,7 @@ def _extra_kwargs_from_config(
     ) -> dict:
         """Extract class detection kwargs from config."""
         return {
+            "s": config.get("detection_s", cls.default_s),
             "filter_by_prediction": config.get(
                 "filter_by_prediction", cls.default_filter_by_prediction
             ),
@@ -158,6 +165,7 @@ class to be used for evaluation.
             model=self.model,
             checkpoints=self.checkpoints,
             train_dataset=self.train_dataset,
+            s=self.s,
             checkpoints_load_func=self.checkpoints_load_func,
             filter_by_prediction=self.filter_by_prediction,
             inference_batch_size=inference_batch_size,

quanda/benchmarks/downstream_eval/subclass_detection.py

@@ -39,11 +39,13 @@ class SubclassDetection(Benchmark):
     ]
     default_use_predictions: bool = True
     default_filter_by_prediction: bool = True
+    default_s: int = 1
 
     def __init__(
         self,
         *args,
         class_to_group: Optional[Dict[int, int]] = None,
+        s: int = 1,
         filter_by_prediction: bool = True,
         **kwargs,
     ):
@@ -55,6 +57,9 @@ def __init__(
             Positional arguments passed to the base class.
         class_to_group : Optional[Dict[int, int]]
             Mapping from class index to group index.
+        s : int, optional
+            Number of top-attributed training points to consider in the
+            same-subclass fraction computation, by default 1.
         filter_by_prediction : bool, optional
             Whether to filter the test samples to only calculate the metric on
             those samples, where the correct superclass is predicted, by
@@ -65,6 +70,7 @@ def __init__(
         """
         super().__init__(*args, **kwargs)
         self.class_to_group = class_to_group
+        self.s = s
         self.filter_by_prediction = filter_by_prediction
 
         # Ensure all datasets use the same class_to_group mapping.
@@ -94,6 +100,7 @@ def _extra_kwargs_from_config(
 
         return {
             "class_to_group": train_dataset.class_to_group,
+            "s": config.get("detection_s", cls.default_s),
             "filter_by_prediction": config.get(
                 "filter_by_prediction", cls.default_filter_by_prediction
             ),
@@ -196,6 +203,7 @@ def evaluate(
             model=self.model,
             checkpoints=self.checkpoints,
             train_dataset=self.train_dataset,
+            s=self.s,
             checkpoints_load_func=self.checkpoints_load_func,
             train_subclass_labels=train_subclass_labels,
             filter_by_prediction=self.filter_by_prediction,

quanda/benchmarks/resources/configs/5d5968d-awa2_resnet50_ClassDetection.yaml

@@ -107,3 +107,4 @@ splits:
     seed: 42
 cfg_output_dir: quanda/benchmarks/resources/configs
 cfg_file_name: GIT_TAG-awa2_resnet50_ClassDetection.yaml
+detection_s: 5

quanda/benchmarks/resources/configs/5d5968d-awa2_resnet50_SubclassDetection.yaml

@@ -142,3 +142,4 @@ splits:
     seed: 42
 cfg_output_dir: quanda/benchmarks/resources/configs
 cfg_file_name: GIT_TAG-awa2_resnet50_SubclassDetection.yaml
+detection_s: 5

quanda/benchmarks/resources/configs/ad1b983-default_ClassDetection.yaml

@@ -65,3 +65,4 @@ splits:
       test: 0.0
     seed: 42
 batch_size: 64
+detection_s: 5

quanda/benchmarks/resources/configs/ad1b983-default_SubclassDetection.yaml

@@ -101,3 +101,4 @@ splits:
       test: 0.0
     seed: 42
 batch_size: 64
+detection_s: 5

quanda/benchmarks/resources/configs/bdb919e-default_ClassDetection.yaml

@@ -68,3 +68,4 @@ splits:
       test: 0.0
     seed: 42
 batch_size: 64
+detection_s: 5

quanda/benchmarks/resources/configs/bdb919e-default_SubclassDetection.yaml

@@ -110,3 +110,4 @@ splits:
       test: 0.0
     seed: 42
 batch_size: 64
+detection_s: 5

quanda/benchmarks/resources/configs/e59b08c-bert_qnli_ClassDetection.yaml

@@ -96,3 +96,4 @@ splits:
     seed: 42
 cfg_file_name: e59b08c-bert_qnli_ClassDetection.yaml
 cfg_output_dir: quanda/benchmarks/resources/configs
+detection_s: 5

quanda/explainers/wrappers/kronfluence.py

@@ -188,6 +188,7 @@ def __init__(
         self.analyzer.fit_all_factors(
             factors_name=self.factors_name,
             dataset=self.train_dataset,
+            per_device_batch_size=self.batch_size,
             factor_args=self.factor_args,
             overwrite_output_dir=not self.load_from_disk,
         )

quanda/metrics/downstream_eval/class_detection.py

@@ -12,13 +12,12 @@
 class ClassDetectionMetric(Metric):
     """Class Detection Metric.
 
-    Metric that measures the performance of a given data attribution method
-    in detecting the class of a test sample from its highest attributed
-    training point.
-
-    Intuitively, a good attribution method should assign the highest
-    attribution to the class of the test sample, as argued by Hanawa et al.
-    (2021) and Kwon et al. (2024).
+    For each test sample, the fraction of training points that share the
+    test sample's class among the ``s`` most influential training points
+    (top-``s`` highest attribution scores). The metric returns the average
+    of this fraction across test samples. With ``s=1`` this reduces to the
+    top-1 class detection accuracy of Hanawa et al. (2021); the recall-style
+    formulation follows Kwon et al. (2024).
 
     References
     ----------
@@ -36,6 +35,7 @@ def __init__(
         self,
         model: torch.nn.Module,
         train_dataset: Union[torch.utils.data.Dataset, datasets.Dataset],
+        s: int = 5,
         checkpoints: Optional[Union[str, List[str]]] = None,
         checkpoints_load_func: Optional[CheckpointLoadFunc] = None,
         filter_by_prediction: bool = False,
@@ -49,6 +49,9 @@ def __init__(
             The model associated with the attributions to be evaluated.
         train_dataset : Union[torch.utils.data.Dataset, datasets.Dataset]
             The training dataset that was used to train `model`.
+        s : int, optional
+            Number of top-attributed training points to consider when
+            computing the same-class fraction, by default 5.
         checkpoints : Optional[Union[str, List[str]]], optional
             Path to the model checkpoint file(s), defaults to None.
         checkpoints_load_func : Optional[CheckpointLoadFunc], optional
@@ -70,6 +73,7 @@ def __init__(
             checkpoints_load_func=checkpoints_load_func,
         )
 
+        self.s = s
         self.scores: List[torch.Tensor] = []
         self.filter_by_prediction = filter_by_prediction
         self.inference_batch_size = inference_batch_size
@@ -93,11 +97,6 @@ def update(
             Only required if `filter_by_prediction` is True during
             initalization.
 
-        Raises
-        ------
-        AssertionError
-            If the number of explanations does not match the number of labels.
-
         """
         if isinstance(test_targets, list):
             test_targets = torch.tensor(test_targets)
@@ -121,14 +120,18 @@ def update(
 
         explanations = explanations[select_idx]
         test_targets = test_targets[select_idx].to(self.device)
-        _, top_one_xpl_indices = explanations.topk(k=1, dim=1)
-        top_one_xpl_targets = torch.tensor(
+        k = min(self.s, explanations.shape[1])
+        _, top_xpl_indices = explanations.topk(k=k, dim=1)
+        top_xpl_targets = torch.tensor(
             [
                 get_targets(self.train_dataset[int(i)])
-                for i in top_one_xpl_indices
-            ]
-        ).to(self.device)
-        scores = (test_targets == top_one_xpl_targets) * 1.0
+                for i in top_xpl_indices.flatten()
+            ],
+            device=self.device,
+        ).reshape(top_xpl_indices.shape)
+        scores = (
+            (top_xpl_targets == test_targets.unsqueeze(1)).float().mean(dim=1)
+        )
         self.scores.append(scores)
 
     def compute(self):
@@ -143,7 +146,7 @@ def compute(self):
         return {"score": torch.cat(self.scores).mean().item()}
 
     def _per_sample_scores(self) -> Optional[torch.Tensor]:
-        """Return per-sample correctness scores."""
+        """Return per-sample same-class fractions among the top-s."""
         return torch.cat(self.scores) if self.scores else torch.empty(0)
 
     def reset(self, *args, **kwargs):

quanda/metrics/downstream_eval/subclass_detection.py

@@ -9,11 +9,14 @@
 
 
 class SubclassDetectionMetric(ClassDetectionMetric):
-    """Subclass Detection Metric as defined in Hanawa et al. (2021).
+    """Subclass Detection Metric.
 
     A model is trained on a dataset where labels are grouped into superclasses.
-    The metric evaluates the performance of an attribution method in detecting
-    the subclass of a test sample from its highest attributed training point.
+    For each test sample, the fraction of training points that share the test
+    sample's (ungrouped) subclass among the ``s`` most influential training
+    points (top-``s`` highest attribution scores). The metric returns the
+    average of this fraction across test samples. With ``s=1`` this reduces to
+    the original subclass detection accuracy of Hanawa et al. (2021).
 
     References
     ----------
@@ -28,6 +31,7 @@ def __init__(
         model: torch.nn.Module,
         train_dataset: torch.utils.data.Dataset,
         train_subclass_labels: torch.Tensor,
+        s: int = 5,
         checkpoints: Optional[Union[str, List[str]]] = None,
         checkpoints_load_func: Optional[CheckpointLoadFunc] = None,
         filter_by_prediction: bool = False,
@@ -43,6 +47,9 @@ def __init__(
             The training dataset that was used to train `model`.
         train_subclass_labels : torch.Tensor
             The subclass labels of the training dataset.
+        s : int, optional
+            Number of top-attributed training points to consider when
+            computing the same-subclass fraction, by default 5.
         checkpoints : Optional[Union[str, List[str]]], optional
             Path to the model checkpoint file(s), defaults to None.
         checkpoints_load_func : Optional[CheckpointLoadFunc], optional
@@ -61,6 +68,7 @@ def __init__(
             model=model,
             checkpoints=checkpoints,
             train_dataset=train_dataset,
+            s=s,
             checkpoints_load_func=checkpoints_load_func,
             inference_batch_size=inference_batch_size,
         )
@@ -134,10 +142,10 @@ def update(
         explanations = explanations[select_idx]
         test_targets = test_targets[select_idx].to(self.device)
 
-        top_one_xpl_indices = explanations.argmax(dim=1)
-        top_one_xpl_targets = torch.stack(
-            [self.subclass_labels[int(i)] for i in top_one_xpl_indices]
-        ).to(self.device)
-
-        score = (test_targets == top_one_xpl_targets) * 1.0
-        self.scores.append(score)
+        k = min(self.s, explanations.shape[1])
+        _, top_xpl_indices = explanations.topk(k=k, dim=1)
+        top_xpl_targets = self.subclass_labels.to(self.device)[top_xpl_indices]
+        scores = (
+            (top_xpl_targets == test_targets.unsqueeze(1)).float().mean(dim=1)
+        )
+        self.scores.append(scores)

scripts/awa2_resnet50_bench/awa2_plot_config.json

@@ -1,11 +1,11 @@
 {
   "methods": [
     "representer_points",
-    "arnoldi",
     "tracincpfast",
     "trak",
     "random",
-    "similarity"
+    "similarity",
+    "kronfluence"
   ],
   "benches": [
     "awa2_class_detection",

scripts/awa2_resnet50_bench/eval_defs.sh

@@ -5,7 +5,8 @@ declare -A EXPL_SWEEP
 
 EXPL_SWEEP[similarity]="explainer.kwargs.layers=flatten explainer.kwargs.batch_size=128 device=cuda:0 hydra.launcher.n_jobs=1 batch_size=128"
 EXPL_SWEEP[representer_points]="explainer.kwargs.features_layer=flatten explainer.kwargs.classifier_layer=fc explainer.kwargs.batch_size=128 device=cuda:0 explainer.kwargs.normalize=true,false hydra.launcher.n_jobs=1 batch_size=128"
-EXPL_SWEEP[tracincpfast]="explainer.kwargs.batch_size=256 batch_size=256 device=0"
-EXPL_SWEEP[arnoldi]="explainer.kwargs.layers=[fc] explainer.kwargs.projection_dim=50 explainer.kwargs.arnoldi_dim=100 explainer.kwargs.batch_size=256 explainer.kwargs.precompute_data_ratio=0.1 device=cuda:1 hydra.launcher.n_jobs=1"
+EXPL_SWEEP[tracincpfast]="explainer.kwargs.batch_size=256 batch_size=256 device=cuda:0"
+EXPL_SWEEP[arnoldi]="explainer.kwargs.layers=[fc] explainer.kwargs.projection_dim=50 explainer.kwargs.arnoldi_dim=100 explainer.kwargs.batch_size=256 +explainer.kwargs.precompute_data_ratio=0.1 device=cuda:1 hydra.launcher.n_jobs=1"
 EXPL_SWEEP[trak]="explainer.kwargs.proj_dim=1024,2048 explainer.kwargs.batch_size=32 device=cuda:0 hydra.launcher.n_jobs=1"
-EXPL_SWEEP[random]="device=cuda:0 explainer.kwargs.seed=0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28,29,30,31,32,33,34,35,36,37,38,39,40 hydra.launcher.n_jobs=1 batch_size=128"
+EXPL_SWEEP[random]="device=cuda:0 explainer.kwargs.seed=0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28,29,30,31,32,33,34,35,36,37,38,39,40 hydra.launcher.n_jobs=5 batch_size=128 device=cuda:1"
+EXPL_SWEEP[kronfluence]="explainer.kwargs.task_module._target_=quanda.explainers.wrappers.kronfluence_tasks.ImageClassificationTask explainer.kwargs.task_module.tracked_modules=[layer4.2.conv3,fc] +explainer.kwargs.score_args._target_=kronfluence.arguments.ScoreArguments +explainer.kwargs.score_args.use_measurement_for_self_influence=true explainer.kwargs.batch_size=64 device=cuda:0 batch_size=1000 inference_batch_size=64 hydra.launcher.n_jobs=1"

scripts/bert_qnli_bench/eval_defs.sh

@@ -5,7 +5,6 @@ declare -A EXPL_SWEEP
 EXPL_SWEEP[similarity]="explainer.kwargs.layers=dropout explainer.kwargs.similarity_metric.path=quanda.utils.functions.cosine_similarity,quanda.utils.functions.dot_product_similarity +explainer.kwargs.task=text_classification device=cuda:0 batch_size=32"
 EXPL_SWEEP[trak]="explainer.kwargs.proj_dim=2048 explainer.kwargs.lambda_reg=1e-5 +explainer.kwargs.task=text_classification device=cuda:0 explainer.kwargs.batch_size=8"
 EXPL_SWEEP[dattri_tracin]="+explainer.kwargs.task=text_classification +explainer.kwargs.layer_name=[classifier.weight,classifier.bias] device=cuda:0 explainer.kwargs.batch_size=8 batch_size=1000"
-EXPL_SWEEP[random]="device=cuda:0 explainer.kwargs.seed=0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28,29,30,31,32,33,34,35,36,37,38,39,40  hydra.launcher.n_jobs=5"
+EXPL_SWEEP[random]="device=cuda:1 explainer.kwargs.seed=0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28,29,30,31,32,33,34,35,36,37,38,39,40  hydra.launcher.n_jobs=5"
 EXPL_SWEEP[kronfluence]="explainer.kwargs.task_module.tracked_modules=[bert.pooler.dense,classifier] +explainer.kwargs.score_args._target_=kronfluence.arguments.ScoreArguments +explainer.kwargs.score_args.use_measurement_for_self_influence=true +explainer.kwargs.task=text_classification device=cuda:0 batch_size=1000 inference_batch_size=32"
-EXPL_SWEEP[representer_points]="explainer.kwargs.features_layer=dropout explainer.kwargs.classifier_layer=classifier explainer.kwargs.normalize=true +explainer.kwargs.task=text_classification explainer.kwargs.batch_size=32 device=cuda:0"
-EXPL_SWEEP[dattri_if_datainf]="explainer.kwargs.task=text_classification explainer.kwargs.loss_func.path=quanda.explainers.wrappers.dattri_losses.bert_classification_batched_loss explainer.kwargs.hf_input_keys=[input_ids,token_type_ids,attention_mask] +explainer.kwargs.layer_name=[classifier.weight,classifier.bias] explainer.kwargs.batch_size=8 device=cuda:0 batch_size=1000"
+EXPL_SWEEP[representer_points]="explainer.kwargs.features_layer=dropout explainer.kwargs.classifier_layer=classifier explainer.kwargs.normalize=true +explainer.kwargs.task=text_classification explainer.kwargs.batch_size=32 device=cuda:0"

scripts/eval.sh

@@ -25,12 +25,12 @@ run_eval() {
 }
 
 # Populate the local cache (metadata + ckpt) once per benchmark 
-for bench in "${benchmarks[@]}"; do
-    python scripts/prefetch_bench.py \
-        --config-name "$EVAL_CONFIG_NAME" \
-        bench="$bench" \
-        >> "${LOG_DIR}/caching.log" 2>&1
-done
+#for bench in "${benchmarks[@]}"; do
+    #python scripts/prefetch_bench.py \
+        #--config-name "$EVAL_CONFIG_NAME" \
+        #bench="$bench" \
+        #>> "${LOG_DIR}/caching.log" 2>&1
+#done
 
 for bench in "${benchmarks[@]}"; do
     for method in "${methods[@]}"; do