BPA: The balanced-pairwise-affinities Feature Transform (ICML 2024)

This repository contains the official PyTorch implementation of BPA (formerly SOT) — the Balanced-Pairwise-Affinities feature transform — as described in our paper The Balanced-Pairwise-Affinities Feature Transform, presented at ICML 2024.

BPA enhances the representation of a set of input features to support downstream tasks such as matching or grouping.

The transformed features capture both direct pairwise similarity and third-party agreement — how other instances in the set influence similarity between a given pair. This enables BPA to model higher-order relations effectively.

---

---

Few-Shot Classification Results

Dataset	Method	5-Way 1-Shot	5-Way 5-Shot
MiniImagenet	PTMAP-BPAp	83.19	89.56
	PTMAP-BPAt	84.18	90.51
	PTMAP-SF-BPA	85.59	91.34
CIFAR-FS	PTMAP-BPAp	87.37	91.12
	PTMAP-SF-BPA	89.94	92.83
CUB	PTMAP-BPAp	91.90	94.63
	PTMAP-SF-BPA	95.80	97.12

---

Setup

Datasets

Instructions for downloading and preparing the few-shot classification datasets are available in the datasets directory.

Pretrained Models

Most results in the paper use fine-tuned models:

• WideResNet-28: PT-MAP checkpoint
• ResNet-12: FEAT checkpoint

BPA Checkpoint

We provide a checkpoint for PTMAP-BPA_t, which yields:

• 84.69% accuracy for 5-way 1-shot (vs. 84.18% in the paper)
• 90.30% accuracy for 5-way 5-shot (vs. 90.51% in the paper)

---

Usage

Quick Start

BPA can be applied in just two lines of code:

import torch
from bpa import BPA

x = torch.randn(100, 128)  # [n_samples, dim]
x = BPA()(x)               # Output shape: [n_samples, n_samples]

---

Training PT-MAP-BPA_t on MiniImagenet

1. Download the pretrained WRN feature extractor.
2. Create an empty checkpoints directory.
3. Extract the downloaded file into the checkpoints folder.

Run the following command to train BPA with PT-MAP:

python train.py \
  --sink_iters 5 \
  --distance_metric cosine \
  --ot_reg 0.2 \
  --method pt_map_bpa \
  --backbone wrn \
  --augment false \
  --lr 5e-5 \
  --weight_decay 0. \
  --max_epochs 50 \
  --train_way 10 \
  --val_way 5 \
  --num_shot 5 \
  --num_query 15 \
  --train_episodes 200 \
  --eval_episodes 400 \
  --checkpoint_dir ./checkpoints/pt_map_bpa \
  --data_path <your_dataset_path/miniimagenet/> \
  --pretrained_path ./checkpoints/miniImagenet/WideResNet28_10_S2M2_R/470.tar

Alternatively, use the trained checkpoint.

---

Evaluation

You can evaluate using the original PT-MAP repository:

1. Clone the PT-MAP repository.
2. Navigate to method/pt_map/evaluation/ and replace the files with the ones from this repository.
3. Edit the following:
4. Set checkpoint_dir in save_plk.py to the location of your BPA checkpoint.
5. Update _datasetFeaturesFiles in FSLTask.py to point to your feature files.
6. Create feature files by running:

python save_plk.py --dataset miniImagenet --method S2M2_R --model WideResNet28_10

1. Then run the evaluation:

   python test_standard.py

---

Alternatively, evaluate directly within this repository

You can run evaluation directly using the same script as training, with a few extra flags:

python train.py \
  --eval true \
  --pretrained_path <path_to_checkpoint> \
  --backbone <backbone_name> \
  --test_episodes 2000 ```

Make sure to include the same arguments you used during training (e.g., --method, --data_path, etc.) so the evaluation runs consistently.

---

Logging

We support logging and visualization via Weights & Biases (wandb). This helps track your training and evaluation metrics in real-time.

To enable logging:

1. Install the wandb package:

pip install wandb

1. Add the following flags to your command:

--wandb true --project <project_name> --entity <wandb_entity>

Replace <project_name> with your wandb project name, and <wandb_entity> with your wandb username or team name.

---

Citation

If you find this repository useful in your research, please cite:

@article{shalam2024balanced,
  title={The Balanced-Pairwise-Affinities Feature Transform},
  author={Shalam, Daniel and Korman, Simon},
  journal={arXiv preprint arXiv:2407.01467},
  year={2024}
}

---

Acknowledgment

Leveraging the Feature Distribution in Transfer-based Few-Shot Learning

S2M2 Charting the Right Manifold: Manifold Mixup for Few-shot Learning

Few-Shot Learning via Embedding Adaptation with Set-to-Set Functions