Title: Differentiable Clustering with Perturbed Spanning Forests
Venue: NeurIPS 2023
Accepted: yes
Keywords: Structured learning

Clustering

Differentiable

weakly supervised

semi-supervised

representation learning
Abstract: We introduce a differentiable clustering method based on stochastic perturbations of minimum-weight spanning forests. This allows us to include clustering in end-to-end trainable pipelines, with efficient gradients. We show that our method performs well even in difficult settings, such as data sets with high noise and challenging geometries. We also formulate an ad hoc loss to efficiently learn from partial clustering data using this operation. We demonstrate its performance on several data sets for supervised and semi-supervised tasks.

Differentiable Clustering with Perturbed Spanning Forests

Abstract

We introduce a differentiable clustering method based on stochastic perturbations of minimum-weight spanning forests. This allows us to include clustering in end-toend trainable pipelines, with efficient gradients. We show that our method performs well even in difficult settings, such as data sets with high noise and challenging geometries. We also formulate an ad hoc loss to efficiently learn from partial clustering data using this operation. We demonstrate its performance on several data sets for supervised and semi-supervised tasks. 2 .

1 Introduction

Clustering is one of the most classical tasks in data processing, and one of the fundamental methods in unsupervised learning (Hastie et al., 2009). In most formulations, the problem consists in partitioning a collection of n elements into k clusters, in a manner that optimizes some criterion, such as intracluster proximity, or some resemblance criterion, such as a pairwise similarity matrix. This procedure is naturally related to other tasks in machine learning, either by using these induced classes in supervised problems, or by either evaluating or looking for well-clustered representations (Caron et al., 2018; Xie et al., 2016). Its performance and flexibility on a wide range of natural dataset, that makes it a good downstream or preprocessing task, also make it a a very important candidate to learn representations in a supervised fashion. Yet, it is a fundamentally 'combinatorial' problem, representing a discrete decision, much like many classical algorithmic methods (e.g., sorting, taking nearest neighbours, dynamic programming).

For these reasons, it is extremely challenging to learn through clustering . As a function, the solution of a clustering problem is piece-wise constant with respect to its inputs (such as a similarity or distance matrix), and its gradient would therefore be zero almost everywhere. This operation is therefore naturally ill-suited to the use of gradient-based approaches to minimize an objective, which are at the center of optimization procedures for machine learning. It does not have the convenient properties of commonly used operations in end-to-end differentiable systems, such as smoothness and differentiability. Another challenge of using clustering as part of a learning pipeline is perhaps its ambiguity : even for a given notion of distance or similarity between elements, there are several valid definitions of clustering, with criteria adapted to different uses. Popular methods include k -means, whose formulation is NP-hard in general even in simple settings (Drineas et al., 2004), and relies on heuristics that depend heavily on their initialization (Arthur and Vassilvitskii, 2007; Bubeck et al., 2012). Several of them rely on proximity to a centroid , or prototype , i.e., a vector representing each cluster. These fail on challenging geometries, e.g., interlaced clusters with no linear separation.

We propose a new method for differentiable clustering that efficiently addresses these difficulties. It is a principled, deterministic operation: it is based on minimum-weight spanning forests, a variant

of minimum spanning trees. We chose this primitive at the heart of our method because it can be represented as a linear program (LP), and this is particularly well-adapted to our smoothing technique. However, we use a greedy algorithm to solve it exactly and efficiently (rather than solving it as an LP or relying on an uncertain heuristic). We observe that this method, often referred to as single linkage clustering (Gower and Ross, 1969), is effective as a clustering method in several challenging settings. Further, we are able to create a differentiable version of this operation, by introducing stochastic perturbations on the similarity matrix (the cost of the LP). This proxy has several convenient properties: it approximates the original function, it is smooth (both of these are controlled by a temperature parameter), and both the perturbed optimizers and its derivatives can be efficiently estimated with Monte-Carlo estimators (see, e.g., Hazan et al., 2016; Berthet et al., 2020, and references therein). This allows us to include this operation in end-to-end differentiable machine learning pipelines, and we show that this method is both efficient and performs well at capturing the clustered aspect of natural data, in several tasks.

Our work is part of an effort to include unconventional operations in model training loops based on gradient computations. These include discrete operators such as optimal transport, dynamic time-warping and other dynamic programs (Cuturi, 2013; Cuturi and Blondel, 2017; Mensch and Blondel, 2018; Blondel et al., 2020b; Vlastelica et al., 2019; Paulus et al., 2020; Sander et al., 2023; Shvetsova et al., 2023), to ease their inclusion in end-to-end differentiable pipelines that can be trained with first-order methods in fields such as computer vision, audio processing, biology, and physical simulators (Cordonnier et al., 2021; Kumar et al., 2021; Carr et al., 2021; Le Lidec et al., 2021; Baid et al., 2023; Llinares-López et al., 2023) and other optimization algorithms (Dubois-Taine et al., 2022). Other related methods, including some based on convex relaxations and on optimal transport, aim to minimize an objective involving a neural network in order to perform clustering, with centroids (Caron et al., 2020; Genevay et al., 2019; Xie et al., 2016). Most of them involve optimization in order to cluster, without explicitly optimizing through the clustering operation, allowing to learn using some supervision. Another line of research is focused on using the matrix tree-theorem to compute marginals of the Gibbs distribution over spanning trees (Koo et al., 2007; Zmigrod et al., 2021). It is related to our perturbation smoothing technique (which yields a different distribution), and can also be used to learn using tree supervision. The main difference is that it does not address the clustering problem, and that optimizing in this setting involves sophisticated linear algebra computations based on determinants, which is not as efficient and convenient to implement as our method based on sampling and greedy algorithms.

There is an important conceptual difference with the methodology described by Berthet et al. (2020), and other recent works based on it that use Fenchel-Young losses (Blondel et al., 2020a): while one of the core discrete operations on which our clustering procedure is based is an LP, the cluster connectivity matrix-which, importantly, has the same structure as any ground-truth label informationis not. This weak supervision is a natural obstacle that we have to overcome: on the one hand, it is reasonable to expect clustering data to come in this form, i.e., only stating whether some elements belong to the same cluster or not, which is weaker than ground-truth spanning forest information would be. On the other hand, linear problems on cluster connectivity matrices, such as MAXCUT, are notoriously NP-hard (Karp, 1972), so we cannot use perturbed linear oracles over these sets (at least exactly). In order to handle these two situations together, we design a partial Fenchel-Young loss , inspired by the literature on weak supervision, that allows us to use a loss designed for spanning forest structured prediction, even though we have the less informative cluster connectivity information. In our experiments, we show that our method enables us to learn through a clustering operation, i.e., that we can find representations of the data for which clustering of the data will match with some ground-truth clustering data. We apply this to both supervised settings, including some illustrations on synthetic challenging data, and semi-supervised settings on real data, focusing on settings with a low number of labeled instances and unlabeled classes.

Main Contributions. In this work, we introduce an efficient and principled technique for differentiable clustering. In summary, we make the following contributions:

Our method is based on using spanning forests, and a differentiable proxy, obtained by adding stochastic perturbations on the edge costs.

Our method allows us to learn through clustering : one can train a model to learn 'clusterable' representations of the data in an online fashion. The model generating the representations is informed by gradients that are transmitted through our clustering operation.

We derive a partial Fenchel-Young loss, which allows us to incorporate weak cluster information , and can be used in any weakly supervised structured prediction problem.

We show that it is a powerful clustering technique, allowing us to learn meaningful clustered representations. It does not require a model to learn linearly separable representations.

Our operation can be incorporated efficiently into any gradient-based ML pipeline. We demonstrate this in the context of supervised and semi-supervised cluster learning tasks .

Notations. For a positive integer n , we denote by [ n ] the set { 1 , . . . , n } . We consider all graphs to be undirected. For the complete graph with n vertices K n over nodes [ n ] , we denote by T the set of spanning trees on K n , i.e., subgraphs with no cycles and one connected component. For any positive integer k ≤ n we also denote by C k the set of k -spanning forests of K n , i.e., subgraphs with no cycles and k connected components. With a slight abuse of notation, we also refer to T and C k for the set of adjacency matrices of these graphs. For two nodes i and j in a general graph, we write i ∼ j if the two nodes are in the same connected component. We denote by S n the set of n × n symmetric matrices.

2 Differentiable clustering

We provide a differentiable operator for clustering elements based on a similarity matrix. Our method is explicit : it relies on a linear programming primitive, not on a heuristic to solve another problem (e.g., k -means). It is label blind : our solution is represented as a connectivity matrix, and is not affected by a permutation of the clusters. Finally it is geometrically flexible : it is based on single linkage, and does not rely on proximity to a centroid, or linear separability to include several elements in the same cluster (as illustrated in Section 4).

In order to cluster n elements in k clusters, we define a clustering operator as a function M ∗ k ( · ) taking an n × n symmetric similarity matrix as input (e.g., negative pairwise square distances) and outputting a cluster connectivity matrix of the same size (see Definition 1). We also introduce its differentiable proxy M ∗ k,ε ( · ) . We use them as a method to learn through clustering . As an example, in the supervised learning setting, we consider the n elements described by features X = ( X 1 , . . . , X n ) in some feature space X , and ground-truth clustering data as an n × n matrix M Ω (with either complete or partial information, see Definition 4). A parameterized model produces a similarity matrix Σ = Σ w ( X ) -e.g., based on pairwise square distances between embeddings Φ w ( X i ) for some model Φ w . Minimizing in w a loss so that M ∗ k (Σ = Σ w ( X )) ≈ M Ω allows to train a model based on the clustering information (see Figure 1). We describe in this section and the following one the tools that we introduce to achieve this goal, and show in Section 4 experimental results on real data.

Features

Model

Embeddings

Similarity

matrix

Prediction

Partial

information

Clustering

Diff. Clustering

Loss gradient

2.1 Clustering with k -spanning forests

In this work, the core operation is to cluster n elements, using a similarity matrix Σ ∈ S n . Informally, pairs of elements ( i, j ) with high similarity Σ ij should be more likely to be in the same cluster than those with low similarity. The clustering is represented in the following manner.

Definition 1 (Cluster connectivity matrix) . Let π : [ n ] → [ k ] be a clustering function, assigning elements to one of k clusters. We represent it with an n × n binary matrix M (the cluster connectivity):

We denote by B k ⊆ { 0 , 1 } n × n the set of binary cluster connectivity matrices with k clusters.

Using this definition, we define an operation M ∗ k ( · ) mapping a similarity matrix Σ to a clustering, in the form of such a membership matrix. Up to a permutation (i.e., naming the clusters), M allows to recover π entirely. It is based on a maximum spanning forest primitive A ∗ k ( · ) , and both are defined below. We recall that a k -forest on a graph is a subgraph with no cycles consisting in k connected components, potentially single nodes (see Figure 2). The cluster connectivity matrix is sometimes referred to as the cluster coincidence matrix matrix in other literature, and the two terms can be used interchangeably.

Definition 2 (Maximum k -spanning forest) . Let Σ be an n × n similarity matrix. We denote by A ∗ k (Σ) the adjacency matrix of the k -spanning forest with maximum similarity, defined as

This defines a mapping A ∗ k : S n →C k . We denote by F k the value of this maximum, i.e.,

Definition 3 (Spanning forest clustering) . Let A be the adjacency matrix of a k -spanning forest. We denote by M ∗ ( A ) the connectivity matrix of the k connected components of A , i.e.,

Given an n × n similarity matrix Σ , we denote by M ∗ k (Σ) ∈ B k the clustering induced by the maximum k -spanning forest, defined by

This defines a mapping M ∗ k : S n →B k , our clustering operator.

Remarks. The solution of the linear program is unique almost everywhere for similarity matrices (relevant for us, as we use stochastic perturbations in learning). Both these operators can be computed using Kruskal's algorithm to find a minimum spanning tree in a graph (Kruskal, 1956). This algorithm builds the tree by iteratively adding edges in a greedy fashion, maintaining non-cyclicity. On a connected graph on n nodes, after n -1 edge additions, a spanning tree (i.e., that covers all nodes) is obtained. The greedy algorithm is optimal for this problem, which can be proved by showing that forests can be seen as the independent sets of a matroid (Cormen et al., 2022). Further, stopping the algorithm after n -k edge additions yields a forest consisting of k trees (possibly singletons), together spanning the n nodes and which is optimal for the problem in Definition 2. As in several other clustering methods, we specify in ours the number of desired clusters k , but a consequence of our algorithmic choice is that one can compute the clustering operator for all k without much computational overhead, and that this number can easily be tuned by validation.

Further, a common manner to run this algorithm is to keep track of the constructed connected components, therefore both A ∗ k (Σ) and M ∗ k (Σ) are actually obtained by this algorithm. The mapping M ∗ is of course many-to-one, and yields a partition of C k into equivalence classes of k -forests that yield the same clusters. This point is actually at the center of our operators of clustering with constraints in Definition 5 and of our designed loss in Section 3.2, both below. We note that as the maximizer of a linear program, when the solution is unique, we have ∇ Σ F k (Σ) = A ∗ k (Σ) .

As described above, our main goal is to enable the inclusion of these operations in end-to-end differentiable pipelines (see Figure 1). We include two important aspects to our framework in order to do so, and to efficiently learn through clustering. The first one is the inclusion of constraints, enforced values coming either from partial or complete clustering information (see below, Definition 5), and the second one is the creation of a differentiable version of these operators, using stochastic perturbations (see Section 2.2).

Clustering with constraints. Given information, we also consider constrained versions of these two problems. We represent the enforced connectivity information as a matrix M Ω , defined as follows.

Definition 4 (Partial cluster coincidence matrix) . Let M be a cluster connectivity matrix, and Ω a subset of [ n ] × [ n ] , representing the set of observations. We denote by M Ω the n × n matrix

Remarks. The symbol ' ∗ ' in this definition is only used as a placeholder, an indicator of 'no information', and does not have a mathematical use. A common setting is where only a subset S ⊆ [ n ] of the data has label information. In this case for any i, j ∈ S , M Ω ,ij = 1 if i and j share the same label, otherwise M Ω ,ij = 0 i.e. elements in the same class are clustered together and separated from elements in other classes.

Definition 5 (Constrained maximum k -spanning forest) . Let Σ be an n × n similarity matrix. We denote by A ∗ k (Σ; M Ω ) the adjacency matrix of the k -spanning forest with maximum similarity, constrained to satisfy the connectivity constraints in M Ω . It is defined as

where for any partial clustering matrix M Ω , C k ( M Ω ) is the set of k -spanning forests whose clusters agree with M Ω , i.e.,

For any partial connectivity matrix M Ω , this defines another mapping A ∗ ( · ; M Ω ) : S n →C k .

We denote by F k (Σ; M Ω ) the value of this maximum, i.e.,

Similarly we define M ∗ k (Σ; M Ω ) = M ∗ ( A ∗ k (Σ; M Ω )) .

Remarks. We consider these operations in order to infer clusterings and spanning forests that are consistent with observed information. That is particularly important when designing a partial loss to learn from observed clustering information. Note that depending on what the set of observations Ω is, these constraints can be more or less subtle to satisfy. For certain sets of constraints Ω , when the matroid structure is preserved, we can obtain A ∗ k (Σ; M Ω ) by the usual greedy algorithm, by additionally checking that any new edge does not violate the constraints defined by Ω . This is for example the case when Ω corresponds to exactly observing a fully clustered subset of points.

2.2 Perturbed clustering

The clustering operations defined above are efficient and perform well (see Figure 2) but by their nature as discrete operators, they have a major drawback: they are piece-wise constant and as such cannot be conveniently included in end-to-end differentiable pipelines, such as those used to train models such as neural networks. To overcome this obstacle, we use a proxy for our operators, by introducing a perturbed version (Abernethy et al., 2016; Berthet et al., 2020; Paulus et al., 2020; Struminsky et al., 2021), obtained by taking the expectation of solutions with stochastic additive noise on the input. In these definitions and the following, we consider Z ∼ µ from a distribution with positive, differentiable density over S n , and ε > 0 .

Definition 6. We define the perturbed maximum spanning forest as the expected maximum spanning forest under stochastic perturbation on the inputs. Formally, for a similarity matrix Σ , we have

We define analogously A ∗ k,ε (Σ; M Ω ) = E [ A ∗ k (Σ + εZ ; M Ω )] and F k,ε (Σ; M Ω ) = E [ F k (Σ + εZ ; M Ω )] , as well as clustering M ∗ k,ε (Σ) = E [ M ∗ k (Σ + εZ )] and M ∗ k,ε (Σ; M Ω ) = E [ M ∗ k (Σ + εZ ; M Ω )] .

We note that this defines operations A ∗ k,ε ( · ) and A ∗ k,ε ( · ; M Ω ) mapping Σ ∈ S n to cvx ( C k ) , the convex hull of C k . These operators have several advantageous features: They are differentiable, and both their values and their derivatives can be estimated using Monte-Carlo methods, by averaging copies of A ∗ k (Σ + εZ ( i ) ) . These operators are the ones used to compute the gradient of the loss that we design to learn from clustering (see Definition 8 and Proposition 1).

This is particularly convenient, as it does not require to implement a different algorithm to compute the differentiable version. Moreover, the use of parallelization in modern computing hardware makes the computational overhead almost nonexistent. These methods are part of a large literature on using perturbations in optimizers such as LP solutions (Papandreou and Yuille, 2011; Hazan et al., 2013; Gane et al., 2014; Hazan et al., 2016), including so-called Gumbel max-tricks (Gumbel, 1954; Maddison et al., 2016; Jang et al., 2017; Huijben et al., 2022; Blondel et al., 2022)

Since M ∗ k,ε ( · ) is a differentiable operator from S n to cvx ( B k ) , it is possible to use any loss function on cvx ( B k ) to design a loss based on Σ and some ground-truth clustering information M Ω , such as L (Σ; M Ω ) = ∥ M ∗ k,ε (Σ) -M Ω ∥ 2 2 . This flexibility is one of the advantages of our method. In the following section, we introduce a loss tailored to be efficient to compute and performant in several learning tasks, that we call a partial Fenchel-Young loss .

3 Learning with differentiable clustering

3.1 Fenchel-Young losses

In structured prediction, a common modus operandi is to minimize a loss between some structured ground truth response or label y ∈ Y and a score θ ∈ R d (the latter often itself the output of a parameterized model). As an example, in logistic regression y ∈ { 0 , 1 } and θ = ⟨ x, β ⟩ ∈ R . The framework of Fenchel-Young losses allows to tackle much more complex structures, such as cluster connectivity information.

Definition 7 (Fenchel-Young loss-Blondel et al. (2020a)) . Let F be a convex function on R d , and F ∗ its Fenchel dual on a convex set C ⊂ R d . The Fenchel-Young loss between θ ∈ R d and y ∈ int ( C ) is

The Fenchel-Young (FY) loss satisfies several properties, making it useful for learning. In particular, it is nonnegative, convex in θ , handles well noisy labels, and its gradient with respect to the score can be efficiently computed, with ∇ θ L FY ( θ ; y ) = ∇ θ F ( θ ) -y (see, e.g., Blondel et al., 2020a).

In the case of linear programs over a polytope C , taking F to be the so-called support function defined as F ( θ ) = max y ∈C ⟨ y, C⟩ (consistent with our notation so far), the dual function F ∗ is the indicator function of C (Rockafellar, 1997), and the Fenchel-Young function is then given by

In the case of a perturbed maximum for F , see, e.g., Berthet et al. (2020), we have

In this setting (under mild regularity assumptions on the noise distribution and C ), we have that ε Ω = F ∗ ε is a strictly convex Legendre function on C and y ∗ ε ( θ ) = ∇ θ F ε ( θ ) . This is actually equivalent to having max and argmax with a -ε Ω( y ) regularization (see Berthet et al., 2020, Propositions 2.1

and 2.2). In this case, the FY loss is given by L FY ,ε ( θ ; y ) = F ε ( θ ) -⟨ θ, y ⟩ + ε Ω( y ) and its gradient by ∇ θ L FY ,ε ( θ ; y ) = y ∗ ε ( θ ) -y . One can also define it as L FY ,ε ( θ ; y ) = E [ L FY ( θ + εZ ; y )] and it has the same gradient. In the perturbations case, it can be easily obtained by Monte-Carlo estimates of the perturbed optimizer, taking 1 B ∑ B i =1 y ∗ ( θ + εZ i ) -y .

3.2 Partial Fenchel-Young losses

As noted above, this loss is widely applicable in structured prediction. As presented here, it requires label data y of the same kind than the optimizer y ∗ . In our setting, the linear program is over spanning k -forests rather than connectivity matrix. This is no accident: linear programs over cut matrices (when k = 2 ) are already NP-hard (Karp, 1972). If one has access to richer data (such as ground-truth spanning forest information), one can ignore the operator M ∗ and focus only on A ∗ k,ε , F k,ε , and the associated FY-loss. However, in most cases, clustering data can reasonably only be expected to be present as connectivity matrix. It is therefore necessary to alter the Fenchel-Young loss, and we introduce the following loss, which allows to work with partial information p ∈ P representing partial information about the unknown ground-truth y . Using this kind of 'inf-loss' is common when dealing with partial label information (see, e.g., Cabannes et al., 2020).

Definition 8 (Partial Fenchel-Young loss) . Let F be a convex function, L FY the associated FenchelYoung loss, and for every p ∈ P a convex constraint subset C ( p ) ⊆ C .

This allows to learn from incomplete information about y . When we do not know its value, but only a subset of C ( p ) ⊆ C to which it might belong, we can minimize the infimum of the FY losses that are consistent with the partial label information y ∈ Y ( p ) .

Proposition 1. When F is the support function of a compact convex set C , the partial Fenchel-Young loss (see Definition 8) satisfies

The loss is a difference of convex functions in θ given explicitly by

The gradient with respect to θ is given by

The perturbed partial Fenchel-Young loss given by ¯ L FY ,ε ( θ ; p ) = E [ ¯ L FY ( θ + εZ ; p )] satisfies ¯ .

Another possibility would be to define the partial loss as min y ∈C ( p ) L FY ,ε ( θ ; y ) , that is, the infimum of smoothed losses instead of the smoothed infimum loss L FY ,ε ( θ ; p ) defined above. However, there is no direct method to minimizing the smoothed loss with respect to y ∈ C ( p ) . Note that we have a relationship between the two losses:

Proposition 2. Letting L FY ,ε ( θ ; y ) = E [ L FY ,ε ( θ + εZ ; y )] and ¯ L FY ,ε as in Definiton 8, we have

The proofs of the above two propositions are detailed in the Appendix.

3.3 Applications to differentiable clustering

We apply this framework, as detailed in the following section and in Section 4, to clustering. This is done naturally by transposing notations and taking C = C k , θ = Σ , y ∗ = A ∗ k , p = M Ω , C ( p ) = C k ( M Ω ) , and y ∗ ( θ ; p ) = A ∗ k (Σ; M Ω ) . In this setting the perturbed partial Fenchel-Young loss satisfies

We learn representations of a data that fit with clustering information (either complete or partial). As described above, we consider settings with n elements described by their features X = X 1 , . . . , X n in some feature space X , and M Ω some clustering information. Our pipeline to learn representations includes the following steps (see Figure 1)

Embed each X i in R d with a parameterized model v i = Φ w ( X i ) ∈ R d , with weights w ∈ R p .

ii) Construct a similarity matrix from these embeddings, e.g. Σ w,ij = -∥ Φ w ( X i ) -Φ w ( X j ) ∥ 2 2 .

iii) Stochastic optimization of the expected loss of ¯ L FY ,ε (Σ w,b , M Ω ,b ) , using mini-batches of X .

Details. To be more specific on each of these phases: i) We embed each X i individually with a model Φ w , using in our application neural networks and a linear model. This allows us to use learning through clustering as a way to learn representations, and to apply this model to other elements, for which we have no clustering information, or for use in other downstream tasks. ii) We focus on cases where the similarity matrix is the negative squared distances between those embeddings. This creates a connection between a model acting individually on each element, and a pairwise similarity matrix that can be used as an input for our differentiable clustering operator. This mapping, from w to Σ , has derivatives that can therefore be automatically computed by backpropagation, as it contains only conventional opereations (at least when Φ w is itself a conventional model, such as commonly used neural networks). iii) We use our proposed smoothed partial Fenchel-Young (Section 3.2) as the objective to minimize between the partial information M Ω and Σ . The full-batch version would be to minimize L FY ,ε (Σ w , M Ω ) as a function of the parameters w ∈ R p of the model. We focus instead on a mini-batch formulation for two reasons: first, stochastic optimization with mini-batches is a commonly used and efficient method for generalization in machine learning; second, it allows to handle larger-scale datasets. As a consequence, the stochastic gradients of the loss are given, for a mini-batch b , by

The simplicity of these gradients is due to our particular choice of smoothed partial Fenchel-Young loss. They can be efficiently estimated automatically, as described in Section 2.2, which results in a doubly stochastic scheme for the loss optimization.

4 Experiments

We apply our framework for learning through clustering in both a supervised and a semi-supervised setting, as illustrated in Figure 1. Formally, for large training datasets of size n , we either have access to a full cluster connectivity matrix M Ω or a partial one (typically built by using partial label information, see below). We use this clustering information M Ω , from which mini-batches can be

NMI: 1.00, SC: 0.11, CoL2: 0.00

NMI: 0.00, SC: 0.35, CoL2: 0.50

Kmeans

NMI: 0.00, SC: 0.36, CoL2: 0.50

MeanShift

NMI: 0.00, SC: 0.35, CoL2: 0.50

GM

NMI: 1.00, SC: 0.38, CoL2: 0.00

NMI: 0.36, SC: 0.50, CoL2: 0.27

NMI: 0.42, SC: 0.50, CoL2: 0.24

NMI: 0.36, SC: 0.50, CoL2: 0.27

NMI: 1.00, SC: 0.80, CoL2: 0.00

NMI: 1.00, SC: 0.80, CoL2: 0.00

NMI: 1.00, SC: 0.80, CoL2: 0.00

NMI: 1.00, SC: 0.80, CoL2: 0.00

NMI: 0.93, SC: 0.40, CoL2: 0.03

NMI: 0.45, SC: 0.49, CoL2: 0.26

NMI: 0.56, SC: 0.57, CoL2: 0.31

NMI: 0.56, SC: 0.48, CoL2: 0.23

0

1000

2000

3000

4000

5000

Number of labelled samples

0.0

0.1

0.2

0.3

0.4

0.5

0.6

0.7

Clustering error

kw

=0

kw

=3

kw

=6

0

1000

2000

3000

4000

5000

Number of labelled samples

0.02

0.04

0.06

0.08

0.10

Classification error

kw

=0

kw

=3

kw

=6

2000

4000

6000

8000

10000

Number of labelled samples

0.0

0.1

0.2

0.3

0.4

0.5

0.6

0.7

0.8

Clustering error

kw

=0

kw

=3

kw

=6

2000

4000

6000

8000

10000

Number of labelled samples

0.20

0.25

0.30

0.35

0.40

0.45

0.50

0.55

0.60

0.65

Classification error

kw

=0

kw

=3

kw

=6

extracted, as supervision. We minimize our partial Fenchel-Young loss with respect to the weights of an embedding model, and evaluate these embeddings in two main manners on a test dataset: first, by evaluating the clustering accuracy (i.e. proportion of correct coefficients in the predicted cluster connectivity matrix), and second by training a shallow model on a classification task (using clusters as classes) on a holdout set, evaluating it on a test set.

4.1 Supervised learning

We apply first our method to synthetic datasets - purely to provide an illustration of both our internal clustering algorithm, and of our learning procedure. In Figure 2, we show how the clustering operator that we use, based on spanning forests (i.e. single linkage), with Kruskal's algorithm is efficient on some standard synthetic examples, even when they are not linearly separable (compared to k -Means, mean-shift, Gaussian mixture-model). We also show that our method allows us to perform supervised learning based on cluster information, in a linear setting. For X signal consisting of n = 60 points in two dimensions consisting of data in four well-separated clusters (see Figure 2), we construct X by appending two noise dimensions, such that clustering based on pairwise square distances between X i mixes the original clusters. We learn a linear de-noising transformation Xθ , with θ ∈ R 4 × 2 through clustering, by minimizing our perturbed partial FY loss with SGD, using the known labels as supervision.

We also show that our method is able to cluster virtually all of some classical datasets. We train a CNN (LeNet-5 LeCun et al. (1998)) on mini-batches of size 64 using the partial Fenchel-Young loss to learn a clustering, with a batch-wise clustering precision of 0.99 for MNIST and 0 . 96 on Fashion MNIST. Using the same approach, we trained a ResNet (He et al., 2016) on CIFAR-10 (with some minor modifications to kernel size and stride), achieving a batch-wise clustering test precision of 0 . 93 . The t-SNE visualization of the embeddings for the validation set are displayed in Figure 5. The experimental setups (as well as visualization of learnt clusters for MNIST), are detailed in the Appendix.

4.2 Semi-supervised learning

We show that our method is particularly useful in settings where labelled examples are scarce, even in the particularly challenging case of having no labelled examples for some classes. To this end, we conduct a series of experiments on the MNIST (LeCun et al., 2010) and Cifar-10 (Krizhevsky et al., 2009) datasets, where we independently vary both the total number of labelled examples n ℓ , as well as the number of withheld classes for which no labelled examples are present in the training set k w ∈ { 0 , 3 , 6 } . For MNIST we consider data sets with n ℓ ∈ { 100 , 250 , 500 , 1000 , 2000 , 5000 } labelled examples, and for Cifar-10 we consider n ℓ ∈ { 1000 , 2500 , 5000 , 7500 , 10 , 000 } .

We train the same embedding models using our partial Fenchel-Young loss with batches of size 64. We use ε = 0 . 1 and B = 100 for the estimated loss gradients, and optimize weights using Adam (Kingma and Ba, 2015).

In left column of Figure 4, we report the test clustering error (evaluated on mini-batches of the same size), for each of the models, and data sets, corresponding to the choice of n ℓ and k w . We compare the performance of each model to a baseline, consisting of the exact same architecture (plus a linear head mapping the embedding to logits), trained to minimize the cross entropy loss.

In addition, we evaluate all models on a down-stream (transfer-learning) classification task, by learning a linear probe on top of the frozen model, trained on hold-out data with all classes present. The results are depicted in the right hand column of Figure 4. Further information regarding the experimental details can be found in the Appendix.

We observe that learning through clustering allows to find a representation where class semantics are easily recoverable from the local topology. On Cifar-10, our proposed approach strikingly achieves a lower clustering error in the most challenging setting ( n ℓ = 1000 labelled examples and k w = 6 withheld classes) than the classification-based baseline in the most lenient setting ( n ℓ = 10 , 000 labelled examples all classes observed). Importantly, this advantage is not limited to clustering metrics: learning through clustering can also lead to lower down-stream classification error, with the gap being most pronounced when few labelled examples are available.

Moreover, besides this pronounced improvement in data efficiency, we found that our method is capable of clustering classes for which no labelled examples are seen during training (see Figure 5, left). Therefore, investigating potential applications of learning through clustering to zero-shot and self-supervised learning are promising avenues for future work.

References

Abadi, M., Barham, P., Chen, J., Chen, Z., Davis, A., Dean, J., Devin, M., Ghemawat, S., Irving, G., Isard, M., et al. (2016). Tensorflow: a system for large-scale machine learning. In Osdi , volume 16, pages 265-283. Savannah, GA, USA.

Abernethy, J., Lee, C., and Tewari, A. (2016). Perturbation techniques in online learning and optimization. Perturbations, Optimization, and Statistics , 233.

Arthur, D. and Vassilvitskii, S. (2007). K-means++: the advantages of careful seeding. In Proceedings of the Annual Symposium on Discrete Algorithms , pages 1027-1035.

Baid, G., Cook, D. E., Shafin, K., Yun, T., Llinares-López, F., Berthet, Q., Belyaeva, A., Töpfer, A., Wenger, A. M., Rowell, W. J., et al. (2023). Deepconsensus improves the accuracy of sequences with a gap-aware sequence transformer. Nature Biotechnology , 41(2):232-238.

Berthet, Q., Blondel, M., Teboul, O., Cuturi, M., Vert, J.-P., and Bach, F. (2020). Learning with differentiable pertubed optimizers. Advances in Neural Information Processing Systems , 33:95089519.

Blondel, M., Berthet, Q., Cuturi, M., Frostig, R., Hoyer, S., Llinares-López, F., Pedregosa, F., and Vert, J.-P. (2022). Efficient and modular implicit differentiation. Advances in Neural Information Processing Systems , 35:5230-5242.

Blondel, M., Martins, A. F., and Niculae, V. (2020a). Learning with Fenchel-Young losses. The Journal of Machine Learning Research , 21(1):1314-1382.

Blondel, M., Teboul, O., Berthet, Q., and Djolonga, J. (2020b). Fast differentiable sorting and ranking. In International Conference on Machine Learning , pages 950-959.

Bradbury, J., Frostig, R., Hawkins, P., Johnson, M. J., Leary, C., Maclaurin, D., Necula, G., Paszke, A., VanderPlas, J., Wanderman-Milne, S., and Zhang, Q. (2018). JAX: composable transformations of Python+NumPy programs.

Bubeck, S., Meil˘ a, M., and von Luxburg, U. (2012). How the initialization affects the stability of the k-means algorithm. ESAIM: Probability and Statistics , 16:436-452.

Cabannes, V., Rudi, A., and Bach, F. (2020). Structured prediction with partial labelling through the infimum loss. In International Conference on Machine Learning , pages 1230-1239.

Caron, M., Bojanowski, P., Joulin, A., and Douze, M. (2018). Deep clustering for unsupervised learning of visual features. In Proceedings of the European conference on computer vision (ECCV) , pages 132-149.

Caron, M., Misra, I., Mairal, J., Goyal, P., Bojanowski, P., and Joulin, A. (2020). Unsupervised learning of visual features by contrasting cluster assignments. Advances in neural information processing systems , 33:9912-9924.

Carr, A. N., Berthet, Q., Blondel, M., Teboul, O., and Zeghidour, N. (2021). Self-supervised learning of audio representations from permutations with differentiable ranking. IEEE Signal Processing Letters , 28:708-712.

Cordonnier, J.-B., Mahendran, A., Dosovitskiy, A., Weissenborn, D., Uszkoreit, J., and Unterthiner, T. (2021). Differentiable patch selection for image recognition. In Proceedings of the Conference on Computer Vision and Pattern Recognition , pages 2351-2360.

Cormen, T. H., Leiserson, C. E., Rivest, R. L., and Stein, C. (2022). Introduction to Algorithms . MIT Press.

Cuturi, M. (2013). Sinkhorn distances: Lightspeed computation of optimal transport. Advances in Neural Information Processing Systems , 26.

Cuturi, M. and Blondel, M. (2017). Soft-DTW: a differentiable loss function for time-series. In International Conference on Machine Learning , pages 894-903.

Defazio, A., Bach, F., and Lacoste-Julien, S. (2014). Saga: A fast incremental gradient method with support for non-strongly convex composite objectives. Advances in neural information processing systems , 27.

Drineas, P., Frieze, A., Kannan, R., Vempala, S., and Vinay, V. (2004). Clustering large graphs via the singular value decomposition. Machine Learning , 56:9-33.

Dubois-Taine, B., Bach, F., Berthet, Q., and Taylor, A. (2022). Fast stochastic composite minimization and an accelerated frank-wolfe algorithm under parallelization. In Advances in Neural Information Processing Systems .

Gane, A., Hazan, T., and Jaakkola, T. (2014). Learning with maximum a-posteriori perturbation models. In Proc. of AISTATS , pages 247-256.

Genevay, A., Dulac-Arnold, G., and Vert, J.-P. (2019). Differentiable deep clustering with cluster size constraints. arXiv preprint arXiv:1910.09036 .

Gower, J. C. and Ross, G. J. (1969). Minimum spanning trees and single linkage cluster analysis. Journal of the Royal Statistical Society: Series C (Applied Statistics) , 18(1):54-64.

Gumbel, E. J. (1954). Statistical Theory of Extreme Values and some Practical Applications: A Series of Lectures . Number 33. US Govt. Print. Office.

Hastie, T., Tibshirani, R., Friedman, J. H., and Friedman, J. H. (2009). The elements of statistical learning: data mining, inference, and prediction , volume 2. Springer.

Hazan, T., Maji, S., and Jaakkola, T. (2013). On sampling from the gibbs distribution with random maximum a-posteriori perturbations. In Advances in Neural Information Processing Systems , pages 1268-1276.

Hazan, T., Papandreou, G., and Tarlow, D. (2016). Perturbations, Optimization, and Statistics . MIT Press.

He, K., Zhang, X., Ren, S., and Sun, J. (2016). Deep residual learning for image recognition. In Proceedings of the IEEE conference on Computer Vision and Pattern Recognition , pages 770-778.

Huijben, I. A., Kool, W., Paulus, M. B., and Van Sloun, R. J. (2022). A review of the gumbel-max trick and its extensions for discrete stochasticity in machine learning. IEEE Transactions on Pattern Analysis and Machine Intelligence , 45(2):1353-1371.

Jang, E., Gu, S., and Poole, B. (2017). Categorical reparameterization with gumbel-softmax. In International Conference on Learning Representations .

Karp, R. M. (1972). Reducibility among combinatorial problems. In Proceedings of a symposium on the Complexity of Computer Computations, held March 20-22, 1972, at the IBM Thomas J. Watson Research Center, Yorktown Heights, New York, USA , pages 85-103.

Kingma, D. P. and Ba, J. (2015). Adam: A method for stochastic optimization. In Bengio, Y. and LeCun, Y., editors, 3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Track Proceedings .

Koo, T., Globerson, A., Carreras Pérez, X., and Collins, M. (2007). Structured prediction models via the matrix-tree theorem. In Joint Conference on Empirical Methods in Natural Language Processing and Computational Natural Language Learning , pages 141-150.

Krizhevsky, A., Hinton, G., et al. (2009). Learning multiple layers of features from tiny images.

Kruskal, J. B. (1956). On the shortest spanning subtree of a graph and the traveling salesman problem. Proceedings of the American Mathematical Society , 7(1):48-50.

Kumar, A., Brazil, G., and Liu, X. (2021). Groomed-nms: Grouped mathematically differentiable nms for monocular 3d object detection. In Proceedings of the Conference on Computer Vision and Pattern Recognition , pages 8973-8983.

Le Lidec, Q., Laptev, I., Schmid, C., and Carpentier, J. (2021). Differentiable rendering with perturbed optimizers. Advances in Neural Information Processing Systems , 34:20398-20409.

LeCun, Y., Bottou, L., Bengio, Y., and Haffner, P. (1998). Gradient-based learning applied to document recognition. Proceedings of the IEEE , 86(11):2278-2324.

LeCun, Y., Cortes, C., and Burges, C. (2010). Mnist handwritten digit database. ATT Labs [Online]. Available: http://yann.lecun.com/exdb/mnist , 2.

Llinares-López, F., Berthet, Q., Blondel, M., Teboul, O., and Vert, J.-P. (2023). Deep embedding and alignment of protein sequences. Nature Methods , 20(1):104-111.

Maddison, C. J., Mnih, A., and Teh, Y. W. (2016). The concrete distribution: A continuous relaxation of discrete random variables. arXiv preprint arXiv:1611.00712 .

Mensch, A. and Blondel, M. (2018). Differentiable dynamic programming for structured prediction and attention. In Proc. of ICML .

Papandreou, G. and Yuille, A. L. (2011). Perturb-and-MAP random fields: Using discrete optimization to learn and sample from energy models. In Proc. of ICCV , pages 193-200.

Paszke, A., Gross, S., Massa, F., Lerer, A., Bradbury, J., Chanan, G., Killeen, T., Lin, Z., Gimelshein, N., Antiga, L., et al. (2019). Pytorch: An imperative style, high-performance deep learning library. Advances in neural information processing systems , 32.

Paulus, M., Choi, D., Tarlow, D., Krause, A., and Maddison, C. J. (2020). Gradient estimation with stochastic softmax tricks. Advances in Neural Information Processing Systems , 33:5691-5704.

Rockafellar, R. T. (1997). Convex analysis , volume 11. Princeton university press.

Sander, M. E., Puigcerver, J., Djolonga, J., Peyré, G., and Blondel, M. (2023). Fast, differentiable and sparse top-k: a convex analysis perspective. arXiv preprint arXiv:2302.01425 .

Shvetsova, N., Petersen, F., Kukleva, A., Schiele, B., and Kuehne, H. (2023). Learning by sorting: Self-supervised learning with group ordering constraints. In Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV) , pages 16453-16463.

Struminsky, K., Gadetsky, A., Rakitin, D., Karpushkin, D., and Vetrov, D. P. (2021). Leveraging recursive gumbel-max trick for approximate inference in combinatorial spaces. Advances in Neural Information Processing Systems , 34:10999-11011.

Vlastelica, M., Paulus, A., Musil, V ., Martius, G., and Rolínek, M. (2019). Differentiation of blackbox combinatorial solvers. arXiv preprint arXiv:1912.02175 .

Xie, J., Girshick, R., and Farhadi, A. (2016). Unsupervised deep embedding for clustering analysis. In International conference on machine learning , pages 478-487.

Yang, B., Fu, X., Sidiropoulos, N. D., and Hong, M. (2017). Towards k-means-friendly spaces: Simultaneous deep learning and clustering. In International Conference on Machine Learning , pages 3861-3870. PMLR.

Zmigrod, R., Vieira, T., and Cotterell, R. (2021). Efficient computation of expectations under spanning tree distributions. Transactions of the Association for Computational Linguistics , 9:675-690.

5 Broader impact

This submission focuses on foundational and exploratory work, with application to general machine learning techniques. We propose methods to extend the range of operations that can be used in end-to-end differentiable systems, by allowing to include clustering as a differentiable operation in machine learning models. We do not foresee societal consequences that are specifically related to these methods, beyond those that are associated with the field in general.

6 Reproducibility and Licensing

Our experiments use data sets that are already open-sourced and cited in the references. All the code written for this project and used to implement our experiments is available at:

https://github.com/LawrenceMMStewart/DiffClust_NeurIPS2023 , and distributed under a Apache 2.0. license.

7 Limitations

At present, our implementation of Kruskal's algorithm is incompatible with processing very large batch sizes at train time. However, as our method can use mini-batches, it is suitable for training deep learning models on large data sets. For example, we use a batch size of 64 (similar to that for standard classification models) for all experiments detailed in our submission. Other methods that focus on cluster assignment, rather than pairwise clustering, often have an n × k parametrization, whose efficiency will depend on the comparison between n (batch size) and k (number of clusters).

At inference time this is not the case, since gradients need not be back-propagated hence, any implementation of Kruskal's algorithm can be used such as the union-find implementation. Faster GPU compatible implementations of Kruskal's algorithm for training are certainly possible ; improvements could be made using a binary heap or a GPU parallelized merge sort. In many learning experiments the batch sizes are typically of values { 32 , 64 , 128 , 256 , 512 } , so our current implementation is sufficient, however, for some specific applications (such as bio-informatics experiments requiring tens of thousands of points in a batch), it would be necessary to improve upon the current implementation. This is mainly an engineering task relating to data-structures and GPU programming, and is beyond the scope of this paper.

Finally, our clustering approach solves a linear program whose argmax solution is a matroid problem with a greedy single-linkage criterion. Whilst we saw that this method is effective for differentiable clustering in both the supervised and semi-supervised setting, there may well be other linkage criteria which also take the form of LP's but whom are more robust to outliers and noise in the train dataset. This line of work is outside the scope of the paper, but is closely related and could help improve our differentiable spanning forests framework.

8 Proofs of Technical Results

Proof of Proposition 1. We show these properties successively

1) We have by Definition 8

We expand this

As required, this implies, following the definitions given

2) By linearity of derivatives and 1) above, we have

Since the argmax of a constrained linear optimization problem is the gradient of its value. We note that this property holds almost everywhere (when the argmax is unique), and almost surely for costs with positive, continuous density, which we always assume (e.g. see the following).

3) By linearity of expectation and 1) above, we have

using the definition

Proof of Proposition 2. By Jensen's inequality and the definition of the Fenchel-Young loss:

9 Algorithms for Spanning Forests

As mentioned in Section 2, both A ∗ k (Σ) and M ∗ k (Σ) are calculated using Kruskal's algorithm (Kruskal, 1956). Our implementation of Kruskal's is tailored to our use: we first initialize both A ∗ k (Σ) and M ∗ k (Σ) as the identity matrix, and then sort the upper triangular entries of Σ . We build the maximumweight spanning forest in the usual greedy manner, using A ∗ k (Σ) to keep track of edges in the forest and M ∗ k (Σ) to check if an edge can be added without creating a cycle, updating both matrices at each step of the algorithm. Once the forest has k connected components, the algorithm terminates. This is done by keeping track of the number of edges that have been added at any time.

We remark that our implementation takes the form as a single loop, with each step of the loop consisting only of matrix multiplications. For this reason it is fully compatible with auto-differentiation engines, such as JAX (Bradbury et al., 2018), Pytorch (Paszke et al., 2019) and TensorFlow (Abadi et al., 2016), and suitable for GPU/TPU acceleration. Therefore, our implementation differs from that of the standard Kruskal's algorithm, which used a disjoint union-find data structure (and hence is not compatible with auto-differentiation frameworks).

9.1 Constrained Spanning Forests

As an heuristic way to solve the constrained problem detailed in Section 5, we make the modifications below to our implementation of Kruskal's, under the assumption that M Ω represents valid clustering information (i.e. with no contradiction):

Regularization (Optional) : It is possible to bias the optimization problem over spanning forests to encourage or discourage edges between some of the nodes, according to the clustering information. Before performing the sort on the upper-triangular of Σ , we add a large value to all entries of Σ ij where ( M Ω ) ij = 1 , and subtract this same value from all entries of Σ ij where ( M Ω ) ij = 1 . Entries Σ ij corresponding to where ( M Ω ) ij = ⋆ are left unchanged. This biasing ensures that any edge between points that are constrained to be in the same cluster will always be processed before unconstrained edges. Similarly, any edge between points that are constrained to not be in the same cluster, will be processed after unconstrained edges. In most cases, i.e. when all clusters are represented in the partial information, such as when Ω = [ n ] × [ n ] (full information), this is not required to solve the constrained linear program, but we have found that this regularization was helpful in practice.

Constraint enforcement : We ensure that adding an edge does not violate the constraint matrix. In other words, when considering adding the edge ( i, j ) to the existing forest, we check that none of the points in the connected component of i are forbidden from joining any of the points in the connected component of j . This is implemented using further matrix multiplications and done alongside the existing check for cycles. The exact implementation is detailed in our code base.

10 Existing Literature on Differentiable Clustering

As discussed in Section 1, there exists many approaches which use clustering during gradient based learning, but these approaches typically use clustering in an offline fashion in order to assign labels to points. The following methods aim to learn through a clustering step (i.e. gradients back-propagate through clustering):

Yang et al. (2017) use a bi-level optimization procedure (alternating between optimizing model weights and centroid clusters). They reported attaining 83% label-wise clustering accuracy on MNIST using a fully-connected deep network. Our method differs from this approach as it is allows for end-to-end online learning.

Genevay et al. (2019) cast k-means as an optimal transport problem, and uses entropic regularization for smoothing. Reported a 85% accuracy on MNIST and 25% accuracy on CIFAR-10 with a CNN.

11 Additional Experimental Information

0.0

0.2

0.4

0.6

0.8

1.0

1.2

We provide the details of the synthetic denoising experiment depicted in Figure 2 and described in Section 4.1. We create the signal data X signal ∈ R 60 × 2 by sampling iid. from four isotropic Gaussians (15 points coming from each of the Gaussians) each having a standard deviation of 0 . 2 . We randomly sample the means of the four Gaussians; for the example seen in Section 4.1 the sampled means were:

.

Let Σ signal be the pairwise euclidean similarity matrix corresponding to X signal, and furthermore let M signal := M ∗ 4 (Σ signal ) be the equivalence matrix corresponding to the signal ( M signal will be the target equivalence matrix to learn).

̸

We append an additional two 'noise dimensions' to X signal in order to form the train data X ∈ R 60 × 4 , where the noise entries were sampled iid from a continuous unit uniform distribution. Similarly letting Σ be the pairwise euclidean similarity matrix corresponding to X , we calculate M ∗ 4 (Σ) = M signal . Both the matrices M signal and M ∗ 4 (Σ) are depicted in Figure 6 ; we remark that adding the noise dimensions leads to most points being assigned one of two clusters, and two points being isolated alone in their own clusters. We also create a validation set of equal size (in exactly the same manner as the train set), to ensure the model has not over-fitted to the train set.

The goal of the experiment is to learn a linear transformation of the data that recovers M signal i.e. a denoising, by minimizing the partial loss. There are multiple solutions to this problem, the most obvious being a transformation that removes the last two noise columns from X :

For any θ ∈ R 4 × 2 , we define Σ θ to be the pairwise similarity matrix corresponding to Xθ , and M ∗ 4 (Σ θ ) to its corresponding equivalence matrix. Then the problem can be summarized as:

We initialized θ from a standard Normal distribution, and minimized the partial loss via stochastic gradient descent, with a learning rate of 0 . 01 and batch size 32.

For perturbations, we took ε = 0 . 1 and B = 1000 , where ε denotes the noise amplitude in randomized smoothing and B denotes the number of samples in the Monte-Carlo estimate. The validation clustering error converged to zero after 25 gradient batches. We verify that the θ attained from training is indeed learning to remove the noise dimensions (see Figure 6).

11.1 Supervised Differentiable Clustering

ReLU

ReLU & Avg Pool: window 2x2, strides = 2x2

ReLU & Avg Pool: window 2x2, strides = 2x2

Conv: features = 32, kernel = 3x3

Conv: features = 64, kernel = 3x3

Dense: d = 256

Dense: d = 256

As mentioned in Section 4.1, our method is able to cluster classical data sets such as MNIST and Fashion MNIST. We trained a CNN with the LeNet-5 architecture LeCun et al. (1998) using our proposed partial loss as the objective function. The exact details of the CNN architecture are depicted in Figure 7. For this experiment and all experiments described below, we trained on a single Nvidia V100 GPU ; training the CNN with our proposed pipeline took < 15 minutes.

The model was trained for 30 k gradient steps on mini-batches of size 64 . We used the Adam optimizer (Kingma and Ba, 2015) with learning rate η = 3 × 10 -4 , momentum parameters ( β 1 , β 2 ) = (0 . 9 , 0 . 999) , and an ℓ 2 weight decay of 10 -4 . We validated / tested the model using the zero-one error between the true equivalence matrix and the equivalence matrix corresponding to the output of the model. We used an early stopping of 10k steps (i.e. training was stopped if the validation clustering error did not improve over 10k steps). For efficiency (and parallelization), we also computed this clustering error batch-wise with batch-size 64. As stated in Section 4.1, we attained a batch-wise clustering precision of 0.99 for MNIST and 0 . 96 on Fashion MNIST.

The t-SNE visualizations of the embedding space of the model trained on MNIST for a collection of train and validation data points are depicted in Figure 7. It can be seen that the model has learnt a distinct cluster for each of the ten classes.

In similar fashion, we trained a ResNet (He et al., 2016) on the Cifar-10 data set. The exact model architecture is similar to that of ResNet-50, but with minor modifications to the input convolutions for compatibility with the dimensions of Cifar images, and is detailed in the code base.

The training procedure was identical to that of the CNN, except the model was trained for 75k steps (with early stopping), and used the standard data augmentation methods for Cifar-10, namely: a

combination of four-pixel padding, random flip followed by a random crop. As mentioned in Section 4.1, the model achieved a batch-wise clustering test precision of 0 . 933 .

11.2 Semi-Supervised Differentiable Clustering

As mentioned in Section 4.2, we show that our method is particularly useful in settings where labelled examples are scarce, even in the particularly challenging case of having no labelled examples for some classes. Our approach allows a model to leverage the semantic information of unlabeled examples when trying to predict a target equivalence matrix M Ω ; this is owing to the fact that the prediction of a class for a single point depends on the values of all other points in the batch, which is in general not the case for more common problems such as classification and regression.

To demonstrate the performance of our approach, we assess our method on two tasks:

Semi-Supervised Clustering: Train a model to learn an embedding of the data which leads to a good clustering error. We can compare our methodology to that of a baseline model trained using the cross-entropy loss. This is to check that our model has leveraged information from the unlabeled data and that our partial loss is indeed leading to good clustering performance.

Downstream Classification: Assess the trained model's capacity to serve as a backbone in a downstream classification task (transfer learning), where its weights are frozen and a linear layer is trained on top of the backbone.

We describe our data processing for both of these tasks below.

11.2.1 Data Sets

In our semi-supervised learning experiments, we divided the standard MNIST and Cifar-10 train splits in the following manner:

We create a balanced hold-out data set consisting of 1k images (100 images from each of the 10 classes). This hold-out data set will be used to assess the utility of the frozen clustering model on a downstream classification problem.

From the remaining 59k images, we select a labeled train set of n ℓ points (detailed in Section 4.2). Our experiments also vary k w ∈ { 0 , 3 , 6 } , the number of digits to withhold all labels from. For example, if k w = 3 , then the labels for the images corresponding to digits { 0 , 1 , 2 } will never appear in the labeled train data.

11.2.2 Semi-Supervised Clustering Task

For each of the choices of n ℓ and k w , we train the architectures described in Section 11.1 using the following two approaches:

Ours: The model is trained on mini-batches, where half the batch is labeled data and half the batch is unlabeled data (i.e. a semi-supervised learning regime), to minimize the partial loss.

Baseline: The baseline model shares the same architecture as that described in Section 11.1, but with an additional dense layer with output dimension 10 (the number of classes). We train the model using mini-batches consisting of labeled points, minimizing the crossentropy loss. The training regime is fully-supervised learning (classification). The baseline backbone refers to all of the model, minus the dense output layer.

Both models were trained with mini-batches of size 64, with points sampled uniformly without replacement. All hyper-parameters and optimization metrics were identical to those detailed in Section 4.1. For MNIST, we repeated training for each model with five different random seeds s ∈ [5] (and with three random seeds s ∈ [3] for Cifar), in order to report population statistics on the clustering error.

11.2.3 Transfer-Learning: Downstream Classification

In this task both models are frozen, and their utility as a foundational backbone is assessed on a downstream classification task using the hold-out data set. We train a linear (a.k.a dense) layer on top of both models using the cross-entropy loss. We refer to this linear layer as the downstream head. Training this linear head is equivalent to performing multinomial logistic regression on the features of the model.

To optimize the weights of the linear head we used the SAGA optimizer (Defazio et al., 2014). The results are depicted in Figure 4. It can be seen that training a CNN backbone using our approach with just 250 labels leads to better downstream classification performance than the baseline trained with 5000 labels. It is worth remarking that the baseline backbone was trained on the same objective function (cross-entropy) and task (classification) as the downstream problem, which is not the case for the backbone corresponding to our approach. This highlights how learning 'cluster-able embeddings' and leveraging unlabeled data can be desirable for transfer learning.

Peer reviews:
comments: SUMMARY:
This paper proposes a differentiable clustering method based on a stochastic perturbation of mnimum-weight spanning forests. They demonstrate its performance on several datasets in supervised and semi-supervised settings.

STRENGTHS:
This paper introduces a differentiable clustering method based on linear programs for minimum-weight spanning forests. 
Their method starts with formulating the clustering method as spanning forest problem, which admits an LP (linear programming) formulation. Then, they make the parametrization differentiable using stochastic perturbation. They also propose a novel partial Fenchel-Young loss, which allows for learning through clustering in weakly supervised problem.

The technical part of this paper is strong. The idea of partial Fenchel-Young loss is novel. The theoretical part is well written and not hard to follow.

WEAKNESSES:
For experiment part, 

- it might be better to elaborate more on the settings and what and how $M_\Omega$ is given in these settings so readers can better understand the tasks. 

- Comparison with other methods is lacking: maybe some SOTA nondifferentiable clustering methods 

- Since the goal is to learn $M$, an $n\times n$ matrix, would it be some computational efficiency/scalability issue of the proposed method when $n$ is large? Other methods might have more condensed parametrization than the proposed framework?

QUESTIONS:
See Weaknesses

LIMITATIONS:
See Weaknesses
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-02

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper presents a novel method for differentiable clustering in combination with a novel loss derived from Fenchel-Young losses.\n\nSTRENGTHS:\nThe paper proposes a new concept for differentiable clustering.\nDifferentiable clustering has been a topic with only few applicable methods so far, so it is great to see a new and at the same time strong method being proposed in this space.\n\nThe method is sound and the utility of the method is experimentally demonstrated.\n\nWEAKNESSES:\nIt would have been great to see a stronger experiment. Maybe CIFAR-10 / -100? Or for self-supervised learning on STL or ImageNet?\nI would increase my score for a stronger experiment. \n\n\nThe main aspect of being able to transmit gradients backward through a clustering operation, a good representation of data points, as computed e.g. by some neural network, can be learned. That is, the original data, which might be difficult to cluster, maybe even due to its dimensionality, is not clustered directly, but a learnable representation of it is clustered, and this representation is informed by the gradients transmitted backwards through the clustering algorithm. I believe this should be emphasized more in the introduction.\n\nThe clustering function \\pi is usually described by a k x n matrix U that is called \"partition matrix\". The matrix that is called a \"cluster membership matrix\" here (not a good name, because it does not state which data points belong to which clusters -- that's the partition matrix) is usually known as \"coincidence matrix\" or \"cluster connectivity matrix\" (e.g. in the context of relative cluster evaluation measures -- it states for each set of data points whether they are assigned to the same cluster or not). I would recommend using this standard terminology. The entries of a coincidence matrix M can be computed from the entries of a partition matrix U as M_jl = \\sum_{i=1}^k U_ij U_il. Note that this works also in the fuzzy or membership degree case (i.e. U_ij \\in [0,1]). In this sense, Def. 4 does not specify membership constraints, but rather coincidence constraints or connectivity constraints.\n\nThe terms \"must link\" and \"must not link\" that occur in the caption of Figure 1 could have been introduced after Def. 4, as \"must link\": M_\\Omega,ij = 1, \"must not link\": M_\\Omega,i,j = 0. In this sense M is a \"cluster connectivity matrix\" (see above).\n\nThe connection to single linkage hierarchical agglomerative clustering with cluster merging carried out until k clusters remain could have been pointed out.\n\n126: a n x n -> an n x n\n\n\n[A] seems to be a relevant but missing reference.\n\n[A] Struminsky, K., Gadetsky, A., Rakitin, D., Karpushkin, D., and Vetrov, D. P. Leveraging recursive gumbel-max trick for approximate inference in combinatorial spaces. Advances in Neural Information Processing Systems, 2021.\n\nQUESTIONS:\nsee above.\n\nLIMITATIONS:\nThe experimental evaluation is limited.", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces a differentiable clustering method based on minimum-weight spanning forests.\nThis method can be integrated into end-to-end trainable pipelines and handle datasets with high noise and challenging geometries.\nThe key idea is to smooth the combinatorial clustering operation by taking expectations over stochastic perturbations of the input similarity matrix. \nBy taking expectations over these perturbations, gradients can be computed and backpropagated through the clustering process, facilitating end-to-end training.\nExperiments on both supervised and semi-supervised tasks demonstrate the effectiveness of the proposed method.\n\nSTRENGTHS:\n1. The authors propose an effective method to address a challenging problem: learning through clustering. \nThe proposed method is well-motivated and reasonable, and it has the potential to be applied in a wide range of applications.\nThis paper addresses the key challenges of incorporating combinatorial operations like clustering into neural network pipelines that can be optimized through gradient descent. The perturbed proxy for discrete operators can address the major issue of piece-wise constant by attaining useful properties like smoothness and computational traceability. \n\n2. The paper is well-written and easy to follow. Detailed remarks are provided for readers to better understand the definitions.\n\nWEAKNESSES:\nMy major concern is about the evaluation. \n\n1. The proposed method is only evaluated on small-scale datasets such as MNIST. It remains unclear whether the proposed method can be scaled to large datasets.\n\n2. The authors did not compare the proposed method with other differentiable clustering methods. Therefore, it remains unclear whether the proposed method is more efficient than other baselines.\n\nQUESTIONS:\nPlease check the weakness.\n\nLIMITATIONS:\nThe authors claim that the proposed method can be used to handle datasets with high noise, but this claim is not verified in the experiments.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-10"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper presents a differentiable clustering approach based on k-spanning forests. It formulates the clustering assignment as an NxN membership matrix while introducing partial Fenchel-Young losses to optimize the between the cluster constraint and similarity/affinity matrix. The study demonstrates the applicability of the proposed approach to both supervised and semi-supervised learning for representation learning through clustering. The results indicate the following: (1) On several synthetic datasets, the proposed method generated qualitatively better clusters compared to baseline clustering methods including k-means, MeanShift, and Gaussian mixture model. (2) In supervised/semi-supervised representation learning experiments, the proposed method effectively learns a robust representation based on the labels (cluster constraint), outperforming the classification approach on the MNIST/Fashion-MNIST datasets.\n\nSTRENGTHS:\n- The proposed approached based on k-spanning tree and the partial Fenchel-Young losses for clustering is technical sound.\n- The proposed approach enable end-to-end optimization for representation learning through clustering which is a novel and interesting direction. Experimental results on MNIST also shows that interestingly the method was able to outperform classification loss on semi-supervised tasks especially on low-data regime.\n\nWEAKNESSES:\n- Experimental results are only limited to toy datasets (synthesized data, MNIST), it would be interesting to see how the methods expands to real world data such as ImageNet.\n- Comparison between traditional clustering methods are also limited to qualitative comparison, it would be good to have more quantitive  comparison between different clustering methods and demonstrate the benefits of the proposed clustering method.\n\nQUESTIONS:\n- How does the proposed clustering methods compared to traditional clustering algorithm (e.g. k-means) on real worlds data? The paper only present qualitative experiments on toy dataset, it would be more interesting to know how it compare quantitatively on real world dataset.\n- How does the proposed method perform on more complicate dataset such as ImageNet compare to traditional classification approach? It would make the paper stronger if we could show that how the proposed method could be incorporate into ML pipeline to demonstrate benefits over more complicate dataset.\n\nLIMITATIONS:\nThe authors have listed the possible broader impact and limitations (e.g. batch size and current implementation efficiency) in the supplemental materials.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-18"}