Differentiable Clustering with Perturbed Spanning Forests
Authors: Lawrence Stewart
Francis Bach
Felipe Llinares-López
Quentin Berthet
Venue: NeurIPS 2023 poster (2023)
Keywords: Structured learning
Clustering
Differentiable
weakly supervised
semi-supervised
representation learning
TL;DR: A method for differentiable clustering and techniques to learn from clustering data, end-to-end.

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
Lawrence Stewart" ENS & INRIA Paris, France
Francis Bach ENS & INRIA Paris, France
Felipe Llinares- Google DeepMind Paris, France
Quentin Berthet Google DeepMind Paris, France
Abstract
We introduce differentiable clustering method based on stochastic perturbations of mninimum-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 hoc loss to efficiently learn from partial clustering data using this operation. We demonstrate its performance on several data sets for supervised and semi sumi-supervised tasks.?
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 partitioning collection of n elements into & clusters, in manner that optimizes cluster proximity, or some resemblance criterion, such as 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 that makes it a good downstream or preprocessing task, also make it a very important candidate to learn representations in supervised fashion. Yet, it is fundamentally "combinatorial" problem, representing a discrete decision, much like many classical algorithmic methods (e.g., sorting, taking nearest neighbours, dynamic programming).
criterion, such intra-
wide range of natural dataset,
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 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 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 smoothness and differentiability. Another challenge of using clustering as part of learning pipeline is perhaps its ambiguity: even for a given notion of distance similarity between elements, there several valid definitions clustering, with criteria adapted to different uses. Popular methods include k-means, whose formulation is NP-hard in general even in simple settings (Drineas al., 2004), and relies heuristics that depend heavily on their initialization (Arthur and Vassilvitskii, 2007; Bubeck et al., 2012). Several of them rely on proximity centroid, or prototype, i.c., vector representing each cluster. These fail on challenging geometries, e.g., interlaced clusters with no linear separation.
We propose new method for differentiable clustering that efficiently addresses these difficulties. It a principled, deterministic operation: is based on minimum menimum--eightt weight spanning forests, a variant
correspondence should be addressed to lawrrence Codee base: https://gi thub. cttp://aarraccccessssaarr/ooott IPS2023
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

of minimum spanning trees. We chose this primitive at the heart of our method because it can be represented as linear program (LP), and this is particularly well-adapted our smoothing technique. However, we use 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 al., 2023) and other optimization algorithms (Dubois- TDineis-Taine et al.,
et al., 2023; Llinarcs-López 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 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 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 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 Tenchel-Youne losses (Blondel et al., 2020a): while one of the core discrete operations on which our clustering procedure is based is LP, the cluster connectivity matrix-which, importantly, has the same structure as any ground-truth label information- not. This weak supervision natural obstacle that we have to overcome: on the one hand, is reasonable to expect clustering data to come in this form, i.c., only stating whether some elements belong the same cluster or not, which weaker than ground-truth spanning forest information would be. On the other hand, linear problems on cluster connectivity matrices, such as MAXCUT, are notoriously Na-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 clustering operation, i.e., that we can find representations of the data for which clustering of the data will match with some ground- -tuuh clustering data. We apply this to both supervised settings, including some illustrations on synthetic challenging data, and semi-supereised settings on real data, focusing settings with low number of labeled instances and unlabeled classes.
Main Contributions. In this work, we introduce an efficient and principled technique for differen- tiable 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 informed by gradients that are transmitted through clustering operation.
2

---

We derive partial Fenchel- Young loss, which allows to incorporate weak cluster information, and can be used in any weakly supervised structured prediction problem.
We show that it is powerful clustering technique, allowing representations. It does not require model to learn linearly separable representations.
learn meaningful clustered
Our operation can 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 be undirected. For the complete graph with vertices Kn nodes [n], denote by 'T the set of spanning frees on K,, i.e., subgraphs with no cycles and one connected component. For any positive integer S
also denote by Cj the set of k-spanning forests of Kn ie.. subgraphs with no cycles and k connected components, With a slight abuse of notation, we also refer to T and C₁ for the of adjacency matrices of these graphs. For nodes and j in general graph, we write i ~ j if the two nodes are in the same connected component. We denote by S, the set of n x symmetric matrices.
Differentiable clustering
We provide differentiable operator for clustering elements based on a similarity matrix. Our method is explicit: it relies linear programming primitive, not heuristic to solve another problem (e.g., k-means). is label blind: our solution is represented as a connectivity matrix, and not affected by permutation of the clusters. Finally is geometrically flexible: it is based 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 clusters, we define a clustering operator as a function Mg taking an 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 method to learn through clustering. As an example, in the
proxy (.). We them supervised learning setting, consider the n elements described by features X (X1 1x..... X,) in some feature space X' and ground-truth clustering data as an n x n matrix Mq (with either complete partial information, see Definition 4). A parameterized model produces similarity matrix D = E ( (X)-c.g.. based on pairwise square distances between embeddings & (X.) for some model a Minimizing in loss so that M₁ 2i, (X)) ss Mq allows train 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 experimental results on real data.
T oo (Xi)
U:
Diff. Clustering
Mn
Similarity matrix
Clustering
Partial information
Features Model
Embeddings
A503) MC(E)
V
Prediction
Loss gradient
Figure 1: Our pipeline, in the semi-supervised learning setting: data points are embedded by a parameterized model, which produces similarity matrix. Partial clustering information may be available, in the form of must-link or must-not-link constraints. Our clustering and differentiable clustering operators used, respectively for prediction and gradient computations.

---

2.1 Clustering with k-spanning forests
In this work, the core operation to cluster n elements, using a similarity matrix Σ E Sn. Informally, pairs of elements (ž.j) with high similarity Eig should be more likely to be in the same cluster than those with low similarity. The clustering represented in the following manner.
Definition 1 (Cluster connectivity matrix). Let [n] + [K] be clustering function, assigning elements one of clusters. We represent with an n n binary matrix M (the cluster connectivity):
Mij =1 if and only if x(i)= = =[j).
We denote by BN S {0,1) 1j"X the set of binary cluster connectivity matrices with & clusters.
Using this definition, we define an operation (-)-) mapping a similarity matrix E to a clustering, in the form of such membership matrix. Up permutation (i.e., naming the clusters), M allows to recover π entirely. It is based on a maximum spanning forest primitive ( and both are defined below. We recall that a k-forest on 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 as the cluster coincidence matrix matrix other literature, and the two terms can be used interchangeably.
Definition 2 (Maximum k-spanning forest). Let E. be AE) the adjacency matrix of the k-spanning forest with maximum similarity, defined as
similarity matrix. We denote by
AE) argmax((, E).
AEC,
This defines a mapping A% : Sn - Ck. We denote by Fg the value of this maximum, i.e., Ex(
mA,
ΛECr.
Definition (Spanning forest clustering). Let be the adjacency matrix of a k-spanning forest. We denote by M*(A) (A) the connectivity matrix of the l connected components of A, i.e.,
Mij =1 if and only if i~ inj.
Given an n X n similarity matrix E, we denote by ME E BR the clustering induced by the maximum k-spanning forest, defined by
Mt(2)= = M (AL()).
This defines a mapping ME S., > BR. 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 minimum spanning tree in graph (Kruskal, 1956). This algorithm builds the tree by iteratively adding edges in greedy fashion, maintaining non-cyclicity. On a connected graph on n nodes, after n 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 A edge additions yields 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 l: without much computational overhead, and that this number can easily be tuned by validation.
Further, a common manner to run this algorithm to keep track of the constructed connected components, therefore both (E) and ME()) are actually obtained by this algorithm. The mapping M is of course many-to-one, and yields partition of C. 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 and of our designed loss in Section 3.2, both below. We note that as the maximizer of linear program, when the solution is unique, we have VyF+(E) = (2..).
As described above, our main goal is to enable the inclusion of these operations in end-to- differentiable pipelines (see Figure l). 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 the creation of a differentiable version of these operators, using stochastic perturbations (see Section 2.2).

---

A30
(E)
MY Mi)
Mo
Figure 2: Method illustration, for (left) Similarity matrix based on pairwise square distance. partial cluster connectivity information. (center) Clustering using spanning forests without partial clustering constraints. (right) Constrained clustering with partial constraints.
Clustering with constraints. Given information, we also consider constrained versions of these two problems. We represent the enforced connectivity information as a matrix Mg. defined as follows.
Definition 4 (Partial cluster coincidence matrix). Let be a cluster connectivity matrix, and Q subset of [n] x [n), representing the set of observations. We denote by Ma the matrix
Ma.ij = Mij (i. EΩ, Ma, = otherwise
Remarks. The symbol "x" in this definition is only used placeholder, indicator of "no information", and does not have mathematical use. A common setting is where only a subset S S [n]] of the data has label information. In this case for any i, S, My,i if and share the same label, otherwise Mn,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 E be an n x n similarity matrix. We denote by (E; Ma) the adjacency matrix of the k-spanning forest with maximum similarity, constrained to satisfy the connectivity constraints in Mq. It is defined as
AE;E; Mn) = arrgaxx (A,E), (Ma)
where for any partial clustering matrix Ma. (x(Ma) is the set of k-spanning forests whose clusters agree with Mg. i.e.,
Ca(Ma) {& MA) Ma,i (i,j) E Q) For any partial connectivity matrix Ma, this defines another mapping Ar Mn): Sn - Ca. We denote by (((E;Ms) the value of this maximum, i.e.,
Ps(E; Mn)
.nax (A, (Ma)
Similarly we define ME Mo) = M'(AA(E;
Remarks. We consider these operations in order to infer clusterings and spanning forests that consistent with observed information. That is particularly important when designing a partial loss to learn from observed clustering information. Note that depending what the set of observations Ω is, these constraints can be more or less subtle to satisfy. For certain sets of constraints S2 when the matroid structure is preserved, additionally checking that any new edge does not violate the constraints defined by Ω. This for example the when Ω corresponds exactly observing a fully clustered subset of points.
obtain AE; Ma) by the usual greedy algorithm, by
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 major drawback: they are piece- wise constant and as such cannot conveniently included in end-to-end differentiable pipelines, such as those used train models such as neural networks. To overcome this obstacle, we use proxy for our operators, by introducing a perturbed version (Abernethy al., 2016; Berthet et al., 2020; Paulus et al., 2020: Struminsky et al., 2021), obtained by taking the expectation of solutions with stochastic additíve noise on the input. In these definitions and the following, we consider Z l from a distribution with positive, differentiable density over Sn, and > 0.
5

---

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
Definition 6.
A (2) =
= E[ argmax A, e FR. E E(Fx eZ))
AEC₁
We define analogously (2;Mp) = E [[A (2) + Mp)] and (2;
EI
2 ; Ma)), as well as clustering M (Σ) = El M + +Z)] and ME Ma) = EJME(E
eZ; Mo)]
We note that this defines operations (-) and (-, Ma) mapping D E S, to cvx(C&), the convex hull of Ck- 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 (E + ) 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, 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; I; 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)
does not require to implement different algorithm to compute
Since M) ()) is a differentiable operator from S,, cVx( it is possible use any loss function on cvx(Ba) ) to design loss based on Σ and some ground-trutt -tuth clustering information Mg, such as L(E; Mn.) = || (Σ) Ma 112- 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 call partial Fenchel- Tenchel-Young loss.
3 Learning with differentiable clustering
3.1 Sinchel-Young losses
In structured prediction, a common modus operandi is to minimize a loss between some structured ground truth response or label y E y and score 0 E R (the latter often itself the output of parameterized model). As an example, in logistic regression y {0, 1} and 0 (a, B) E E. The framework of Fenchel- Young losses allows to tackle much more complex structures, such as cluster connectivity information.
Definition 7 (Fenchel-Yourg loss-Blondel al. (2020a)). Let F be convex function on Rdd and F its Fenchel dual on convex set C C. Rd. The Fenchel- Yennhe loss between 0 E IRd and y E inr(C) is
(0, = F()) - (0,4) + (y).
The Fenchel- Young (FY) loss satisfies several properties, making it useful for learning. In particular, it is nonnegative, convex in 0, handles well noisy labels, and its gradient with respect to the score can be efficiently computed, with VoLFy(0: y) F(F(
(see, e.g., Blondel et al., 2020a).
In the case of linear programs over polytope C, taking F to be the so-called support function defined as F(0)= maxyec (v, C)) (consistent with notation so far), the dual function F is the indicator function (Rockafellar, 1997), and the Fenchel- Fenchel-Young function is then given by
[F(0) - -0,9) if C, -1+o + otherwise.
(Fr =
In the case of perturbed maximum for F, see, e.g. Berthet et al. (2020), we have
(0) = El Immax (y, + +)), and ye = E E(argmax(y, + eZ)). yEC
yec
In this setting (under mild regularity assumptions on the noise distribution and C), we have that elt F a strictly convex Legendre function on C and y (0) = (o)= This actually equivalent to having max and argmax with a -Es2(y) regularization (see Berthet et al., 2020, Propositions 2.1

---

and 2.2). In this case, the FY loss is given by LFY.e (0:1) = (0) (0. + = and its gradient by VoLFY.c((::) y) y5 (0) y. One can also define as (e: ( + =Z; y)] and has the same gradient. In the perturbations case, it can be easily obtained by Monte-Carlo estimates of the perturbed optimizer, taking i E v" (0 + eZi) y-
3.2 Partial Fenchel- Fenchel-Young losses
As noted above, this loss is widely applicable structured prediction. As presented here, 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 & = 2) are already NP-hard (Karp, 1972). If one has access richer data (such ground-truth spanning forest information), one can ignore the operator M and focus only on A Fik 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- Srachel-Youg loss, and we introduce the following loss, which allows to work with partial information p E 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 convex function, LEY the associated Fenchel- Fenchel-Young
Young loss, and for every E C a convex constraint subset C(p) C.
LFy(0:p) min (Py yEC(p)
This allows to learn from incomplete information about y. When we do not know its value, but only a subset of C(p) S C to which it might belong, we can minimize the infimum of the FY losses that are consistent with the partial label information y E Y(p).
Proposition 1. When F is the support function of compact convex set C, the partial Fenchel-Your Young loss (see Definition 8) satisfies
1. The loss is a difference of convex functionn in O given explicitly by
LFY(0:P)= = F(0) - p),:p), where F(0:P) = max {y. f),
vEC(p)
2. The gradient with respect is given VoLFv(0:p) = (0) (0;p), where (0:pp) = arginax(,, 0).
yEC(p)
3. The perturbed partial Fenchel- Fenchel-Young loss given by LFY.s (0;p) = EI (0 (0+ +22;p]] satisfies
VaLFY,e = (0) 3. (#:p), where ((: = Elargmax 0 + E EZ)
vEC(p)
Another possibility would be to define the partial loss as min minyeC(p) LFY. (0: v), that is, the infimum of smoothed losses instead of the smoothed infimum loss T,y,c (0:p) p) defined above. However, there no direct method minimizing the smoothed loss with respect to y E C(p). Note that have relationship between the two losses:
Proposition Letting TAY, (0:y) = ([ + e7;y)) and Toy, as in Definiton 8, we have
LFY, min LEY.s (0;y).
yEC(p)
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 done naturally by transposing notations and taking C Cx. Σ, 3² A p Ma. C(p) = C. (M2), and y" (0:p) AE; Mp). In this setting the perturbed partial Fenchel-Young loss satisfies
V2LFY. (2; At) (2;
7

---

We learn representations of data that fit with clustering information (either complete or partial). As described above, we consider settings with n elements described by their features X I.....,X Xx, in some feature space X, and Mn some clustering information. Our pipeline to learn representations includes the following steps (see Figure 1)
i) Embed each X; in 1² with a parameterized model = (w, E R with weights w E IRP.
ii) Construct a similarity matrix from these embeddings, c.g. Sw.yj = - - (Xx)
ii) Stochastic optimization of the expected loss of & Ma,a), using mini-batches of X.
Details. To be more specific on each of these phases: i) We embed each X; individually with model ww. using in our application neural networks and a linear model. This allows us to use learning through clustering which we have no clustering information, or for use in other downstream tasks. ii) We focus on cases where the similarity matrix 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 E. has derivatives that can therefore be automatically computed by backpropagation, as it contains only conventional opereations (at least when .. is itself a conventional model, such commonly used neural networks). tii) We use our proposed smoothed partial Fenchel- Young (Section 3.2) as the objective to minimize between the partial information Ms and Σ. The full-batch version would to minimize IeF,,c(En, (Em. Ma) as function of the parameters E RRP of the model. We focus instead on mini-batch formulation for two reasons: first, stochastic optimization with mini-batches is 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 6, by
way learn representations, and to apply this model to other elements, for
LL (Ew. bi Ma,s) = d Ew VIIFY.:<EW,.y Ma,0) = Ew (Ak (Ew; Ma..))
The simplicity of these gradients 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 doubly stochastic scheme for the loss optimization.
4 Experiments
We apply our framework for learning through clustering both a supervised and a semi supervupe setting, as illustrated in Figure 1. Formally, for large training datasets of size we either have access full cluster connectivity matrix Mn or partial one (typically built by using partial label information, see below). We use this clustering information Mo. from which mini-batches can be
Kruskal's
Kmeans
MeanShift
GM
SMITER CaE
TM L., END
16
Figure 3: (left) Illustration of clustering methods on toy-data sets. (right) Using cluster information to learn linear de-noising (bottom) of a noised signal (top).
8

---

0.7
0.10
Classiticttion artor
8.08
Cludering
0.4
02
0.04
0.2
0.02
2000
5000
1000
2000
3000
4000
5000
Number samples
Number of labelled samples
0.83
Cluster tng error
Clessification costiitttiionn errer
6.0
Number of labelled samples
Number labelled samples
Figure 4: Semi-supervised learning: Performance of CNN trained partially labeled MNIST data (top row), and ResNet trained on partially labeled Cifar- 10 data (bottom row). Our trained model (full line) is evaluated clustering (left column) and compared to model entirely trained on the classification task (dashed line). Both models are evaluated on a down-stream classification problem i.e. transfer learning, via linear probe (right column).
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 test dataset: first, by evaluating the clustering accuracy (i.e. proportion of correct coefficients in the predicted cluster connectivity matrix), and second by training shallow model classification task (using clusters as classes) on 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.c. single linkage), with Kruskal's algorithm is efficient on some standard synthetic examples, even when they are not linearly separable (compared k- k-Means, mean-shift, Gaussian mixture-model, We also show that our method allows to perform supervised learning based on cluster information, in linear setting. For X consisting of 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; mixes the original clusters. We learn linear de-noising transformation X0, with 0 E R4x22 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 cluster virtually all of some classical datasets. We train CNN (LeNet-5 LeCun et al. (1998)) on mini-batches of size 64 using the partial Penchel-Young loss learn clustering, with batch-wise clustering precision of 0.99 for MNIST and 0.96 on Fashion MNIST. Using the same approach, we trained ResNet (He et al., 2016) on CIFAR-10 10 (with some minor modifications to kernel size and stride), achieving 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.

---

Figure 5: (left) t-SNE of learnt embeddings for the CNN trained on MNIST with ky withheld classes (highlighted). (right) t-SNE of learnt embeddings for the ResNet trained CIFAR-10.
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 al., 2010) and Cifar-10 (Krizhevsky et al., 2009) datasets, where we independently vary both the total number of labelled examples ne, as well as the number of withheld classes for which no labelled examples are present in the training set ki, E {0, 3, 6). For MNIST we consider data sets with E {100, 250, 500, 1000, 2000, 5000) labelled examples, and for Cifar- we consider nig (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 mini-batches of the same size), for each of the models, and data sets, corresponding to the choice of n, and hn- We compare the performance of each model baseline, consisting of the exact same architecture (plus linear head mapping the embedding to logits), trained to minimize the cross entropy loss.
In addition, we evaluate all models on down-stream (transfer-leamiing) classification task, by learning linear probe on top of the frozen model, trained on hold-out data with 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 representation where class semantics are easily recoverable from the local topology. On Cifar- 10, our proposed approach strikingly achieves a lower clustering eтrоr in the most challenging setting (ng 1000 labelled examples and l 6 withheld classes) than the classification-base baseline 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 lower down-stream classification error, with the gap being most pronounced when few labelled examples are available.
Moreover, besides this pronounced improvement 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.
10

---

Acknowledgements
We thank Lean-Philippe Vert for discussions relating to this work, Simon Legrand for conversations relating to the CLEPS computing cluster (Cluster pour l'Expérimentation et le Prototypage Scien- tifique), and Matthieu Blondel for providing references regarding learning via Kirchhoff's 's theorem. We also thank the reviewers and Kyle Heuton, for providing helpful corrections for the camera-reat paper. We also acknowledge support from the French government under the management of the Agence Nationale de la Recherche as part of the "Investissements d' t'avenir" program, reference ANRR1--3IA-000 (PRAIRIE 3IA Institute), as well as from the European Research Council (grant SEQUOIA 724063), and monetary don from Google.
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:9508- 9519.
Blondel, M., Berthet, Q., Cuturi, M., Frostig, R., Hoyer, S., -liiarrss-óóez. F., Pedregosa, F., and Vert, J.-P. (2022). Efficient and modular implicit differentiation. Advances in Neural Information Processing Systems, -22-22.......
Blondel, M., Martins, A. F., and Niculae, V. (2020a). Leaming with Fenchel- losses. The Journal of Machine Learning Research, 1144)2131--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 PythonnNumPyy programs.
Bubeck, S., Meilá, 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 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-711 -22

---

Cordonnier, J.-B., Mahendran, A., Dosovitskiy, A.. Weissenbomn 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-2366
Cormen, T. H., Leiserson, C. E., Rivest, R. L.,, and Stein, C. (2022). Introduction 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:
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 fralk-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-Arold,, G., and Vert, J.-P. (2019). Differentiable deep clustering with cluster size constraints. arXiv preprint arXiv: 191.:999..
Gower, J. C. and Ross, G. (1969). Minimum spanning trees and single linkage cluster analysis. Journal of the Royal Statistical Society: Series C (Applied Statistics), 18(1):54 64.
Gumbel, (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, 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, 135331133--137.
Jang, E., Gu, S., and Poole, B. (2017). Categorical reparameterization with gumbel-softtax. In International Conference 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, the IBM Thomas J. Watson Research Center, Yorktown Heights, New York, USA, pages 85-103.
Kingma, D. P. and Ba, (2015). Adam: A method for stochastic optimization. In Bengio, Y. and LeCun, Y., editors, 3rd International Conference on Learning Representations, ICLR 2015, Son Diego, CA. USA. May 7-9, 2015, Conference Track Proceedings.
12

---

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, 701):40-50.
Kumar, A., Brazil, G., and Liu, X. (2021). Groomed- 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 14:20398-20409..
LeCun, Y., Bottou, L., Bengio, Y., and Haffner, P. (1998). Gradient-base learning applied to document recognition. Proceedings of the IEEE, (6(11))2278-23333
LeCun, Y., Cortes, C., and Burges, C. (2010). Mnist handwritten digit database. ATT Labs [Online]. 2.
Available: http://yann.lecue www.::ewww....cc/eeeeeee
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 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, 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: 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 combinatorial spaces. Advances in Neural Information Processing Systems, 34: 10999-11000..
Vlastelica, M., Paulus, A., Musil, V., Martius, G., and Rolínek, M. (2019). Differentiation of blackbox combinatorial solvers. arXiv preprint arXiv: 19/2.02/75
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- k-means-friendly spaces: Simultaneous deep learning and clustering. In International Conference on Machine Learning, pages 3861-3870. PMLR.
7migrod, 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.
13

---

Broader impact
This submission focuses on foundational and exploratory work, with application 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.
Reproducibility and Licensing
Our experiments use data sets that are already written for this project and used to implement experiments is available at:
and cited in the references. All the code
attpss:/githii
and distributed under a Apache 2.0. license.
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, suitable for training deep learning models on large data sets. For example, standard classification models) for all experiments detailed our submission. Other methods that focus on cluster assignment, rather than pairwise clustering, often have X k parametrization, whose efficiency will depend on the comparison between n (batch size) and l: (number of clusters).
batch size of 64 (similar to that for
At inference time this is not the case, since gradients need not be back- propapppaa hence, any imple- mentation of Kruskal's algorithm can be used such as the union-find implementation. Faster GPU compatible implementations of Kruskal's ss algorithm for training are certainly possible improvements could be made using binary heap or a GPU parallelized merge sort. In many learning experiments the batch sizes 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 batch), it would be necessary improve upon the current implementation. This is mainly an engincering task relating to structursss 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 greedy lingkee 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 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
Loy(O:))=
min LFr(0:y). yEC(p)
We expand this
(F: min (F(0) () F(B) max (y. 8). yEe(p)
yEC(p)
As required. this implies. following the definitions given Lov(O:p) F(0) F(0:p), where F(O:p) max (v. 0) yEC(p)
2) By linearity of derivatives and 1) above, we have
voIry((;:)) = voF(0) -.F(()-vFF((;p)) y (0) - 3² (0:p), where y²(0:p = argmax(y. 0).
vEC(p)
14

---

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
(8: VoE + =Z) F(B - =2: = (OF VeF.(8:p) =0) (0:
using the definition
V2 (0:p) = Elargmax 0 + =Z,]
vEC(p)
Proof of Proposition 2. By Jensen's 's inequality and the definition of the Fenchel- Fenchel-Young loss:
(0: p) E[ min (F eZ: vEC(p)
s min E[ - =ZZ; v)l vEC(p)
= min vEC(p)
(U) (0, v)
s min LEY,e (0; v): SEC(p)
Algorithms for Spanning Forests
As mentioned in Section 2, both A (E) and ME) are calculated using Kruskal's algorithm (Kruskal, 1956). Our implementation of Kruskal's 's is tailored to use: we first initialize both A (E) and Mt (E) as the identity matrix, and then sort the upper triangular entries of Σ. We build the maximum- weight spanning forest in the usual greedy manner, using AE(E) keep track of edges in the forest and Mc (E) 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 & 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 single loop, with each step of the loop con- sisting only matrix multiplications. For this reason is fully compatible with auto-differentiation engines, such as JAX (Bradbury et al., 2018), 8), 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 's algorithm, which used a disjoint union-fing 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, make the modifications below to our implementation of Kruskal's 's under the assumption that M1 represents valid clustering information (i.e. with contradiction):
1. Regularization (Optional): It possible 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 large value to all entries of Mig where (Ma), = 1, and subtract this same value from all entries of Mi where (Ma)) = 1. Entries Σij corresponding to where (Ma) = 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. most cases, i.e. when all clusters are represented information, such as when Ω= [m] x [n] (full information). this is not required to solve the constrained linear program, but we have found that this regularization was helpful in practice.
the partial
15

---

2. Constraint enforcement: We ensure that adding an edge does not violate the constraint matrix. In other words, when considering adding the edge (i.j) the existing forest, check that none of the points in the connected component of are forbidden from joining any of the points in the connected component This is implemented using further matrix multiplications and done alongside the existing check for eycles. The exact implementation is detailed in our code base.
10 Existing Literature Differentiable Clustering
As discussed in Section I, there exists many approaches which use clustering during gradient based learning, but these approaches typically 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- pack-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 fully-connected deep network. Our method differs from this approach as allows for end-to-ent online learning.
Genevay al. (2019) cast k-means as an optimal transport problem, and uses entropic regularization for smoothing. Reported 85% accuracy on MNIST and 25% accuracy on CIFAR-10 with a CNN.
11 Additional Experimental Information
Figure 6:
Leftmost figure: M signal+
Center Figure: (2)... Rightmost Figure: 0 after training.
We provide the details of the synthetic denoising experiment depicted in Figure and described in Section 4.1. We create the signal data äignal E R2 by sampling iid. from four isotropic Gaussians (15 points coming from each of the Gaussians) each having 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:
0.97627008
4.30378733
2. 2.05526752
0.89766366
- 1.52690401
2.91788226
1.24825577
7.83546002)
Let Esignal be the pairwise euclidean similarity matrix corresponding X Signaa and furthermore let Magnal := Mi (Esigna)) be the equivalence matrix corresponding to the signal (Msignal will be the target equivalence matrix to learn..
We append an additional two 'noise dimensions' to Xsignal in order to form the train data X g60x4 where the noise entries were sampled iid from continuous unit uniform distribution. Similarly letting be the pairwise euclidean similarity matrix corresponding to X, we calculate MI) & Msignal- Both the matrices Maggall and (2)) are depicted in Figure 6 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 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 linear transformation of the data that recovers Magnal i.e. denoising, by minimizing the partial loss. There are multiple solutions to this problem, the most obvious being transformation that removes the last two noise columns from X:
16

---

0
0
0"
for which xO-
X
sigual
0
For any E Ixxx we define Eg to be the pairwise similarity matrix corresponding to X8, and (Eg) to its corresponding equivalence matrix. Then the problem can be summarized as:
min TF M Msignal). 0ER4X2
(1)
We initialized 9 from a standard Normal distribution, and minimized the partial loss via stochastic gradient descent, with learning rate of 0.01 and batch size 32.
For perturbations, we took ş = 0.1 and 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 indeed learning to remove the noise dimensions (see Figure 6).
1000, where denotes the noise amplitude randomized
11.1 Supervised Differentiable Clustering
kernel
ReLU & Aug Poot: kindon 2x2, stitides 2x2
features
kernel
RCLU & Aug Pool: windon 2x2, strides 2x2
Dense: 256
RetU
Dense:
Figure 7: (left) Architecture of the CNN, (middle) t-SNE visualization of train data embeddings,
(right) (SNE visualization of validation data embeddings.
As mentioned in Section 4.1, our method able to cluster classical data sets such as MNIST and Fashion MNIST. We trained CNN with the LeNet-5 architecture LeCun et al. (1998) using our proposed partial loss as the objective function. The exact details the CNN architecture are depicted in Figure 7. For this experiment and all experiments described below, we trained on single Nvidia V100 GPU; training the CNN with our proposed pipeline took < 15 minutes.
The model was trained for 30k gradient steps on mini-batches of size 64. We used the Adam optimizer (Kingma and Ba, 2015) with learning rate 7 = 3 10-4, momentum parameters (51, 32) (0.9, 0.999), and an E2 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 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-wit wise with batch-size 64. As stated in Section 4.1, we attained 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 It can be seen that the model has learnt distinct cluster for each of the ten classes.
In similar fashion, we trained 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 that the CNN, except the model was trained for 75k steps (with early stopping), and used the standard data augmentation methods for Cifar-10, namely:
17

---

combination of four-pixt sour-pieel padding, random flip followed by random crop. As mentioned in Section 4.1, the model achieved batch-wise clustering test precision 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 searce, even 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 Ma this is owing to the fact that the prediction of class for single point depends 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:
Train 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 check that our model has leveraged information from the unlabeled data and that our partial loss is indeed leading to good clustering performance.
Semi-Supervised Clustering:
2. Downstream Classification: Assess the trained model's capacity to serve as backbone in a downstream classification task (transfer learning), where its weights are frozen and 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 k 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 downstream classification problem.
From the remaining 59k images, we select a labeled train set of ny points (detailed in Section 4.2). Our experiments also vary lo E {0, 3, 6], the number of digits to withhold all labels from. For example, if kw = 3, then the labels for the images corresponding to digits {0, 1, will never appear in the labeled train data.
11.2.2 Semi-Supervised Clustering Task
For each of the choices of rig and kig we train the architectures described in Section 11.1 using the following two approaches:
1. Ours: The model is trained on mini-batches, where half the batch is labeled data and half the batch unlabeled data (i.e. semi-supervised learning regime), minimize the partial loss.
2. 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 labeled points, minimizing the cross- entropy loss. The training regime is fully-supervised learning (classification). The baseline backbone refers 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 E [5] (and with three random seeds s E [3] for Cifar), in order to report population statistics on the clustering error.
18

---

11.2.3 Transfer-Learning: Downstream Classification
In this task both models are frozen, and their utility foundational backbone is assessed on 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 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 the backbone corresponding to our approach. This highlights how learning 'cluster-able embeddings and leveraging unlabeled data can be desirable for transfer learning.
the case for
19


Reviews:
[{"review_id": "ht9MaDTFPT", "reviewer": "Reviewer_cJCp", "summary": "The 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.", "questions": "- 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.", "limitations": "The authors have listed the possible broader impact and limitations (e.g. batch size and current implementation efficiency) in the supplemental materials.", "rating": 5, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- 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.", "weaknesses": "- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you very much for your positive evaluation of our work, as well as the interesting points that you raise. We have taken this opportunity to address them in the following manner.\n\n### *Comparison with other clustering algorithms*\n\nWe have added a quantitative comparison for different clustering methods (k-means, Gaussian Mixture Model), both on the standard synthetic data showcased in Figure 3 of the submission, as well as on the CIFAR-10 dataset, using pre-trained features. \n\nThe results on the standard synthetic data showcased in Figure 3 of our original submission have been added to our manuscript - see Figure 1.(a) in the attached document.\n\nThe results for CIFAR-10 are as follows, they have been added to our manuscript\n\n+++++\n\nCIFAR-10 (ResNet50 for embeddings) [metric: L2 distance in coincidence matrix]:\n\nDifferentiable clustering: 0.097 (ours)\n\nClassification embeddings + k-means: 0.104 (non-differentiable)\n\nClassification embeddings + GMM: 0.107 (non-differentiable)\n\n+++++\n\nThis allows us to highlight that our method not only allows us to perform differentiable clustering, it allows us to learn good representations of the data (see following point)\n\n### *Performance on more complicated datasets*\n\nWe have added experiments on the CIFAR-10 dataset for both the fully supervised and semi-supervised settings (as well as a downstream classification transfer learning task, as in Figure 4 of the submission). \n\nFor supervised clustering, we obtain a clustering error of 6.7% on the test set, in line with our results on MNIST and Fashion-MNIST, we also provide in Figure 1.(b) the t-SNE of embeddings obtained in this manner in the feature space.\n\nFor semi-supervised clustering, the results are in Figure 2 in the attached document and have been included in the manuscript (analogous to Figure 4 on MNIST in our original submission). We observe analogous trends for clustering error to those detailed in Sections 4.2 (& 10.2) of the submission, with our model outperforming the baseline. The model attained a comparable classification error via a linear head (downstream transfer learning), to that of a baseline classification model for k=0, and for k=3, 6 outperformed the baseline.\n"}, {"review_id": "X0J893ZfQi", "reviewer": "Reviewer_mX78", "summary": "This 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.", "questions": "Please check the weakness.", "limitations": "The authors claim that the proposed method can be used to handle datasets with high noise, but this claim is not verified in the experiments.", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. 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.", "weaknesses": "My 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you very much for your positive evaluation of our work, as well as the interesting points that you raise. We have taken this opportunity to address them in the following manner.\n\n### *Scaling to more complicated datasets*\n\nAs our method can use mini-batches, it is suitable for training deep learning models on large datasets (we use a batch size of 64 for all experiments detailed in our submission). \n\nWe have added experiments on the CIFAR-10 dataset for both the fully supervised and semi-supervised settings (as well as a downstream classification transfer learning task, as in Figure 4 of the submission). \n\nFor supervised clustering, we obtain a clustering error of 6.7% on the test set, in line with our results on MNIST and Fashion-MNIST, we also provide in Figure 1.(b) the t-SNE of embeddings obtained in this manner in the feature space.\n\nFor semi-supervised clustering, the results are in Figure 2 in the attached document and have been included in the manuscript (analogous to Figure 4 on MNIST in our original submission). We observe analogous trends for clustering error to those detailed in Sections 4.2 (& 10.2) of the submission, with our model outperforming the baseline. The model attained a comparable classification error via a linear head (downstream transfer learning), to that of a baseline classification model for k=0, and for k=3, 6 outperformed the baseline. \n\n\n### *Comparison with other differentiable clustering methods*\n\nWe have added a comparison between our method’s results on MNIST and CIFAR-10 and those reported by [Yang et al. 2017] and [Genevay et al. 2019], because they were the most comparable to our own experimental evaluations. As a reminder, using our methodology, we obtain a clustering error of 1.0% on MNIST and 6.7% on CIFAR-10 (and 4.0% on Fashion-MNIST).\nA remark describing this comparison has been included in the manuscript.\n\n+ [Yang et al. 20177] - Towards K-means-friendly Spaces: Simultaneous Deep Learning and Clustering, ICML 2017\n\nUses an 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. \n\n+ [Genevay et al. 2019] -  ​​Differentiable Deep Clustering with Cluster Size Constraints, 2019\n\nCasts 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.\n"}, {"review_id": "DNEV2mrsaR", "reviewer": "Reviewer_ukVi", "summary": "The paper presents a novel method for differentiable clustering in combination with a novel loss derived from Fenchel-Young losses.", "questions": "see above.", "limitations": "The experimental evaluation is limited.", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The 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.", "weaknesses": "It 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you very much for your positive evaluation of our work, as well as the interesting points that you raise. We have taken this opportunity to address them in the following manner.\n\n ### *Scaling to more complicated datasets*\n\nThank you very much for your suggestion, and the indication that you would raise your score. As suggested, we have added experiments on the CIFAR-10 dataset for both the fully supervised and semi-supervised settings (as well as a downstream classification transfer learning task, as in Figure 4 of the submission). \n\nFor supervised clustering, we obtain a clustering error of 6.7% on the test set, in line with our results on MNIST and Fashion-MNIST, we also provide in Figure 1.(b) the t-SNE of embeddings obtained in this manner in the feature space.\n\nFor semi-supervised clustering, the results are in Figure 2 in the attached document and have been included in the manuscript (analogous to Figure 4 on MNIST in our original submission). We observe analogous trends for clustering error to those detailed in Sections 4.2 (& 10.2) of the submission, with our model outperforming the baseline. The model attained a comparable classification error via a linear head (downstream transfer learning), to that of a baseline classification model for k=0, and for k=3, 6 outperformed the baseline. \n\n### *On learning representations as the main aspect of being able to transmit gradients backward through a clustering operation*\n\nThank you very much for bringing this up. Indeed, this is a very interesting way to present our work, and an important point, that we had not clearly expressed. We have added a remark highlighting this point in the introduction.\n\n### *Notation and presentation* \n\nWe have made all of the suggested terminology and presentation changes.\n\n### *Suggested reference*\n\nThank you for this very interesting reference, we have added a citation to it, as well as discussion of this work.\n"}, {"review_id": "gDaBpfUw7F", "reviewer": "Reviewer_sM5n", "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.", "questions": "See Weaknesses", "limitations": "See Weaknesses", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "This paper introduces a differentiable clustering method based on linear programs for minimum-weight spanning forests. \nTheir 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.\n\nThe 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, \n\n- 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. \n\n- Comparison with other methods is lacking: maybe some SOTA nondifferentiable clustering methods \n\n- 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?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you very much for your positive evaluation of our work, as well as the interesting points that you raise. We have taken this opportunity to address them in the following manner.\n\n### *Elaborating on the setting*\n\nWe have added a remark clarifying how $M_\\Omega$ is given, from label information. To summarize, for all pairs of elements $(i, j)$ who both have label information, $M_{\\Omega, ij}$ is equal to $1$ if they have the same label (must-link) and $0$ if they have a different one (must-not-link). The coefficient $M_{\\Omega, ij}$ is ‘$*$’ (no information) if one of the two labels is not known.\n\n### *Comparison with nondifferentiable clustering methods*\n\nWe have added a quantitative comparison for different clustering methods (k-means, Gaussian Mixture Model), both on the standard synthetic data showcased in Figure 3 of the submission, as well as on the CIFAR-10 dataset, using pre-trained features. \n\nThe results on the standard synthetic data showcased in Figure 3 have been added to our manuscript - see Figure 1.(a) in the attached document.\n\nThe results for CIFAR-10 are as follows, they have been added to our manuscript\n\n+++++\n\nCIFAR-10 (ResNet50 for embeddings) [metric: L2 distance in coincidence matrix]:\n\nDifferentiable clustering: 0.097 (ours)\n\nClassification embeddings + k-means: 0.104 (non-differentiable)\n\nClassification embeddings + GMM: 0.107 (non-differentiable)\n\n+++++\n\nThis allows us to highlight that our method not only allows us to perform differentiable clustering, it allows us to learn good representations of the data.\n\n### *Parametrization of our method*\n\nOur method indeed scales quadratically in $n$, similarly to attention mechanisms. 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.\n\nOther methods that focus on cluster assignment, rather than pairwise clustering, often have an $n \\times k$ parametrization, which can be smaller or larger, depending on the comparison between $n$ (batch size) and $k$ (number of clusters). We have added a remark to clarify both of these points.\n"}]

Author rebuttal:
We would like to thank the reviewers and the area chair for their overall positive evaluation of our work. 

We are thankful for the high-quality, informative reviews, and their constructive suggestions. We have made the following changes, further described in individual replies to reviewers:

- We have added an additional clustering experiment on the CIFAR 10 dataset, in both the fully-supervised and semi-supervised settings. The results are in the attached document (Figure 2) and have been included in our manuscript.

- We have added a quantitative comparison for the clustering methods contained in Figure 3 of the submission, as well as evaluating different clustering methods on features of the pre-trained CIFAR-10 features. We detail these results in replies below and they have been included in our manuscript.

- We have made suggested changes: clarifying some points, and improving our terminology and presentation (see replies below for details).

Meta-review: AC agreed with the reviewers that the approach is novel and interesting. However, AC remains concerned about the paper has only MNIST experiments and then in the rebuttal added cifar10. CIFAR10 is not a convincing dataset either. Ac recommends larger datasets to affirm the validity of the approach in the camera ready but recommends accepting the paper due to the reviewers' alignment to accept the paper.
Decision: Accept (poster)