The Balanced-Pairwise-Affinities Feature Transform

TL;DR

The paper introduces Balanced-Pairwise-Affinities (BPA), a parameterless, differentiable feature transform that re-embeds a set of item features by solving a self-transport problem between the set and itself. BPA yields a symmetric, doubly-stochastic embedding ${W}$ whose rows encode each item’s distribution of affinities to all others, providing a probabilistic and interpretable representation for set-input tasks. The authors demonstrate state-of-the-art or near state-of-the-art results across few-shot classification, unsupervised image clustering, and person re-identification by inserting BPA into diverse hosting networks, both in transductive and inductive settings, and provide extensive ablations and practical guidance. The method sits at the intersection of set-to-set processing, attention, spectral methods, and optimal transport, offering a distinct self-transport perspective with a simple, drop-in, end-to-end trainable pipeline. Code is released to facilitate adoption and replication of results, underscoring BPA’s practical impact for grouping and matching tasks.

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}.

Figures (10)

• Figure 1: Generic designs of networks that act on sets of items. These cover relevant architectures, e.g. for few-shot-classification and clustering. Left: A generic network for processing a set of input items typically follows the depicted structure: (i) Each item separately goes through a common feature extractor $F$. (ii) The set of extracted features is the input to a downstream task processing module $G$. ; Right: A more general structure in which the extracted features undergo a joint processing by a transform $T$. Our BPA transform (as well as other attention mechanisms) is of this type and its high-level design (within the 'green' module) is detailed in Fig. \ref{['fig.BPA_illustration']}.
• Figure 2: The BPA transform: illustrated on a toy 7 image 3-class MNIST example.
• Figure 3: The min-cost max-flow perspective: Costs are shown.
• Figure 4: The (symmetric) embedding matrix $\mathcal{W}$ and the absolute difference between its $i$th and $j$th rows.
• Figure 5: BPA scaling in terms of accuracy and efficiency.