Add QB

examples/glue/README.md

@@ -25,8 +25,8 @@ python train.py --dataset_name sst2 \
 To obtain a pairwise influence scores on maximum of 2000 query data points using `ekfac`, run the following command:
 ```bash
 python analyze.py --dataset_name sst2 \
-    --query_batch_size 8 \
-    --train_batch_size 32 \
+    --query_batch_size 100 \
+    --train_batch_size 512 \
     --checkpoint_dir ./checkpoints \
     --factor_strategy ekfac
 ```

examples/glue/analyze.py

@@ -150,7 +150,7 @@ def main():
         dataset=train_dataset,
         per_device_batch_size=None,
         factor_args=factor_args,
-        overwrite_output_dir=True,
+        overwrite_output_dir=False,
         initial_per_device_batch_size_attempt=512,
     )
     analyzer.compute_pairwise_scores(

examples/imagenet/README.md

@@ -24,7 +24,7 @@ larger query batch size.
 python analyze.py --dataset_dir /mfs1/datasets/imagenet_pytorch/ \
     --query_gradient_rank 32 \
     --query_batch_size 500 \
-    --train_batch_size 512 \
+    --train_batch_size 1024 \
     --factor_strategy ekfac
 ```
 On A100 (80GB), it takes less than 4 hours to compute the pairwise scores with query batching (including computing EKFAC factors).
@@ -35,8 +35,7 @@ contains code the compute the correlations between the full rank prediction and
 <p align="center">
 <a href="#"><img width="380" img src="figure/query_batching.png" alt="Counterfactual"/></a>
 </p>
-The averaged correlations between the low-rank and full rank for 100 data points is 0.95.
-
+The averaged correlations between the low-rank and full rank scores for 100 data points is 0.95.
 
 ## Computing Pairwise Influence Scores with DDP