Skip to content

colure/GeometricMeanLoss

Repository files navigation

[ICIP2025] Geometric Mean Improves Loss For Few-Shot Learning

Python PyTorch arXiv

Official implementation of Geometric Mean Improves Loss For Few-Shot Learning, by Tong Wu and Takumi Kobayashi.

Abstract

Few-shot learning (FSL) is a challenging task in machine learning, demanding a model to render discriminative classification by using only a few labeled samples. In the literature of FSL, deep models are trained in a manner of metric learning to provide metric in a feature space which is well generalizable to classify samples of novel classes; in the space, even a few amount of labeled training examples can construct an effective classifier. In this paper, we propose a novel FSL loss based on geometric mean to embed discriminative metric into deep features. In contrast to the other losses such as utilizing arithmetic mean in softmax-based formulation, the proposed method leverages geometric mean to aggregate pair-wise relationships among samples for enhancing discriminative metric across class categories. The proposed loss is not only formulated in a simple form but also is thoroughly analyzed in theoretical ways to reveal its favorable characteristics which are favorable for learning feature metric in FSL. In the experiments on few-shot image classification tasks, the method produces competitive performance in comparison to the other losses.

Installation

We recommend using uv for fast dependency management.

git clone https://github.com/colure/GeometricMeanLoss.git
cd GeometricMeanLoss

# Initialize and install dependencies automatically
uv sync

Usage

1. Data Preparation

Datasets

Place datasets in the ./data directory (e.g., ./data/mini_imagenet, ./data/cifar_fs). You can adjust default paths in ./config.toml.

2. Configuration

  • ./config.toml: Centralized configuration for all hyperparameters and paths.
  • CLI Overrides: Override any TOML setting via command line (e.g., --batch-size 64).

3. Run

We utilize Distributed Data Parallel (DDP) for multi-GPU training. Use uv run to execute with the managed environment:

# To train the model:
uv run torchrun --nproc_per_node=${num_gpus} train.py --loss GMLoss --logit l1_dist --output-dir ./result/my_experiment
# To evaluate a checkpoint:
uv run python train.py --resume ./result/my_experiment/checkpoints/checkpoint_best_shot5.pth --test-only

More examples can be found in ./train_script.sh.

Results

Loss miniImageNet CIFAR-FS tieredImageNet
1-shot 5-shot 1-shot 5-shot 1-shot 5-shot
PN 62.42±0.20 79.13±0.15 67.14±0.22 82.36±0.16 66.74±0.23 82.14±0.17
NCA 62.68±0.20 78.93±0.54 69.20±0.21 84.24±0.16 67.21±0.22 83.77±0.16
GML 65.51±0.20 81.13±0.14 71.09±0.22 85.08±0.16 69.61±0.23 84.04±0.16

Citation

If you find this project or paper useful in your research, please consider citing the original paper:

@misc{wu2025geometricmeanimprovesloss,
      title={Geometric Mean Improves Loss For Few-Shot Learning},
      author={Tong Wu and Takumi Kobayashi},
      year={2025},
      eprint={2501.14593},
      archivePrefix={arXiv},
      primaryClass={cs.CV},
      url={https://arxiv.org/abs/2501.14593},
}

Acknowledgments

Parts of this code inspired by the following projects:

License

MIT License.

About

Official PyTorch implementation of Geometric Mean Improves Loss For Few-Shot Learning.

Topics

Resources

License

Stars

Watchers

Forks

Releases

No releases published

Contributors