SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Lawrence Stewart*
ENS & INRIA
Francis Bach
Felipe Llinares-Lopez
Quentin Berthet
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.[2].
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 intra-cluster 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 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
*All correspondence should be addressed to lawrence.stewart@ens.fr.
$^{2}$ Code base: https://github.com/LawrenceMMStewart/DiffClust_NeurIPS2023
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-Lopez 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; Genevaay 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 information—is 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, \dots, n\}$ . We consider all graphs to be undirected. For the complete graph with  $n$  vertices  $K_{n}$  over nodes  $[n]$ , we denote by  $\mathcal{T}$  the set of spanning trees on  $K_{n}$ , i.e., subgraphs with no cycles and one connected component. For any positive integer  $k \leq n$  we also denote by  $\mathcal{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  $\mathcal{T}$  and  $\mathcal{C}_k$  for the set of adjacency matrices of these graphs. For two nodes  $i$  and  $j$  in a general graph, we write  $i \sim j$  if the two nodes are in the same connected component. We denote by  $S_{n}$  the set of  $n \times 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}^{*}(\cdot)$  taking an  $n \times 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,\varepsilon}^{*}(\cdot)$ . 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  $\mathbf{X} = (X_1, \dots, X_n)$  in some feature space  $\mathcal{X}$ , and ground-truth clustering data as an  $n \times n$  matrix  $M_{\Omega}$  (with either complete or partial information, see Definition 4). A parameterized model produces a similarity matrix  $\Sigma = \Sigma_w(\mathbf{X})$ -e.g., based on pairwise square distances between embeddings  $\Phi_w(X_i)$  for some model  $\Phi_w$ . Minimizing in  $w$  a loss so that  $M_{k}^{*}(\Sigma = \Sigma_w(\mathbf{X})) \approx M_{\Omega}$  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.
Figure 1: Our pipeline, in the semi-supervised learning setting: data points are embedded by a parameterized model, which produces a 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 are used, respectively for prediction and gradient computations.
2.1 Clustering with  $k$ -spanning forests
In this work, the core operation is to cluster  $n$  elements, using a similarity matrix  $\Sigma \in S_{n}$ . Informally, pairs of elements  $(i,j)$  with high similarity  $\Sigma_{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  $\pi :[n]\to [k]$  be a clustering function, assigning elements to one of  $k$  clusters. We represent it with an  $n\times n$  binary matrix  $M$  (the cluster connectivity):
$$
M _ {i j} = 1 \quad i f a n d o n l y i f \quad \pi (i) = \pi (j).
$$
We denote by  $\mathcal{B}_k\subseteq \{0,1\}^{n\times n}$  the set of binary cluster connectivity matrices with  $k$  clusters.
Using this definition, we define an operation  $M_{k}^{*}(\cdot)$  mapping a similarity matrix  $\Sigma$  to a clustering, in the form of such a membership matrix. Up to a permutation (i.e., naming the clusters),  $M$  allows to recover  $\pi$  entirely. It is based on a maximum spanning forest primitive  $A_{k}^{*}(\cdot)$ , 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  $\Sigma$  be an  $n \times n$  similarity matrix. We denote by  $A_k^*(\Sigma)$  the adjacency matrix of the  $k$ -spanning forest with maximum similarity, defined as
$$
A_{k}^{*}(\Sigma) = \operatorname *{argmax}_{A\in \mathcal{C}_{k}}\langle A,\Sigma \rangle .
$$
This defines a mapping  $A_k^* : \mathcal{S}_n \to \mathcal{C}_k$ . We denote by  $F_k$  the value of this maximum, i.e.,
$$
F_{k}(\Sigma) = \max_{A\in \mathcal{C}_{k}}\langle A,\Sigma \rangle .
$$
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.,
$$
M _ {i j} = 1 \quad i f a n d o n l y i f \quad i \sim j.
$$
Given an  $n \times n$  similarity matrix  $\Sigma$ , we denote by  $M_k^*(\Sigma) \in \mathcal{B}_k$  the clustering induced by the maximum  $k$ -spanning forest, defined by
$$
M _ {k} ^ {*} (\Sigma) = M ^ {*} (A _ {k} ^ {*} (\Sigma)).
$$
This defines a mapping  $M_k^*:\mathcal{S}_n\to \mathcal{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}^{*}(\Sigma)$  and  $M_{k}^{*}(\Sigma)$  are actually obtained by this algorithm. The mapping  $M^{*}$  is of course many-to-one, and yields a partition of  $\mathcal{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  $\nabla_{\Sigma}F_{k}(\Sigma) = A_{k}^{*}(\Sigma)$ .
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).
Figure 2: Method illustration, for  $k = 2$ : (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  $M_{\Omega}$ , defined as follows.
Definition 4 (Partial cluster coincidence matrix). Let  $M$  be a cluster connectivity matrix, and  $\Omega$  a subset of  $[n] \times [n]$ , representing the set of observations. We denote by  $M_{\Omega}$  the  $n \times n$  matrix
$$
M _ {\Omega , i j} = M _ {i j} \quad \text {i f} (i, j) \in \Omega , \quad M _ {\Omega , i j} = * \quad \text {o t h e r w i s e}.
$$
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 \subseteq [n]$  of the data has label information. In this case for any  $i, j \in S$ ,  $M_{\Omega, ij} = 1$  if  $i$  and  $j$  share the same label, otherwise  $M_{\Omega, 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  $\Sigma$  be an  $n \times n$  similarity matrix. We denote by  $A_k^*(\Sigma; M_{\Omega})$  the adjacency matrix of the  $k$ -spanning forest with maximum similarity, constrained to satisfy the connectivity constraints in  $M_{\Omega}$ . It is defined as
$$
A _ {k} ^ {*} (\Sigma ; M _ {\Omega}) = \operatorname * {a r g m a x} _ {A \in \mathcal {C} _ {k} (M _ {\Omega})} \langle A, \Sigma \rangle ,
$$
where for any partial clustering matrix  $M_{\Omega}$ ,  $\mathcal{C}_k(M_\Omega)$  is the set of  $k$ -spanning forests whose clusters agree with  $M_{\Omega}$ , i.e.,
$$
\mathcal {C} _ {k} (M _ {\Omega}) = \left\{A \in \mathcal {C} _ {k}: M ^ {*} (A) _ {i j} = M _ {\Omega , i j} \forall (i, j) \in \Omega \right\}.
$$
For any partial connectivity matrix  $M_{\Omega}$ , this defines another mapping  $A^{*}(\cdot ;M_{\Omega}):S_{n}\to \mathcal{C}_{k}$
We denote by  $F_{k}(\Sigma ;M_{\Omega})$  the value of this maximum, i.e.,
$$
F _ {k} (\Sigma ; M _ {\Omega}) = \max  _ {A \in \mathcal {C} _ {k} (M _ {\Omega})} \langle A, \Sigma \rangle .
$$
Similarly we define  $M_k^*(\Sigma; M_{\Omega}) = M^*(A_k^*(\Sigma; M_{\Omega}))$ .
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  $\Omega$  is, these constraints can be more or less subtle to satisfy. For certain sets of constraints  $\Omega$ , when the matroid structure is preserved, we can obtain  $A_{k}^{*}(\Sigma; M_{\Omega})$  by the usual greedy algorithm, by additionally checking that any new edge does not violate the constraints defined by  $\Omega$ . This is for example the case when  $\Omega$  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 \sim \mu$  from a distribution with positive, differentiable density over  $S_{n}$ , and  $\varepsilon > 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  $\Sigma$ , we have
$$
A _ {k, \varepsilon} ^ {*} (\Sigma) = \mathbf {E} [ A _ {k} ^ {*} (\Sigma + \varepsilon Z) ] = \mathbf {E} \left[ \operatorname * {a r g m a x} _ {A \in \mathcal {C} _ {k}} \langle A, \Sigma + \varepsilon Z \rangle \right], \quad F _ {k, \varepsilon} (\Sigma) = \mathbf {E} \left[ F _ {k} (\Sigma + \varepsilon Z) \right].
$$
We define analogously  $A_{k,\varepsilon}^{*}(\Sigma; M_{\Omega}) = \mathbf{E}[A_{k}^{*}(\Sigma + \varepsilon Z; M_{\Omega})]$  and  $F_{k,\varepsilon}(\Sigma; M_{\Omega}) = \mathbf{E}[F_{k}(\Sigma + \varepsilon Z; M_{\Omega})]$ , as well as clustering  $M_{k,\varepsilon}^{*}(\Sigma) = \mathbf{E}[M_{k}^{*}(\Sigma + \varepsilon Z)]$  and  $M_{k,\varepsilon}^{*}(\Sigma; M_{\Omega}) = \mathbf{E}[M_{k}^{*}(\Sigma + \varepsilon Z; M_{\Omega})]$ .
We note that this defines operations  $A_{k,\varepsilon}^{*}(\cdot)$  and  $A_{k,\varepsilon}^{*}(\cdot; M_{\Omega})$  mapping  $\Sigma \in S_n$  to  $\mathrm{cvx}(\mathcal{C}_k)$ , the convex hull of  $\mathcal{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^*(\Sigma + \varepsilon 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,\varepsilon}^{*}(\cdot)$  is a differentiable operator from  $S_{n}$  to  $\operatorname{cvx}(\mathcal{B}_k)$ , it is possible to use any loss function on  $\operatorname{cvx}(\mathcal{B}_k)$  to design a loss based on  $\Sigma$  and some ground-truth clustering information  $M_{\Omega}$ , such as  $L(\Sigma; M_{\Omega}) = \| M_{k,\varepsilon}^{*}(\Sigma) - M_{\Omega} \|_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