Attention-Based Clustering: Learning a Kernel from Context

1 INTRODUCTION . Many problems in machine learning involve modelling the relations between elements of a set . A notable example , and the focus of this paper , is clustering , in which the elements are grouped according to some shared properties . A common approach uses kernel methods : a class of algorithms that operate on pairwise similarities , which are obtained by evaluating a specific kernel function ( Filippone et al. , 2008 ) . However , for data points that are not trivially comparable , specifying the kernel function is not straightforward . With the advent of deep learning , this gave rise to metric learning frameworks where a parameterized binary operator , either explicitly or implicitly , is taught from examples how to measure the distance between two data points ( Bromley et al. , 1993 ; Koch et al. , 2015 ; Zagoruyko & Komodakis , 2015 ; Hsu et al. , 2018 ; Wojke & Bewley , 2018 ; Hsu et al. , 2019 ) . These cases operate on the assumption that there exists a global metric , that is , the distance between points depends solely on the two operands . This assumption disregards situations where the underlying metric is contextual , by which we mean that the distance between two data points may depend on some structure of the entire dataset . We hypothesize that the context provided by a set of data points can be helpful in measuring the distance between any two data points in the set . As an example of where context might help , consider the task of clustering characters that belong to the same language . There are languages , like Latin and Greek , that share certain characters , for example the Latin T and the Greek upper case τ .1 However , given two sentences , one from the Aeneid and one from the Odyssey , we should have less trouble clustering the same character in both languages correctly due to the context , even when ignoring any structure or meaning derived from the sentences themselves . Indeed , a human performing this task will not need to rely on prior knowledge of the stories of Aeneas or Odysseus , nor on literacy in Latin or Ancient Greek . As a larger principle , it is well recognized that humans perceive emergent properties in configurations of objects , as documented in the Gestalt Laws of Perceptual Organization ( Palmer , 1999 , Chapter 2 ) . We introduce Attention-Based Clustering ( ABC ) which uses context to output pairwise similarities between the data points in the input set . Like other approaches in the literature ( Hsu et al. , 2018 ; 2019 ; Han et al. , 2019 ; Lee et al. , 2019b ) , our model is trained with ground-truth labels in the form of pairwise constraints , but in contrast to other methods , ours can be used with an unsupervised clustering method to obtain cluster labels . To demonstrate the benefit of using ABC over pairwise 1To the extend that there is not even a LaTeX command \Tau metric learning methods , we propose a clustering problem that requires the use of properties emerging from the entire input set in order to be solved . The task is to cluster a set of points that lie on a number of intersecting circles , which is a generalization of the Olympic circles problem ( Anand et al. , 2014 ) . Pairwise kernel methods for clustering perform poorly on the circles problem , whereas our ABC handles it with ease , as displayed in Figure 1 . We use the circles dataset for an ablation study in Section 6.1 . In recent years , numerous deep neural network architectures have been proposed for clustering ( Xie et al. , 2016 ; Min et al. , 2018 ) . The idea of using more than pairwise interactions between elements of an input set in order to improve clustering has been pursued recently in Lee et al . ( 2019a ; b ) , and is motivated by the problem of amortized clustering ( Gershman & Goodman , 2014 ; Stuhlmüller et al. , 2013 ) . Our architecture is inspired by the Transformer ( Vaswani et al. , 2017 ) , which was used by Lee et al . ( 2019a ) as the Set Transformer to improve clustering ( Lee et al. , 2019b ) . We inherit its benefits such as being equivariant under permutations as well as agnostic to input size . However , our approach is motivated by the use of context to improve metric learning , giving us a model that is moreover agnostic to the number of clusters in the sense that neither a prediction nor a bound on the number of clusters needs to be specified for the architecture definition . We also provide theoretical evidence that the Transformer architecture is effective for metric learning and clustering , and to our knowledge , are the first to do so . The idea of using deep metric learning to improve clustering has been pursued in Koch et al . ( 2015 ) ; Zagoruyko & Komodakis ( 2015 ) ; Hsu et al . ( 2018 ; 2019 ) ; Han et al . ( 2019 ) , but without considering the use of context . We use ground-truth labels , only in the form of pairwise constraints , to train a similarity kernel , making our approach an example of constrained clustering . These algorithms are often categorized by whether they use the constraints to only learn a metric or to also generate cluster labels ( Hsu et al. , 2018 ) . Our architecture belongs to the former category , where we only use the constraints to learn a metric and rely on an unconstrained clustering process to obtain cluster labels . Despite this , we achieve nearly state-of-the-art clustering results on the Omniglot , embedded ImageNet , and CIFAR-100 datasets , comparable to sophisticated methods that synthesize clusters , either using the constraints ( Hsu et al. , 2018 ; 2019 ; Han et al. , 2019 ) or otherwise ( Lee et al. , 2019a ; b ) . Our main contributions are : • ABC incorporates context in a general and flexible manner to improve metric learning for clustering . Our competitive results on Omniglot , embedded ImageNet , and CIFAR-100 , as well as our ablation study on our circles dataset provide support for the use of context in metric learning algorithms . • We provide theoretical evidence of why the self-attention module in the Transformer architecture is well suited for clustering , justifying its effectiveness for this task . 2 RELATED WORKS . Our method is similar to a line of research where a distance metric , rather than a similarity score , is learned in a supervised manner , which can then be used as input to off-the-shelf clustering methods ( Xing et al. , 2003 ; Shalev-Shwartz et al. , 2004 ; Davis et al. , 2007 ) . Only certain classes of distances , such as the Mahalanobis distance , are learned . In general , deep neural nets offer the ability to learn a more general class of distances , and have been used to learn a pairwise metric in numerous works ( Zagoruyko & Komodakis , 2015 ; Hsu et al. , 2018 ; Wojke & Bewley , 2018 ; Hsu et al. , 2019 ) , most notably in the Siamese network ( Bromley et al. , 1993 ; Koch et al. , 2015 ) . The idea of using contextual information has not been explored in any of these papers . Many models go further than metric learning by also learning how to synthesize clusters . An example of constrained clustering can be found in Anand et al . ( 2014 ) ; Amid et al . ( 2015 ) , where pairwise constraints are used to transform a predefined kernel in an iterative manner , which is used in a kernel mean shift clustering algorithm . Constrained clustering algorithms have been implemented using deep neural nets as well . In Hsu et al . ( 2018 ; 2019 ) , the authors train a similarity metric and transfer learning to a secondary clustering model . Both models are trained using only pairwise constraints , and any available context information remains unused in both components of their architecture . In Han et al . ( 2019 ) , a constrained clusterer inspired by the deep embedded clustering idea ( Xie et al. , 2016 ) is proposed , along with a number of best practices such as temporal ensembling and consistency constraints in the loss function . These techniques are fairly generic and can perhaps be applied to any other clustering algorithm to improve its results . Their model generates clusters by slowly annealing them , requiring optimization and back-propagation even at test time . The models from Hsu et al . ( 2018 ) and Hsu et al . ( 2019 ) also have this requirement . This may not be feasible during deployment . A more detailed discussion of the differences between our approach and that of Lee et al . ( 2019a ; b ) is in order . The Set Transformer architecture ( Lee et al. , 2019a ) uses the Transformer as a contextual encoder , followed by a pooling layer that uses a fixed number of seed vectors as queries . This architecture is used to cluster a mixture of Gaussians , but is less flexible than ours for two reasons : it requires the number of clusters in advance in setting the number of seed vectors , and those seed vectors being learned makes their approach less adaptable to unseen classes . The first limitation is addressed in a follow-up paper ( Lee et al. , 2019b ) , with the use of an iterated process to filter out clusters and a stopping condition . Our architecture , due to its use of metric learning in place of the pooling layer with learned seed vectors , is inductive and better suited to handle new classes . We also present a mathematical justification for the use of the Transformer in clustering applications . Lastly , Lee et al . ( 2019a ) contains no clustering results on real-world data . Lee et al . ( 2019b ) does and our results on embedded ImageNet are similar , while ours on Omniglot are significantly better . 3 BACKGROUND . Taking inspiration from kernel methods , we aim to compute a similarity matrix from a sequence of data points . Our architecture is inspired by ideas from two streams : the metric learning literature and the Siamese network ( Bromley et al. , 1993 ; Koch et al. , 2015 ) on how to learn compatibility scores , and the Transformer architecture ( Vaswani et al. , 2017 ) and the Set Transformer ( Lee et al. , 2019a ) on how to use context to make decisions . We discuss a few concepts from the literature which will form building blocks of our architecture in the next section . 3.1 COMPATIBILITY . In this section we introduce some compatibility functions which compute a similarity score between two vector arguments , called the query and key respectively . We present the forms of compatibility used in this paper in Table 1 and for both of these forms , keys and queries are required to have equal dimension d. In Siamese Networks ( Koch et al. , 2015 ) , compatibility between two input images is measured by the sigmoid of a weighted L1-distance between representations of the input . This can be seen as a special case of additive compatibility above . The Transformer ( Vaswani et al. , 2017 ) and Set Transformer ( Lee et al. , 2019a ; b ) make use of multiplicative compatibility .

Summary: This paper proposes a method for producing representations for clustering that take into account global trends in the dataset, rather than considering each pair of instances in isolation. They claim to achieve competitive clustering performance on omniglot, which they attribute to the use of these contextualised embeddings. They also present a theoretical justification for using transformers in metric learning