SlowGuess/NeurIPS_2023_OCR huggingface dataset:
In structured prediction, a common modus operandi is to minimize a loss between some structured ground truth response or label  $y \in \mathcal{V}$  and a score  $\theta \in \mathbb{R}^d$  (the latter often itself the output of a parameterized model). As an example, in logistic regression  $y \in \{0,1\}$  and  $\theta = \langle x,\beta \rangle \in \mathbb{R}$ . The framework of Fenchel-Young losses allows to tackle much more complex structures, such as cluster connectivity information.
Definition 7 (Fenchel-Young loss–Blondel et al. (2020a)). Let  $F$  be a convex function on  $\mathbb{R}^d$ , and  $F^*$  its Fenchel dual on a convex set  $\mathcal{C} \subset \mathbb{R}^d$ . The Fenchel-Young loss between  $\theta \in \mathbb{R}^d$  and  $y \in \operatorname{int}(\mathcal{C})$  is
$$
L _ {F Y} (\theta ; y) = F (\theta) - \langle \theta , y \rangle + F ^ {*} (y).
$$
The Fenchel-Young (FY) loss satisfies several properties, making it useful for learning. In particular, it is nonnegative, convex in  $\theta$ , handles well noisy labels, and its gradient with respect to the score can be efficiently computed, with  $\nabla_{\theta}L_{\mathrm{FY}}(\theta ;y) = \nabla_{\theta}F(\theta) - y$  (see, e.g., Blondel et al., 2020a).
In the case of linear programs over a polytope  $\mathcal{C}$ , taking  $F$  to be the so-called support function defined as  $F(\theta) = \max_{y\in \mathcal{C}}\langle y,\mathcal{C}\rangle$  (consistent with our notation so far), the dual function  $F^{*}$  is the indicator function of  $\mathcal{C}$  (Rockafellar, 1997), and the Fenchel-Young function is then given by
$$
L _ {\mathsf {F Y}} (\theta , y) = \left\{ \begin{array}{l l} F (\theta) - \langle \theta , y \rangle & \text {i f} y \in \mathcal {C}, \\ + \infty & \text {o t h e r w i s e}. \end{array} \right.
$$
In the case of a perturbed maximum for  $F$ , see, e.g., Berthet et al. (2020), we have
$$
F _ {\varepsilon} (\theta) = \mathbf {E} [ \max  _ {y \in \mathcal {C}} \langle y, \theta + \varepsilon Z \rangle ], \quad \text {a n d} \quad y _ {\varepsilon} ^ {*} (\theta) = \mathbf {E} [ \operatorname * {a r g m a x} _ {y \in \mathcal {C}} \langle y, \theta + \varepsilon Z \rangle ].
$$
In this setting (under mild regularity assumptions on the noise distribution and  $\mathcal{C}$ ), we have that  $\varepsilon \Omega = F_{\varepsilon}^{*}$  is a strictly convex Legendre function on  $\mathcal{C}$  and  $y_{\varepsilon}^{*}(\theta) = \nabla_{\theta}F_{\varepsilon}(\theta)$ . This is actually equivalent to having max and argmax with a  $-\varepsilon \Omega(y)$  regularization (see Berthet et al., 2020, Propositions 2.1
and 2.2). In this case, the FY loss is given by  $L_{\mathsf{FY},\varepsilon}(\theta ;y) = F_{\varepsilon}(\theta) - \langle \theta ,y\rangle +\varepsilon \Omega (y)$  and its gradient by  $\nabla_{\theta}L_{\mathsf{FY},\varepsilon}(\theta ;y) = y_{\varepsilon}^{*}(\theta) - y$  . One can also define it as  $L_{\mathsf{FY},\varepsilon}(\theta ;y) = \mathbf{E}[L_{\mathsf{FY}}(\theta +\varepsilon Z;y)]$  and it has the same gradient. In the perturbations case, it can be easily obtained by Monte-Carlo estimates of the perturbed optimizer, taking  $\frac{1}{B}\sum_{i = 1}^{B}y^{*}(\theta +\varepsilon Z_{i}) - y$
3.2 Partial Fenchel-Young losses
As noted above, this loss is widely applicable in structured prediction. As presented here, it requires label data  $y$  of the same kind than the optimizer  $y^{*}$ . In our setting, the linear program is over spanning  $k$ -forests rather than connectivity matrix. This is no accident: linear programs over cut matrices (when  $k = 2$ ) are already NP-hard (Karp, 1972). If one has access to richer data (such as ground-truth spanning forest information), one can ignore the operator  $M^{*}$  and focus only on  $A_{k,\varepsilon}^{*}$ ,  $F_{k,\varepsilon}$ , and the associated FY-loss. However, in most cases, clustering data can reasonably only be expected to be present as connectivity matrix. It is therefore necessary to alter the Fenchel-Young loss, and we introduce the following loss, which allows to work with partial information  $p \in \mathcal{P}$  representing partial information about the unknown ground-truth  $y$ . Using this kind of "inf-loss" is common when dealing with partial label information (see, e.g., Cabannes et al., 2020).
Definition 8 (Partial Fenchel-Young loss). Let  $F$  be a convex function,  $L_{\mathsf{FY}}$  the associated Fenchel-Young loss, and for every  $p \in \mathcal{P}$  a convex constraint subset  $\mathcal{C}(p) \subseteq \mathcal{C}$ .
$$
\bar {L} _ {\mathsf {F Y}} (\theta ; p) = \min  _ {y \in \mathcal {C} (p)} L _ {\mathsf {F Y}} (\theta ; y).
$$
This allows to learn from incomplete information about  $y$ . When we do not know its value, but only a subset of  $\mathcal{C}(p) \subseteq \mathcal{C}$  to which it might belong, we can minimize the infimum of the FY losses that are consistent with the partial label information  $y \in \mathcal{V}(p)$ .
Proposition 1. When  $F$  is the support function of a compact convex set  $\mathcal{C}$ , the partial Fenchel-Young loss (see Definition 8) satisfies
1. The loss is a difference of convex functions in  $\theta$  given explicitly by
$$
\bar {L} _ {\mathsf {F Y}} (\theta ; p) = F (\theta) - F (\theta ; p), \quad w h e r e \quad F (\theta ; p) = \max  _ {y \in \mathcal {C} (p)} \langle y, \theta \rangle ,
$$
2. The gradient with respect to  $\theta$  is given by
$$
\nabla_ {\theta} \bar {L} _ {\mathbb {F Y}} (\theta ; p) = y ^ {*} (\theta) - y ^ {*} (\theta ; p), \quad w h e r e \quad y ^ {*} (\theta ; p) = \operatorname * {a r g m a x} _ {y \in \mathcal {C} (p)} \langle y, \theta \rangle .
$$
3. The perturbed partial Fenchel-Young loss given by  $\bar{L}_{\mathsf{FY},\varepsilon}(\theta ;p) = \mathbf{E}\big[\bar{L}_{\mathsf{FY}}(\theta +\varepsilon Z;p)\big]$  satisfies
$$
\nabla_ {\theta} \bar {L} _ {\mathsf {F Y}, \varepsilon} (\theta ; p) = y _ {\varepsilon} ^ {*} (\theta) - y _ {\varepsilon} ^ {*} (\theta ; p), \quad w h e r e \quad y _ {\varepsilon} ^ {*} (\theta ; p) = \mathbf {E} [ \operatorname * {a r g m a x} _ {y \in \mathcal {C} (p)} \langle y, \theta + \varepsilon Z \rangle ].
$$
Another possibility would be to define the partial loss as  $\min_{y\in \mathcal{C}(p)}L_{\mathsf{FY},\varepsilon}(\theta ;y)$ , that is, the infimum of smoothed losses instead of the smoothed infimum loss  $L_{\mathsf{FY},\varepsilon}(\theta ;p)$  defined above. However, there is no direct method to minimizing the smoothed loss with respect to  $y\in \mathcal{C}(p)$ . Note that we have a relationship between the two losses:
Proposition 2. Letting  $L_{\mathsf{FY},\varepsilon}(\theta ;y) = \mathbf{E}[L_{\mathsf{FY},\varepsilon}(\theta +\varepsilon Z;y)]$  and  $\bar{L}_{\mathsf{FY},\varepsilon}$  as in Definition 8, we have
$$
\bar {L} _ {\mathsf {F Y}, \varepsilon} (\theta ; p) \leq \min  _ {y \in \mathcal {C} (p)} L _ {\mathsf {F Y}, \varepsilon} (\theta ; y).
$$
The proofs of the above two propositions are detailed in the Appendix.
3.3 Applications to differentiable clustering
We apply this framework, as detailed in the following section and in Section 4, to clustering. This is done naturally by transposing notations and taking  $\mathcal{C} = \mathcal{C}_k$ ,  $\theta = \Sigma$ ,  $y^{*} = A_{k}^{*}$ ,  $p = M_{\Omega}$ ,  $\mathcal{C}(p) = \mathcal{C}_k(M_\Omega)$ , and  $y^{*}(\theta; p) = A_{k}^{*}(\Sigma; M_{\Omega})$ . In this setting the perturbed partial Fenchel-Young loss satisfies
$$
\nabla_ {\Sigma} \bar {L} _ {\mathsf {F Y}, \varepsilon} (\Sigma ; M _ {\Omega}) = A _ {k, \varepsilon} ^ {*} (\Sigma) - A _ {k, \varepsilon} ^ {*} (\Sigma ; M _ {\Omega}).
$$
We learn representations of a data that fit with clustering information (either complete or partial). As described above, we consider settings with  $n$  elements described by their features  $\mathbf{X} = X_{1},\ldots ,X_{n}$  in some feature space  $\mathcal{X}$ , and  $M_{\Omega}$  some clustering information. Our pipeline to learn representations includes the following steps (see Figure 1)
i) Embed each  $X_{i}$  in  $\mathbb{R}^d$  with a parameterized model  $v_{i} = \Phi_{w}(X_{i})\in \mathbb{R}^{d}$ , with weights  $w\in \mathbb{R}^p$
ii) Construct a similarity matrix from these embeddings, e.g.  $\Sigma_{w,ij} = -\|\Phi_w(X_i) - \Phi_w(X_j)\|_2^2$ .
iii) Stochastic optimization of the expected loss of  $\bar{L}_{\mathbb{F}\mathcal{V},\varepsilon}(\Sigma_{w,b},M_{\Omega ,b})$  , using mini-batches of  $\mathbf{X}$
Details. To be more specific on each of these phases:  $i$  ) We embed each  $X_{i}$  individually with a model  $\Phi_w$ , using in our application neural networks and a linear model. This allows us to use learning through clustering as a way to learn representations, and to apply this model to other elements, for which we have no clustering information, or for use in other downstream tasks.  $ii$  ) We focus on cases where the similarity matrix is the negative squared distances between those embeddings. This creates a connection between a model acting individually on each element, and a pairwise similarity matrix that can be used as an input for our differentiable clustering operator. This mapping, from  $w$  to  $\Sigma$ , has derivatives that can therefore be automatically computed by backpropagation, as it contains only conventional operations (at least when  $\Phi_w$  is itself a conventional model, such as commonly used neural networks). iii) We use our proposed smoothed partial Fenchel-Young (Section 3.2) as the objective to minimize between the partial information  $M_{\Omega}$  and  $\Sigma$ . The full-batch version would be to minimize  $L_{\mathrm{FY},\varepsilon}(\Sigma_w,M_\Omega)$  as a function of the parameters  $w\in \mathbb{R}^p$  of the model. We focus instead on a mini-batch formulation for two reasons: first, stochastic optimization with mini-batches is a commonly used and efficient method for generalization in machine learning; second, it allows to handle larger-scale datasets. As a consequence, the stochastic gradients of the loss are given, for a mini-batch  $b$ , by
$$
\nabla_ {w} \bar {L} _ {\mathsf {F Y}, \varepsilon} (\Sigma_ {w, b}; M _ {\Omega , b}) = \partial_ {w} \Sigma_ {w} \cdot \nabla_ {\Sigma} \bar {L} _ {\mathsf {F Y}, \varepsilon} (\Sigma_ {w, b}; M _ {\Omega , b}) = \partial_ {w} \Sigma_ {w} \cdot \left(A _ {k, \varepsilon} ^ {*} (\Sigma_ {w}) - A _ {k, \varepsilon} ^ {*} (\Sigma_ {w}; M _ {\Omega , b})\right).
$$
The simplicity of these gradients is due to our particular choice of smoothed partial Fenchel-Young loss. They can be efficiently estimated automatically, as described in Section 2.2, which results in a doubly stochastic scheme for the loss optimization.
We apply our framework for learning through clustering in both a supervised and a semi-supervised setting, as illustrated in Figure 1. Formally, for large training datasets of size  $n$ , we either have access to a full cluster connectivity matrix  $M_{\Omega}$  or a partial one (typically built by using partial label information, see below). We use this clustering information  $M_{\Omega}$ , from which mini-batches can be
Figure 3: (left) Illustration of clustering methods on toy-data sets. (right) Using cluster information to learn a linear de-noising (bottom) of a noised signal (top).
Figure 4: Semi-supervised learning: Performance of a CNN trained on partially labeled MNIST data (top row), and a ResNet trained on partially labeled Cifar-10 data (bottom row). Our trained model (full line) is evaluated on clustering (left column) and compared to a model entirely trained on the classification task (dashed line). Both models are evaluated on a down-stream classification problem i.e. transfer learning, via a 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 on a test dataset: first, by evaluating the clustering accuracy (i.e. proportion of correct coefficients in the predicted cluster connectivity matrix), and second by training a shallow model on a classification task (using clusters as classes) on a holdout set, evaluating it on a test set.
4.1 Supervised learning
We apply first our method to synthetic datasets - purely to provide an illustration of both our internal clustering algorithm, and of our learning procedure. In Figure 2, we show how the clustering operator that we use, based on spanning forests (i.e. single linkage), with Kruskal's algorithm is efficient on some standard synthetic examples, even when they are not linearly separable (compared to  $k$ -Means, mean-shift, Gaussian mixture-model). We also show that our method allows us to perform supervised learning based on cluster information, in a linear setting. For  $X_{\mathrm{signal}}$  consisting of  $n = 60$  points in two dimensions consisting of data in four well-separated clusters (see Figure 2), we construct  $X$  by appending two noise dimensions, such that clustering based on pairwise square distances between  $X_{i}$  mixes the original clusters. We learn a linear de-noising transformation  $X\theta$ , with  $\theta \in \mathbb{R}^{4\times 2}$  through clustering, by minimizing our perturbed partial FY loss with SGD, using the known labels as supervision.
We also show that our method is able to cluster virtually all of some classical datasets. We train a CNN (LeNet-5 LeCun et al. (1998)) on mini-batches of size 64 using the partial Fenchel-Young loss to learn a clustering, with a batch-wise clustering precision of 0.99 for MNIST and 0.96 on Fashion MNIST. Using the same approach, we trained a ResNet (He et al., 2016) on CIFAR-10 (with some minor modifications to kernel size and stride), achieving a batch-wise clustering test precision of 0.93. The t-SNE visualization of the embeddings for the validation set are displayed in Figure 5. The experimental setups (as well as visualization of learnt clusters for MNIST), are detailed in the Appendix.
Figure 5: (left) t-SNE of learnt embeddings for the CNN trained on MNIST with  $k_w = 3$  withheld classes (highlighted). (right) t-SNE of learnt embeddings for the ResNet trained on 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 et al., 2010) and Cifar-10 (Krizhevsky et al., 2009) datasets, where we independently vary both the total number of labelled examples  $n_{\ell}$ , as well as the number of withheld classes for which no labelled examples are present in the training set  $k_{w} \in \{0,3,6\}$ . For MNIST we consider data sets with  $n_{\ell} \in \{100,250,500,1000,2000,5000\}$  labelled examples, and for Cifar-10 we consider  $n_{\ell} \in \{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  $\varepsilon = 0.1$  and  $B = 100$  for the estimated loss gradients, and optimize weights using Adam (Kingma and Ba, 2015).
In left column of Figure 4, we report the test clustering error (evaluated on mini-batches of the same size), for each of the models, and data sets, corresponding to the choice of  $n_{\ell}$  and  $k_{w}$ . We compare the performance of each model to a baseline, consisting of the exact same architecture (plus a linear head mapping the embedding to logits), trained to minimize the cross entropy loss.
In addition, we evaluate all models on a down-stream (transfer-learning) classification task, by learning a linear probe on top of the frozen model, trained on hold-out data with all classes present. The results are depicted in the right hand column of Figure 4. Further information regarding the experimental details can be found in the Appendix.
We observe that learning through clustering allows to find a representation where class semantics are easily recoverable from the local topology. On Cifar-10, our proposed approach strikingly achieves a lower clustering error in the most challenging setting ( $n_{\ell} = 1000$  labelled examples and  $k_{w} = 6$  withheld classes) than the classification-based baseline in the most lenient setting ( $n_{\ell} = 10,000$  labelled examples all classes observed). Importantly, this advantage is not limited to clustering metrics: learning through clustering can also lead to lower down-stream classification error, with the gap being most pronounced when few labelled examples are available.
Moreover, besides this pronounced improvement in data efficiency, we found that our method is capable of clustering classes for which no labelled examples are seen during training (see Figure 5, left). Therefore, investigating potential applications of learning through clustering to zero-shot and self-supervised learning are promising avenues for future work.
We thank Jean-Philippe Vert for discussions relating to this work, Simon Legrand for conversations relating to the CLEPS computing cluster (Cluster pour l'Experimentation et le Prototypage Scientifique), and Matthieu Blondel for providing references regarding learning via Kirchhoff's theorem. We also thank the reviewers and Kyle Heuton, for providing helpful corrections for the camera-ready 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'avoir" program, reference ANR19-P3IA-0001 (PRAIRIE 3IA Institute), as well as from the European Research Council (grant SEQUOIA 724063), and a monetary don from Google.
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., Topfer, 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 perturbed optimizers. Advances in Neural Information Processing Systems, 33:9508-9519.
Blondel, M., Berthet, Q., Cuturi, M., Frostig, R., Hoyer, S., Llinares-Lopez, F., Pedregosa, F., and Vert, J.-P. (2022). Efficient and modular implicit differentiation. Advances in Neural Information Processing Systems, 35:5230-5242.
Blondel, M., Martins, A. F., and Niculae, V. (2020a). Learning with Fenchel-Young losses. The Journal of Machine Learning Research, 21(1):1314-1382.
Blondel, M., Teboul, O., Berthet, Q., and Djolonga, J. (2020b). Fast differentiable sorting and ranking. In International Conference on Machine Learning, pages 950-959.
Bradbury, J., Frostig, R., Hawkins, P., Johnson, M. J., Leary, C., Maclaurin, D., Necula, G., Paszke, A., VanderPlas, J., Wanderman-Milne, S., and Zhang, Q. (2018). JAX: composable transformations of Python+NumPy programs.
Bubeck, S., Meilă, M., and von Luxburg, U. (2012). How the initialization affects the stability of the k-means algorithm. *ESAIM: Probability and Statistics*, 16:436-452.
Cabannes, V., Rudi, A., and Bach, F. (2020). Structured prediction with partial labelling through the infimum loss. In International Conference on Machine Learning, pages 1230-1239.
Caron, M., Bojanowski, P., Joulin, A., and Douze, M. (2018). Deep clustering for unsupervised learning of visual features. In Proceedings of the European conference on computer vision (ECCV), pages 132-149.
Caron, M., Misra, I., Mairal, J., Goyal, P., Bojanowski, P., and Joulin, A. (2020). Unsupervised learning of visual features by contrasting cluster assignments. Advances in neural information processing systems, 33:9912-9924.
Carr, A. N., Berthet, Q., Blondel, M., Teboul, O., and Zeghidour, N. (2021). Self-supervised learning of audio representations from permutations with differentiable ranking. IEEE Signal Processing Letters, 28:708-712.
This submission focuses on foundational and exploratory work, with application to general machine learning techniques. We propose methods to extend the range of operations that can be used in end-to-end differentiable systems, by allowing to include clustering as a differentiable operation in machine learning models. We do not foresee societal consequences that are specifically related to these methods, beyond those that are associated with the field in general.
6 Reproducibility and Licensing
Our experiments use data sets that are already open-sourced and cited in the references. All the code written for this project and used to implement our experiments is available at:
https://github.com/LawrenceMMStewart/DiffClust_NeurIPS2023,
and distributed under a Apache 2.0. license.
At present, our implementation of Kruskal's algorithm is incompatible with processing very large batch sizes at train time. However, as our method can use mini-batches, it is suitable for training deep learning models on large data sets. For example, we use a batch size of 64 (similar to that for standard classification models) for all experiments detailed in our submission. Other methods that focus on cluster assignment, rather than pairwise clustering, often have an  $n \times k$  parametrization, whose efficiency will depend on the comparison between  $n$  (batch size) and  $k$  (number of clusters).
At inference time this is not the case, since gradients need not be back-propagated hence, any implementation of Kruskal's algorithm can be used such as the union-find implementation. Faster GPU compatible implementations of Kruskal's algorithm for training are certainly possible ; improvements could be made using a binary heap or a GPU parallelized merge sort. In many learning experiments the batch sizes are typically of values  $\{32,64,128,256,512\}$  , so our current implementation is sufficient, however, for some specific applications (such as bio-informatics experiments requiring tens of thousands of points in a batch), it would be necessary to improve upon the current implementation. This is mainly an engineering task relating to data-structures and GPU programming, and is beyond the scope of this paper.
Finally, our clustering approach solves a linear program whose argmax solution is a matroid problem with a greedy single-linkage criterion. Whilst we saw that this method is effective for differentiable clustering in both the supervised and semi-supervised setting, there may well be other linkage criteria which also take the form of LP's but whom are more robust to outliers and noise in the train dataset. This line of work is outside the scope of the paper, but is closely related and could help improve our differentiable spanning forests framework.
8 Proofs of Technical Results
Proof of Proposition 1. We show these properties successively
1) We have by Definition 8
We expand this
$$
\bar {L} _ {\mathsf {F Y}} (\theta ; p) = \min  _ {y \in \mathcal {C} (p)} \left\{F (\theta) - \langle y, \theta \rangle \right\} = F (\theta) - \max  _ {y \in \mathcal {C} (p)} \langle y, \theta \rangle .
$$
As required, this implies, following the definitions given
$$
\bar {L} _ {\mathsf {F Y}} (\theta ; p) = F (\theta) - F (\theta ; p), \quad \text {w h e r e} \quad F (\theta ; p) = \max  _ {y \in \mathcal {C} (p)} \langle y, \theta \rangle ,
$$
2) By linearity of derivatives and 1) above, we have
$$
\nabla_ {\theta} \bar {L} _ {\mathsf {F Y}} (\theta ; p) = \nabla_ {\theta} F (\theta) - \nabla_ {\theta} F (\theta ; p) = y ^ {*} (\theta) - y ^ {*} (\theta ; p), \quad \text {w h e r e} \quad y ^ {*} (\theta ; p) = \operatorname * {a r g m a x} _ {y \in \mathcal {C} (p)} \langle y, \theta \rangle ,
$$