Minimizing FLOPs to Learn Efficient Sparse Representations

1 INTRODUCTION . Learning semantic representations using deep neural networks ( DNN ) is now a fundamental facet of applications ranging from visual search ( Jing et al. , 2015 ; Hadi Kiapour et al. , 2015 ) , semantic text matching ( Neculoiu et al. , 2016 ) , oneshot classification ( Koch et al. , 2015 ) , clustering ( Oh Song et al. , 2017 ) , and recommendation ( Shankar et al. , 2017 ) . The high-dimensional dense embeddings generated from DNNs however pose a computational challenge for performing nearest neighbor search in large-scale problems with millions of instances . In particular , when the embedding dimension is high , evaluating the distance of any query to all the instances in a large database is expensive , so that efficient search without sacrificing accuracy is difficult . Representations generated using DNNs typically have a higher dimension compared to hand-crafted features such as SIFT ( Lowe , 2004 ) , and moreover are dense . The key caveat with dense features is that unlike bag-of-words features they can not be efficiently searched through an inverted index , without approximations . Since accurate search in high dimensions is prohibitively expensive in practice ( Wang , 2011 ) , one has to typically sacrifice accuracy for efficiency by resorting to approximate methods . Addressing the problem of efficient approximate Nearest-Neighbor Search ( NNS ) ( Jegou et al. , 2011 ) or Maximum Inner-Product Search ( MIPS ) ( Shrivastava and Li , 2014 ) is thus an active area of research , which we review in brief in the related work section . Most approaches ( Charikar , 2002 ; Jegou et al. , 2011 ) aim to learn compact lower-dimensional representations that preserve distance information . While there has been ample work on learning compact representations , learning sparse higher dimensional representations have been addressed only recently ( Jeong and Song , 2018 ; Cao et al. , 2018 ) . As a seminal instance , Jeong and Song ( 2018 ) propose an end-to-end approach to learn ∗Part of the work was done when BP was a research intern at Snap Inc. 1The implementation is available at https : //github.com/biswajitsc/sparse-embed sparse and high-dimensional hashes , showing significant speed-up in retrieval time on benchmark datasets compared to dense embeddings . This approach has also been motivated from a biological viewpoint ( Li et al. , 2018 ) by relating to a fruit fly ’ s olfactory circuit , thus suggesting the possibility of hashing using higher dimensions instead of reducing the dimensionality . Furthermore , as suggested by Glorot et al . ( 2011 ) , sparsity can have additional advantages of linear separability and information disentanglement . In a similar vein , in this work , we propose to learn high dimensional embeddings that are sparse and hence efficient to retrieve using sparse matrix multiplication operations . In contrast to compact lowerdimensional ANN-esque representations that typically lead to decreased representational power , a key facet of our higher dimensional sparse embeddings is that they can have the same representational capacity as the initial dense embeddings . The core idea behind our approach is inspired by two key observations : ( i ) retrieval of d ( high ) dimensional sparse embeddings with fraction p of non-zero values on an average , can be sped up by a factor of 1/p . ( ii ) The speed up can be further improved to a factor of 1/p2 by ensuring that the non-zero values are evenly distributed across all the dimensions . This indicates that sparsity alone is not sufficient to ensure maximal speedup ; the distribution of the non-zero values plays a significant role as well . This motivates us to consider the effect of sparsity on the number of floating point operations ( FLOPs ) required for retrieval with an inverted index . We propose a penalty function on the embedding vectors that is a continuous relaxation of the exact number of FLOPs , and encourages an even distribution of the non-zeros across the dimensions . We apply our approach to the large scale metric learning problem of learning embeddings for facial images . Our training loss consists of a metric learning ( Weinberger and Saul , 2009 ) loss aimed at learning embeddings that mimic a desired metric , and a FLOPs loss to minimize the number of operations . We perform an empirical evaluation of our approach on the Megaface dataset ( Kemelmacher-Shlizerman et al. , 2016 ) , and show that our proposed method successfully learns high-dimensional sparse embeddings that are orders-of-magnitude faster . We compare our approach to multiple baselines demonstrating an improved or similar speed-vs-accuracy trade-off . The rest of the paper is organized as follows . In Section 3 we analyze the expected number of FLOPs , for which we derive an exact expression . In Section 4 we derive a continuous relaxation that can be used as a regularizer , and optimized using gradient descent . We also provide some analytical justifications for our relaxation . In Section 5 we then compare our method on a large metric learning task showing an improved speed-accuracy trade-off compared to the baselines . 2 RELATED WORK . Learning compact representations , ANN . Exact retrieval of the top-k nearest neighbours is expensive in practice for high-dimensional dense embeddings learned from deep neural networks , with practitioners often resorting to approximate nearest neighbours ( ANN ) for efficient retrieval . Popular approaches for ANN include Locality sensitive hashing ( LSH ) ( Gionis et al. , 1999 ; Andoni et al. , 2015 ; Raginsky and Lazebnik , 2009 ) relying on random projections , Navigable small world graphs ( NSW ) ( Malkov et al. , 2014 ) and hierarchical NSW ( HNSW ) ( Malkov and Yashunin , 2018 ) based on constructing efficient search graphs by finding clusters in the data , Product Quantization ( PQ ) ( Ge et al. , 2013 ; Jegou et al. , 2011 ) approaches which decompose the original space into a cartesian product of low-dimensional subspaces and quantize each of them separately , and Spectral hashing ( Weiss et al. , 2009 ) which involves an NP hard problem of computing an optimal binary hash , which is relaxed to continuous valued hashes , admitting a simple solution in terms of the spectrum of the similarity matrix . Overall , for compact representations and to speed up query times , most of these approaches use a variety of carefully chosen data structures , such as hashes ( Neyshabur and Srebro , 2015 ; Wang et al. , 2018 ) , locality sensitive hashes ( Andoni et al. , 2015 ) , inverted file structure ( Jegou et al. , 2011 ; Baranchuk et al. , 2018 ) , trees ( Ram and Gray , 2012 ) , clustering ( Auvolat et al. , 2015 ) , quantization sketches ( Jegou et al. , 2011 ; Ning et al. , 2016 ) , as well as dimensionality reductions based on principal component analysis and t-SNE ( Maaten and Hinton , 2008 ) . End to end ANN . Learning the ANN structure end-to-end is another thread of work that has gained popularity recently . Norouzi et al . ( 2012 ) propose to learn binary representations for the Hamming metric by minimizing a margin based triplet loss . Erin Liong et al . ( 2015 ) use the signed output of a deep neural network as hashes , while imposing independence and orthogonality conditions on the hash bits . Other end-to-end learning approaches for learning hashes include ( Cao et al. , 2016 ; Li et al. , 2017 ) . An advantage of end-to-end methods is that they learn hash codes that are optimally compatible to the feature representations . Sparse representations . Sparse representations have been previously studied from various viewpoints . Glorot et al . ( 2011 ) explore sparse neural networks in modeling biological neural networks and show improved performance , along with additional advantages such as better linear separability and information disentangling . Ranzato et al . ( 2008 ; 2007 ) ; Lee et al . ( 2008 ) propose learning sparse features using deep belief networks . Olshausen and Field ( 1997 ) explore sparse coding with an overcomplete basis , from a neurobiological viewpoint . Sparsity in auto-encoders have been explored by Ng et al . ( 2011 ) ; Kavukcuoglu et al . ( 2010 ) . Arpit et al . ( 2015 ) provide sufficient conditions to learn sparse representations , and also further provide an excellent review of sparse autoencoders . Dropout ( Srivastava et al. , 2014 ) and a number of its variants ( Molchanov et al. , 2017 ; Park et al. , 2018 ; Ba and Frey , 2013 ) have also been shown to impose sparsity in neural networks . High-dimensional sparse representations . Sparse deep hashing ( SDH ) ( Jeong and Song , 2018 ) is an end-to-end approach that involves starting with a pre-trained network and then performing alternate minimization consisting of two minimization steps , one for training the binary hashes and the other for training the continuous dense embeddings . The first involves computing an optimal hash best compatible with the dense embedding using a min-cost-max-flow approach . The second step is a gradient descent step to learn a dense embedding by minimizing a metric learning loss . A related approach , k-sparse autoencoders ( Makhzani and Frey , 2013 ) learn representations in an unsupervised manner with at most k non-zero activations . The idea of high dimensional sparse embeddings is also reinforced by the sparse-lifting approach ( Li et al. , 2018 ) where sparse high dimensional embeddings are learned from dense features . The idea is motivated by the biologically inspired fly algorithm ( Dasgupta et al. , 2017 ) . Experimental results indicated that sparse-lifting is an improvement both in terms of precision and speed , when compared to traditional techniques like LSH that rely on dimensionality reduction . ` 1 regularization , lasso . The Lasso ( Tibshirani , 1996 ) is the most popular approach to impose sparsity and has been used in a variety of applications including sparsifying and compressing neural networks ( Liu et al. , 2015 ; Wen et al. , 2016 ) . The group lasso ( Meier et al. , 2008 ) is an extension of lasso that encourages all features in a specified group to be selected together . Another extension , the exclusive lasso ( Kong et al. , 2014 ; Zhou et al. , 2010 ) , on the other hand , is designed to select a single feature in a group . Our proposed regularizer , originally motivated by idea of minimizing FLOPs closely resembles exclusive lasso . Our focus however is on sparsifying the produced embeddings rather than sparsifying the parameters . Sparse matrix vector product ( SpMV ) . Existing work for SpMV computations include ( Haffner , 2006 ; Kudo and Matsumoto , 2003 ) , proposing algorithms based on inverted indices . Inverted indices are however known to suffer from severe cache misses . Linear algebra back-ends such as BLAS ( Blackford et al. , 2002 ) rely on efficient cache accesses to achieve speedup . Haffner ( 2006 ) ; Mellor-Crummey and Garvin ( 2004 ) ; Krotkiewski and Dabrowski ( 2010 ) propose cache efficient algorithms for sparse matrix vector products . There has also been substantial interest in speeding up SpMV computations using specialized hardware such as GPUs ( Vazquez et al. , 2010 ; Vázquez et al. , 2011 ) , FPGAs ( Zhuo and Prasanna , 2005 ; Zhang et al. , 2009 ) , and custom hardware ( Prasanna and Morris , 2007 ) . Metric learning . While there exist many settings for learning embeddings ( Hinton and Salakhutdinov , 2006 ; Kingma and Welling , 2013 ; Kiela and Bottou , 2014 ) in this paper we restrict our attention to the context of metric learning ( Weinberger and Saul , 2009 ) . Some examples of metric learning losses include large margin softmax loss for CNNs ( Liu et al. , 2016 ) , triplet loss ( Schroff et al. , 2015 ) , and proxy based metric loss ( Movshovitz-Attias et al. , 2017 ) .

Summary: This paper proposes to learn sparse representation in neural networks for retrieval in large database of vectors. Such sparse representation, when the fraction of non-zeros is high, can be computed using sparse matrix multiplication, or variants of inverted index scoring and lead to potentially lower FLOPs needed. This paper proposes to induce sparsity by adding a regularization term, which counts the expected number of FLOPs needed for sparse scoring.