Differentiable Clustering and Partial Fenchel-Young Losses

Lawrence Stewart 1 Francis Bach 1 Felipe Llinares-López 2 Quentin Berthet 2

Abstract

the clusters), M allows to recover π entirely. It is based ona maximum spanning forest primitive A∗k (·), and both aredefined below.

We introduce a differentiable clustering method
based on stochastic perturbations of minimumweight spanning forests. This allows us to include
clustering in end-to-end trainable pipelines, withefficient 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.

Definition 1.2 (Maximum k-spanning forest). Let Σ be ann × n similarity matrix. We denote by A∗k (Σ) the adjacency
matrix of the k-spanning forest with maximum similarity,
defined asA∗k (Σ) = argmax⟨A, Σ⟩.
A∈Ck

This defines a mapping A∗k : Sn → Ck . We denote by Fk
the value of this maximum, i.e.,
Fk (Σ) = max ⟨A, Σ⟩.
A∈Ck

Definition 1.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 ofA, i.e.,Mij = 1 if and only if i ∼ j .

1. Differentiable clustering
1.1. Clustering with k-spanning forests
The complete set of notations used throughout this paper canbe found in Appendix A. In this work, the core operation
is to cluster n elements, using a similarity matrix Σ ∈ Sn .
Informally, pairs of elements (i, j) with high similarity Σijshould be more likely to be in the same cluster than thosewith low similarity. The clustering is represented in thefollowing manner.

Given an n × n similarity matrix Σ, we denote by Mk∗ (Σ) ∈Bk the clustering induced by the maximum k-spanning forest, defined byMk∗ (Σ) = M ∗ (A∗k (Σ)) .

Definition 1.1 (Cluster connectivity matrix). Let π : [n] →[k] be a partition function, assigning n elements to one of kclusters. We represent it with a n × n binary cluster connectivity matrix M (also known as the coincidence matrix)

This defines a mapping Mk∗ : Sn → Bk , our clustering
operator.
Remarks. Both A∗k (Σ) and Mk∗ (Σ) can be obtained usinga modified version of Kruskal’s algorithm (Kruskal, 1956).
This is equivalent to single linkage hierarchical agglomerative clustering with cluster merging carried out until kclusters remain (for more details see Appendix B). The mapping M ∗ is of course many-to-one, and yields a partition of
Ck into equivalence classes of k-forests that yield the sameclusters. Further illustrations are given in Appendix D.

Mij = 1 if and only if π(i) = π(j) .
We denote by Bk ⊆ {0, 1}n×n the set of binary cluster
connectivity matrices with k clusters.
Using this definition, we define an operation Mk∗ (·) mappinga similarity matrix Σ to a clustering, in the form of sucha connectivity matrix. Up to a permutation (i.e., naming*

Clustering with constraints. We introduce constrained
versions of the Mk∗ and A∗k in order to infer clusterings andspanning forests that are consistent with observed information e.g. weak supervision. We represent these enforced
connectivity values as a matrix MΩ defined as follows.

1

Equal contribution Ecole Normale Superieure and INRIA,
Paris. 2 Google DeepMind, Paris.. Correspondence to: Lawrence
Stewart <lawrence.stewart@ens.fr>.
Published at the Differentiable Almost Everything Workshop of the
40 th International Conference on Machine Learning, Honolulu,
Hawaii, USA. July 2023. Copyright 2023 by the author(s).

Definition 1.4 (Partial cluster connectivity matrix). Let M
be a cluster connectivity matrix, and Ω a subset of [n] × [n],1Differentiable Clustering and Partial Fenchel-Young Losses

We note that this defines operations A∗k,ε (·) and A∗k,ε (·, MΩ )mapping Σ ∈ Sn to cvx(Ck ), 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 copiesof A∗k (Σ + εZ (i) ).

representing the set of observations. We denote by MΩ then × n matrix
MΩ,ij = Mij

if (i, j) ∈ Ω ,MΩ,ij = ∗ otherwise .

Remarks. The symbol “∗” in this definition is only used
as a placeholder, an indicator of “no information”, and does
not have a mathematical use. MΩ,ij = 1 corresponds toa must-link constraint, whilst MΩ,ij = 0 corresponds to a
must-not-link constraint.
Definition 1.5 (Constrained maximum k-spanning forest).
Let Σ be an n × n similarity matrix. We denote byA∗k (Σ, MΩ ) the adjacency matrix of the k-spanning forest
with maximum similarity, constrained to satisfy the connectivity constraints in MΩ . It is defined as

This is particularly convenient, as it does not require to implement a different algorithm to compute the differentiableversion. Moreover, the use of parallelization in modern computing hardware makes the computational overhead almost
nonexistent.
Our method allows for clustering on learn-able representations, without breaking the differentiability of a pipeline.
These representations are informed by gradients transmitted∗(·) isbackward through the clustering algorithm. Since Mk,εa differentiable operator from Sn to cvx(Bk ), it is possible
to use any loss function on cvx(Bk ) to design a loss basedon Σ and some ground-truth clustering information MΩ ,∗such as L(Σ; MΩ ) = ∥Mk,ε(Σ) − MΩ ∥22 . This flexibility
is one of the advantages of our method.

A∗k (Σ, MΩ ) = argmax ⟨A, Σ⟩ ,A∈Ck (MΩ )where for any partial clustering matrix MΩ , Ck (MΩ ) is theset of k-spanning forests whose clusters agree with MΩ ,i.e.,
Ck (MΩ ) = {A ∈ Ck : M ∗ (A)ij = MΩ,ij ∀ (i, j) ∈ Ω} .

In the following section, we introduce a loss tailored tobe efficient to compute and performant in several learning
tasks, that we call a partial Fenchel-Young loss.

For any partial connectivity matrix MΩ , this defines anothermapping A∗ (·, MΩ ) : Sn → Ck .
We denote by Fk (Σ, MΩ ) the value of this maximum, i.e.,
Fk (Σ, MΩ ) =max

A∈Ck (MΩ )2. Learning with differentiable clustering⟨A, Σ⟩ .

2.1. Partial Fenchel-Young losses

1.2. Perturbed Clustering

The framework of Fenchel-Young losses (Blondel et al.,2020) can tackle complex structures which can be encoded
as an LP, such as spanning k-forests. However, in mostcases, clustering data can reasonably only be expected to
be present as a connectivity matrix, which unfortunately
does not take the form of an argmax of an LP. To address
this problem, we introduce the Partial Fenchel-Young loss,
which leverages the relation between Mk∗ and A∗k (the priorpartitions the latter into equivalence classes).

The clustering operations defined above are efficient andperform 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 endto-end differentiable pipelines, such as those used to trainmodels such as neural networks. To overcome this obstacle,
we use a perturbed proxy (Abernethy et al., 2016; Berthet
et al., 2020; Paulus et al., 2020; Struminsky et al., 2021).
In these definitions and the following, we consider Z ∼ µfrom a distribution with positive, differentiable density over
Sn , and ε > 0, following the methodology of Berthet et al.
(2020).
Definition 1.6. We define the perturbed maximum spanning forest as the expected maximum spanning forest understochastic perturbation on the inputs. Formally, for a similarity matrix Σ, we haveA∗k,ε (Σ) = E[A∗k (Σ + εZ)] = E argmax⟨A, Σ + εZ⟩ ,Definition 2.1 (Partial Fenchel-Young loss). Let F be aconvex function, C a convex set and LFY the associatedFenchel-Young loss. For every p ∈ P a convex constraint
subset C(p) ⊆ C, we define:
L̄FY (θ, p) = min LFY (θ; y) .
y∈C(p)

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

A∈Ck

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

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

Differentiable Clustering and Partial Fenchel-Young Losses

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

Diff. Clustering

L̄FY (θ, p) = F (θ) − F (θ; p) ,where F (θ; p) = max ⟨y, θ⟩,

Similarity
matrixClustering

Partialinformationy∈C(p)

Features Model

2. The gradient with respect to θ is given byEmbeddings
Prediction∇θ L̄FY (θ, p) = y ∗ (θ) − y ∗ (θ; p) ,Loss gradientFigure 1. Our pipeline, in a semi-supervised example: data points
are embedded by a parameterized model, which produces a similarity matrix. Partial clustering information may be available, inthe form of must-link or must-not-link constraints.

where y ∗ (θ; p) = argmax⟨y, θ⟩.
y∈C(p)

3. The perturbed partial Fenchel-Young loss given by
L̄FY,ε (θ; p) = E[L̄FY (θ + εZ; p)] satisfies∇θ L̄FY,ε (θ, p) = yε∗ (θ) − yε∗ (θ; p) ,3. Experiments

where yε∗ (θ; p) = E[argmax⟨y, θ + εZ⟩] .

We apply our framework for learning through clustering inboth a supervised and a semi-supervised setting, as illustrated in Figure 1. Formally, for large training data sets ofsize n, we either have access to a full cluster connectivity
matrix MΩ or a partial one (typically built by using partial
label information, see below).

y∈C(p)

Another possibility would be to define the partial loss as
miny∈C(p) LFY,ε (θ; y), that is, the infimum of smoothed
losses instead of the smoothed infimum loss LFY,ε (θ; p)defined above. However, there is no direct method to minimizing the smoothed loss with respect to y ∈ C(p).
Proposition 2.3. Letting LFY,ε (θ; y) = E[LFY,ε (θ +εZ; y)] and L̄FY,ε as in Definiton 2.1, we have

We use this clustering information MΩ , from which minibatches can be extracted, as supervision. We minimize ourpartial Fenchel-Young loss with respect to the weights ofan embedding model, and evaluate these embeddings in
two main manners on a test dataset: first, by evaluating theclustering accuracy (i.e. proportion of correct coefficientsin the predicted cluster connectivity matrix), and secondby training a shallow model on a classification task (using
clusters as classes) on a holdout set, evaluating it on a test
set.

L̄FY,ε (θ; p) ≤ min LFY,ε (θ; y) .
y∈C(p)The proofs of the above propositions are given in Appendix E.
2.2. Applications to differentiable clustering
We apply this framework, as detailed in the following section and in Section 3, to clustering. This is done naturally bytransposing notations and taking C = Ck , θ = Σ, y ∗ = A∗k ,p = MΩ , C(p) = Ck (MΩ ), and y ∗ (θ, p) = A∗k (Σ; MΩ ). In
this setting the perturbed partial Fenchel-Young loss satisfies

3.1. Supervised learningiii) Stochastic optimization of the expected loss of
L̄FY,ε (Σw,b , MΩ,b ), using mini-batches of X.

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

Further details on each of these steps is detailed in Appendix C.

We also show that our method is able to cluster virtually
all of some classical data sets. We train a CNN (LeNet-5∇Σ L̄FY,ε (Σ, MΩ ) = A∗k,ε (Σ) − A∗k,ε (Σ; MΩ ) .
We learn representations of a data that fit with clustering
information (either complete or partial). As described above,
we consider settings with n elements described by theirfeatures X = X1 , . . . , Xn in some feature space X , andMΩ some clustering information. Our pipeline to learn
representations includes the following steps (see Figure 1)i) Embed each Xi in Rd with a parameterized modelvi = Φw (Xi ) ∈ Rd , with weights w ∈ Rp .
ii) Construct a similarity matrix from these embeddings,
e.g. Σw,ij = −∥Φw (Xi ) − Φw (Xj )∥22 .

3

Differentiable Clustering and Partial Fenchel-Young Losses

Kruskal's

KmeansMeanShift

GM

Figure 2. Comparing clustering methods (left), and learning a linear representation with clustering information from noised signal (middle)
and after training (right).
0.7kw = 0kw = 3kw = 6

0.6

Clustering error

Classification error

0.5

0.08

0.4

0.06

0.3
0.2

0.04

0.1
0.0

kw = 0kw = 3kw = 6

0.10

0

1000

2000

3000

Number of labelled samples

4000

5000

0.02

0

1000

2000

3000

Number of labelled samples

4000

5000

Figure 3. Performance on semi-supervised learning task. Our trained model (full line) is evaluated on clustering (left) and classification
errors (middle), compared to a model entirely trained on the classification task (dashed line). The rightmost figure shows the t-SNE oflearnt embeddings for MNIST with kw = 3 withheld classes (highlighted).

(LeCun et al., 1998)) on mini-batches of size 64 using thepartial Fenchel-Young loss to learn a clustering, with a batchwise clustering precision of 0.99 for MNIST and 0.96 on
Fashion MNIST (Xiao et al., 2017). The experimental setup
and visualization of learnt clusters is detailed in Appendix F.

the classes present. We compare both metrics to the performance of a model composed of an identical CNN, with anadditional linear head, trained on the same data using thecross-entropy loss (see Appendix F for further details).
We observe that learning through clustering allows to find arepresentation where class semantics are easily recoverablefrom the local topology. Strikingly, our proposed approachachieves a lower clustering error in the most challenging
setting (nℓ = 100 labelled examples and kw = 6 withheldclasses) than the classification-based baseline in the mostlenient setting (nℓ = 5000 labelled examples and all classes
observed). Most importantly, this advantage is not limited
to clustering metrics: learning through clustering also leads
to a lower classification error for all conditions under study,
with the gap being most pronounced when few labelled
examples are available.

3.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 someclasses. To this end, we conduct a series of experimentson the MNIST dataset (LeCun et al., 2010) where we independently vary both the total number of labelled examples
nℓ = {100, 250, 500, 1000, 2500, 5000} as well as the number of withheld classes for which no labelled examples arepresent in the training set kw = {0, 3, 6}. We train the sameembedding model Φw as above (CNN LeNet-5), using ourpartial 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 & Ba, 2015).
The two metrics reported in Figure 3, both on a test set are
clustering error (evaluated on mini-batches of the same size)and classification error, evaluated by training an additional
linear layer (freezing our CNN) on hold-out data with all

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

4

Differentiable Clustering and Partial Fenchel-Young Losses

References

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, 34:
10999–11011, 2021.

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

Xiao, H., Rasul, K., and Vollgraf, R. Fashion-mnist: a
novel image dataset for benchmarking machine learningalgorithms, 2017.

Abernethy, J., Lee, C., and Tewari, A. Perturbation techniques in online learning and optimization. Perturbations,
Optimization, and Statistics, 233, 2016.
Berthet, Q., Blondel, M., Teboul, O., Cuturi, M., Vert, J.P., and Bach, F. Learning with differentiable pertubed
optimizers. Advances in Neural Information Processing
Systems, 33:9508–9519, 2020.
Blondel, M., Martins, A. F., and Niculae, V. Learning with
Fenchel-Young losses. The Journal of Machine Learning
Research, 21(1):1314–1382, 2020.
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. JAX: composable
transformations of Python+NumPy programs, 2018. URL
http://github.com/google/jax.
Defazio, A., Bach, F., and Lacoste-Julien, S. Saga: A fast incremental gradient method with support for non-strongly
convex composite objectives. Advances in Neural Information Processing Systems, 27, 2014.
Kingma, D. P. and Ba, J. Adam: A method for stochastic
optimization. In Bengio, Y. and LeCun, Y. (eds.), 3rdInternational Conference on Learning Representations,
ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Track Proceedings, 2015.
Kruskal, J. B. On the shortest spanning subtree of a graphand the traveling salesman problem. Proceedings of theAmerican Mathematical Society, 7(1):48–50, 1956.
LeCun, Y., Bottou, L., Bengio, Y., and Haffner, P. Gradientbased learning applied to document recognition. Proceedings of the IEEE, 86(11):2278–2324, 1998.
LeCun, Y., Cortes, C., and Burges, C. Mnist handwritten digit database. ATT Labs [Online]. Available:
http://yann.lecun.com/exdb/mnist, 2, 2010.
Paszke, A., Gross, S., Massa, F., Lerer, A., Bradbury, J.,
Chanan, G., Killeen, T., Lin, Z., Gimelshein, N., Antiga,L., et al. Pytorch: An imperative style, high-performance
deep learning library. Advances in Neural Information
Processing Systems, 32, 2019.
Paulus, M., Choi, D., Tarlow, D., Krause, A., and Maddison,
C. J. Gradient estimation with stochastic softmax tricks.
Advances in Neural Information Processing Systems, 33:
5691–5704, 2020.
5

Differentiable Clustering and Partial Fenchel-Young Losses

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

B. Algorithms for Spanning Forests
As mentioned in Section 1, both A∗k (Σ) and Mk∗ (Σ) are calculated using Kruskal’s algorithm (Kruskal, 1956). Our
implementation of Kruskal’s is tailored to our use: we first initialize both A∗k (Σ) and Mk∗ (Σ) as the identity matrix, andthen sort the upper triangular entries of Σ. We build the maximum-weight spanning forest in the usual greedy manner, usingA∗k (Σ) to keep track of edges in the forest and Mk∗ (Σ) to check if an edge can be added without creating a cycle, updating
both matrices at each step of the algorithm. Once the forest has k connected components, the algorithm terminates. This isdone by keeping track of the number of edges that have been added at any time.
We remark that our implementation takes the form as a single loop, with each step of the loop consisting only of matrix
multiplications. For this reason it is fully compatible with auto-differentiation engines, such as JAX (Bradbury et al., 2018),
Pytorch (Paszke et al., 2019) and TensorFlow (Abadi et al., 2016), and suitable for GPU/TPU acceleration. Therefore, ourimplementation differs from that of the standard Kruskal’s algorithm, which used a disjoint union-find data structure (andhence is not compatible with auto-differentiation frameworks). Finally, our implementation of Kruskal’s compatible withthe perturbations module of the JAXOPT library, so that this operation and the Fenchel-Young losses can be automaticallydifferentiated without requiring to be implemented.
B.1. Constrained Spanning Forests
To solve the constrained problem detailed in Section 1.5, we make the modifications below to our implementation of
Kruskal’s, under the assumption that MΩ represents valid clustering information (i.e. with no contradiction):
1. Regularization (Optional) : It is possible to bias the optimization problem over spanning forests to encourage
or discourage edges between some of the nodes, according to the clustering information. Before performing thesort on the upper-triangular of Σ, we add a large value to all entries of Σij where (MΩ )ij = 1, and subtract
this same value from all entries of Σij where (MΩ )ij = 1. Entries Σij corresponding to where (MΩ )ij = ⋆are left unchanged. This biasing ensures that any edge between points that are constrained to be in the samecluster will always be processed before unconstrained edges. Similarly, any edge between points that are
constrained to not be in the same cluster, will be processed after unconstrained edges. In most cases, i.e.
when all clusters are represented in the parial information, such as when Ω = [n] × [n] (full information), thisis not required to solve the constrained linear program, but we have found that this regularization was helpful in practice.
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) to the existing forest, we check that none of the points in the connected componentof i are forbidden from joining any of the points in the connected component of j. This is implemented using further
matrix multiplications and done alongside the existing check for cycles. The exact implementation is detailed in ourcode base.

C. Details of Pipeline
To be more specific on each of the three phases listed in Section 2.2 :1. We embed each Xi individually with a model Φw , using in our application neural networks and a linear model. This
allows us to use learning through clustering as a way to learn representations, and to apply this model to other elements,
6

Differentiable Clustering and Partial Fenchel-Young Losses

for which we have no clustering information, or for use in other downstream tasks.
2. We focus on cases where the similarity matrix is the negative squared distances between those embeddings. This
creates a connection between a model acting individually on each element, and a pairwise similarity matrix that canbe used as an input for our differentiable clustering operator. This mapping, from w to Σ, has derivatives that cantherefore be automatically computed by backpropagation, as it contains only conventional opereations (at least when
Φw is itself a conventional model, such as commonly used neural networks).
3. We use our proposed smoothed partial Fenchel-Young (Section 2.1) as the objective to minimize between the partial
information MΩ and Σ. The full-batch version would be to minimize LFY,ε (Σw , MΩ ) as a function of the parametersw ∈ Rp of the model. We focus instead on a mini-batch formulation for two reasons: first, stochastic optimization withmini-batches is a commonly used and efficient method for generalization in machine learning; second, it allows tohandle larger-scale data sets. As a consequence, the stochastic gradients of the loss are given, for a mini-batch b, by∇w L̄FY,ε (Σw,b , MΩ,b ) = ∂w Σw · ∇Σ L̄FY,ε (Σw,b , MΩ,b )= ∂w Σw · A∗k,ε (Σw ) − A∗k,ε (Σw ; MΩ,b ) .
The simplicity of these gradients is due to our particular choice of smoothed partial Fenchel-Young loss. They can be
efficiently estimated automatically, as described in Section 1.2, which results in a doubly stochastic scheme for the loss
optimization.

D. Illustration of Method
We recall that a k-forest on a graph is a subgraph with no cycles consisting in k connected components, potentially single
nodes.

Figure 4. Method illustration, for k = 2 - top Similarity matrix based on pairwise square distance, partial cluster connectivity information- center Clustering using spanning forests without (partial) clustering constraints - bottom Constrained clustering.

E. Proofs of Technical Results
Proof of Proposition 2.2. We show these properties successively
7Differentiable Clustering and Partial Fenchel-Young Losses

1) We have by Definition 2.1
L̄FY (θ, p) = min LFY (θ; y) .
y∈C(p)We expand thisL̄FY (θ, p) = miny∈C(p)

F (θ) − ⟨y, θ⟩ = F (θ) − max ⟨y, θ⟩ .
y∈C(p)As required, this implies, following the definitions given
L̄FY (θ, p) = F (θ) − F (θ; p) ,where F (θ; p) = max ⟨y, θ⟩,y∈C(p)

2) By linearity of derivatives and 1) above, we have∇θ L̄FY (θ, p) = ∇θ F (θ) − ∇θ F (θ; p) = y ∗ (θ) − y ∗ (θ; p) ,where y ∗ (θ; p) = argmax⟨y, θ⟩,y∈C(p)

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 wealways assume (e.g. see the following).
3) By linearity of expectation and 1) above, we have∇θ L̄FY,ε (θ, p) = ∇θ E[F (θ + εZ) − F (θ + εZ; p)] = ∇θ Fε (θ) − ∇θ Fε (θ; p) = yε∗ (θ) − yε∗ (θ; p) ,using the definitionyε∗ (θ; p) = E[argmax⟨y, θ + εZ⟩] .
y∈C(p)

Proof of Proposition 2.3. By Jensen’s inequality and the definition of the Fenchel-Young loss:
L̄FY,ε (θ; p) = E[ min LFY (θ + εZ; y)]y∈C(p)≤ min E[LFY (θ + εZ; y)]y∈C(p)= min Fε (y) − ⟨θ, y⟩y∈C(p)≤ min LFY,ε (θ; y) .
y∈C(p)

F. Additional Experimental Information
1.2

1.0

0.8

0.6

0.4

0.2

0.0

Figure 5. left: Msignal , middle: M4∗ (Σ),

right: θ after training.

We provide the details of the synthetic denoising experiment depicted in Figure 2 and described in Section 3.1. We create thesignal data Xsignal ∈ R60×2 by sampling iid. from four isotropic Gaussians (15 points coming from each of the Gaussians)
8Differentiable Clustering and Partial Fenchel-Young Losses

each having a standard deviation of 0.2. We randomly sample the means of the four Gaussians; for the example seen in
Section 3.1 the sampled means were:
0.97627008 4.30378733
 2.05526752 0.89766366
−1.52690401 2.91788226 .
−1.24825577 7.83546002
Let Σsignal be the pairwise euclidean similarity matrix corresponding to Xsignal , and furthermore let Msignal := M4∗ (Σsignal )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 ∈ R60×4 , where the noise
entries were sampled iid from a continuous unit uniform distribution. Similarly letting Σ be the pairwise euclidean similarity
matrix corresponding to X, we calculate M4∗ (Σ) ̸= Msignal . Both the matrices Msignal and M4∗ (Σ) are depicted in Figure 2 ;we remark that adding the noise dimensions leads to most points being assigned one of two clusters, and two points being
isolated alone in their own clusters. We also create a validation set of equal size (in exactly the same manner as the train set),to ensure the model has not over-fitted to the train set.
The goal of the experiment is to learn a linear transformation of the data that recovers Msignal i.e. a denoising, by minimizing
the partial loss. There are multiple solutions to this problem, the most obvious being a transformation that removes the last
two noise columns from X:
1
0θ∗ = 
0
0


0
1
,
0
0

for which Xθ∗ = Xsignal

For any θ ∈ R4×2 , we define Σθ to be the pairwise similarity matrix corresponding to Xθ, and M4∗ (Σθ ) to its corresponding
equivalence matrix. Then the problem can be summarized as:
min L̄FY,ε (Σθ , Msignal ).

θ∈R4×2(1)

We initialized θ from a standard Normal distribution, and minimized the partial loss via stochastic gradient descent, with alearning rate of 0.01 and batch size 32.
For perturbations, we took ε = 0.1 and B = 1000, where ε denotes the noise amplitude in randomized smoothing and B
denotes the number of samples in the Monte-Carlo estimate. The validation clustering error converged to zero after 25
gradient batches. We verify that the θ attained from training is indeed learning to remove the noise dimensions (see Figure5).
F.1. Supervised Differentiable Clustering
As mentioned in Section 3.1, our method is able to cluster classical data sets such as MNIST and Fashion MNIST. We
trained a CNN with the LeNet-5 architecture (LeCun et al., 1998) using our proposed partial loss as the objective function.
The exact details of the CNN architecture are depicted in Figure 6. For this experiment and all experiments described below,
we trained on a single Nvidia V100 GPU ; training the CNN with our proposed pipeline took < 15 minutes.
The model was trained for 30k gradient steps on mini-batches of size 64. We used the Adam optimizer (Kingma & Ba,2015) with learning rate η = 3 × 10−4 , momentum parameters (β1 , β2 ) = (0.9, 0.999), and an ℓ2 weight decay of 10−4 .
We validated / tested the model using the zero-one error between the true equivalence matrix and the equivalence matrix
corresponding to the output of the model. We used an early stopping of 10k steps (i.e. training was stopped if the validation
clustering error did not improve over 10k steps). For efficiency (and parallelization), we also computed this clustering error
batch-wise with batch-size 64. As stated in Section 3.1, we attained a batch-wise clustering precision of 0.99 for MNIST
and 0.96 on Fashion MNIST.
9Differentiable Clustering and Partial Fenchel-Young Losses
Conv: features = 32, kernel = 3x3
ReLU & Avg Pool: window 2x2, strides = 2x2

Conv: features = 64, kernel = 3x3
ReLU & Avg Pool: window 2x2, strides = 2x2

Dense: d = 256
ReLU

Dense: d = 256

Figure 6. left: Architecture of the CNN, middle: t-SNE visualization of train data embeddings, right : tSNE visualization of validation
data embeddings.

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 6. It can be seen that the model has learnt a distinct cluster for each of the ten classes.
F.2. Semi-Supervised Differentiable Clustering
As mentioned in Section 3.2, we show that our method is particularly useful in settings where labelled examples are scarce,even in the particularly challenging case of having no labelled examples for some classes. Our approach allows a model toleverage the semantic information of unlabeled examples when trying to predict a target equivalence matrix MΩ ; this isowing to the fact that the prediction of a class for a single point depends on the values of all other points in the batch, whichis 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:
1. Semi-Supervised Clustering: Train a CNN to learn an embedding of the data which leads to a good clustering error.
We can compare our methodology to that of a baseline model trained using the cross-entropy loss. This is to checkthat our model has leveraged information from the unlabeled data and that our partial loss is indeed leading to goodclustering performance.
2. Downstream Classification: Assess the trained model’s capacity to serve as a backbone in a downstream classification
task (transfer learning), where its weights are frozen and a linear layer is trained on top of the backbone.
We describe our data processing for both of these tasks below.
F.2.1. DATA S ETSIn our semi-supervised learning experiments, we divided the standard MNIST train split in the following manner:• We create a balanced hold-out data set consisting of 1k images (100 images from each of the 10 classes). This hold-out
data set will be used to assess the utility of the frozen clustering model on a downstream classification problem.
• From the remaining 59k images, we select a labeled train set of nℓ = {100, 250, 500, 1000, 2500, 5000} images. Our
experiments also vary kw ∈ {0, 3, 6}, the number of digits to withhold all labels from. For example, if kw = 3, thenthe labels for the images corresponding to digits {0, 1, 2} will never appear in the labeled train data.
F.2.2. S EMI -S UPERVISED C LUSTERING TASKFor each of the 18 data sets above, we train the same CNN architecture as that described in Section F.1 using the following
two approaches:
10Differentiable Clustering and Partial Fenchel-Young Losses

1. Ours: The CNN is trained on mini-batches, where half the batch is labeled data and half the batch is unlabeled data
(i.e. a semi-supervised learning regime), to minimize the partial loss.
2. Baseline: The baseline model consists of a CNN with the same architecture as that described in Section F.1, but withan additional dense layer with output dimension 10 (the number of classes). We train the model using mini-batches
consisting of labeled points, minimizing the cross-entropy loss. The training regime is fully-supervised learning(classification). The baseline backbone refers to all of the model, minus the dense output layer.
Both models were trained with mini-batches of size 64, with points sampled uniformly without replacement. All hyperparameters and optimization metrics were identical to those detailed in Section 3.1. We ran each experiment on each dataset
for five different random seeds s ∈ {1, . . . , 5}, in order to report population statistics on the clustering error.
F.2.3. T RANSFER -L EARNING : D OWNSTREAM C LASSIFICATIONIn this task both models are frozen, and their utility as a foundational backbone is assessed on a downstream classification
task using the hold-out data set. We train a linear (a.k.a dense) layer on top of both models using the cross-entropy loss. We
refer to this linear layer as the downstream head. Training this linear head is equivalent to performing multinomial logistic
regression on the features of the model.
To optimize the weights of the linear head we used the SAGA optimizer (Defazio et al., 2014). The results are depicted in
Figure 3. It can be seen that training a CNN backbone using our approach with just 250 labels leads to better downstream
classification performance than the baseline trained with 5000 labels. It is worth remarking that the baseline backbone
was trained on the same objective function (cross-entropy) and task (classification) as the downstream problem, which is
not the case for the backbone corresponding to our approach. This highlights how learning ‘cluster-able embeddings’ andleveraging unlabeled data can be desirable for transfer learning.

11