The Balanced-Pairwise-Affinities Feature Transform

TL;DR

The paper introduces the Balanced-Pairwise-Affinities (BPA) feature transform, a novel, efficient, and differentiable method that encodes high-order relations among features, improving performance in various tasks like classification and clustering.

Contribution

It presents a new feature transform based on optimal transport that is efficient, differentiable, and improves downstream task performance.

Findings

Achieves state-of-the-art results in few-shot classification.

Enhances unsupervised image clustering performance.

Improves person re-identification accuracy.

Abstract

The Balanced-Pairwise-Affinities (BPA) feature transform is designed to upgrade the features of a set of input items to facilitate downstream matching or grouping related tasks. The transformed set encodes a rich representation of high order relations between the input features. A particular min-cost-max-flow fractional matching problem, whose entropy regularized version can be approximated by an optimal transport (OT) optimization, leads to a transform which is efficient, differentiable, equivariant, parameterless and probabilistically interpretable. While the Sinkhorn OT solver has been adapted extensively in many contexts, we use it differently by minimizing the cost between a set of features to $itself$ and using the transport plan's $rows$ as the new representation. Empirically, the transform is highly effective and flexible in its use and consistently improves networks it is inserted into, in a variety of tasks and training schemes. We demonstrate state-of-the-art results in few-shot classification, unsupervised image clustering and person re-identification. Code is available at \url{github.com/DanielShalam/BPA}.