Adaptive Mixture of Low-Rank Factorizations for Compact Neural Modeling

1 INTRODUCTION . Modern neural networks usually contain millions of parameters ( Krizhevsky et al. , 2012 ; Simonyan & Zisserman , 2014 ) , and they are difficult to be deployed on mobile devices with limited computation resources . To solve this problem , model compression techniques are proposed in recent years . For example , ( Wu et al. , 2018 ; Li et al. , 2016 ; Han et al. , 2015a ) try to limit the weights ( and activations ) of neural networks to lower precisions by quantization . This can save the model size by 4 to 16 times . While quantization can reduce the number of bits per weight , it can not reduce the number of weights . To reduce this redundancy , ( Han et al. , 2015b ; Ullrich et al. , 2017 ; Louizos et al. , 2017 ) propose pruning the weight matrices , leading to sparse neural networks that require less computation . However , sparse neural networks often require specialized ASIC or FPGA to accelerate ( Han et al. , 2016 ; 2017 ) . Alternatively , low-rank factorization is a popular way of reducing the matrix size . It has been extensively explored in the literature ( Lu et al. , 2016 ; Nakkiran et al. , 2015 ; Jaderberg et al. , 2014 ; Yu et al. , 2017 ) . Mathematically , a large weight matrix W ∈ Rm×n is factorized to two small rank-d matrices U ∈ Rm×d , V ∈ Rn×d with W = UV T . Since both U and V are dense , no sparsity support is required from specialized hardware . It naturally fits the general-purpose , off-the-shelf CPUs and GPUs . To significantly reduce the model size and computation , the rank d in the low-rank factorization needs to be small . However , a small rank can limit the expressiveness of the model ( Yang et al. , 2018 ) and lead to worse performance . To understand the limitations , given a n-dim feature vector h , we observe that V Th , as in U ( V Th ) , is a linear projection from a high-dimensional space ( n dims ) to a low-dimensional space ( d dims ) . This can lead to a significant loss of information . The conflict between the rank d and the model expressiveness prevents us from obtaining a both compact and accurate model . To address the dilemma , we propose to increase the expressiveness by learning an adaptive , inputdependent factorization , rather than performing a fixed factorization of a weight matrix . To do so , we use a mixture of multiple low-rank factorizations . The mixing weights are computed based on the input . This creates an adaptive linear projection from a high-dimensional space to a low-dimensional space . Compared to the conventional low-rank factorization , the proposed approach can significantly improve its performance while only introducing a small additional cost . To demonstrate the effectiveness of adaptive low-rank factorization , we experiment with both recurrent and convolutional neural networks on language modeling and image classification tasks . Experimental results on both tasks show that our method consistently improves upon conventional low-rank factorization . On the Penn Tree Bank dataset , we achieved 40 % reduction in FLOPs , and 1.7 better perplexity than the full rank baseline LSTM for language modeling . On ImageNet dataset , we use 48 % less computation , 12 % less parameters , but achieve 2.5 % better Top-1 accuracy than MobileNetV1 ( Howard et al. , 2017 ) . Compared to MobileNet-V2 which utilizes a standard low-rank bottleneck structure , our proposed method achieves 1.95 % better Top-1 accuracy with less than 1 % extra FLOPs , which is significant given that MobileNet-V2 is already very compact . 2 LOW-RANK FACTORIZATION AND THE LINEAR BOTTLENECK . A common linear transformation between two spaces can be represented by a linear function F : Rn → Rm , F ( h ; W ) = Wh where W ∈ Rm×n and h ∈ Rn . To reduce the size and computation of this linear transformation , a low-rank factorization of W , i.e . W = UV > , can be applied , where U ∈ Rm×d , V ∈ Rn×d , and d < min ( m , n ) . With this factorization , we can compute the transformation with h′ = UV > h . This reduces computations from O ( mn ) to O ( ( m+ n ) d ) . In the context of neural networks where W represents a weight matrix for a layer , both U and V can be learned using gradient-based algorithms . From the model compression perspective , we want to minimize the rank d , since that relates to smaller model size and computation 1 . However , the expressiveness of the transformation is limited by the rank d. By applying the factorized transformation h′ = UV > h in the reverse order , i.e . h′ = U ( V > h ) , we observe that the first transformation for h is to project it from a high-dimensional to a low-dimensional space since d < n. The latent feature distribution in high dimensional space may be either high dimensional , or lie on a non-linear manifold . In either case , projecting it into the low-dimensional space can lead to significant information loss if d is small . This is less appealing for preserving information for latter layers . To demonstrate the expressiveness issue , we conduct a toy classification task in 2D spaces . We first generate a 2D dataset with XNOR labels ( two diagonal blocks are labeled with the same class ) 1a . A non-linear classifier with one hidden layer is trained to predict the output probability by P ( y|x ) = softmax ( W2σ ( W1x ) ) where W1 ∈ R2×2 , and we factorize W1 using two 2× 1 matrices , i.e . W1 = UV T . Since the rank d = 1 , the 2D data points are first projected into 1D space and then projected into class probabilities . The visualization of one-dim space is shown in Figure 1b . We observe a significant amount of previously separated data points are now overlapped , and the class information is lost . We attribute this loss of information to the linear bottleneck . This limitation can not be solved by adding non-linear activation at the bottleneck ( after the linear transformation ) , since the information is already lost before the application of the non-linearity . Even worse , adding lossy non-linearity to a low-dimensional manifold will further negatively impact the 1The compression ratio , as well as the ratio of FLOPS , is mn/ ( ( m+ n ) d ) . A smaller d is better . network ’ s performance , as pointed out in ( Howard et al. , 2017 ) . Therefore , a new approach to boost the expressiveness of the linear bottleneck without much overhead is in demand . 3 ADAPTIVE MIXTURE OF LOW-RANK FACTORIZATIONS . To overcome the linear bottleneck in the low-rank factorization approach presented above , we propose to use an unnormalized learned mixture of low-rank factorizations whose mixing weights are computed adaptively based on the input . More specifically , denoting the input by h and the number of mixture components by K , we decompose a large weight matrix by W ( h ) = K∑ k=1 πk ( h ) U ( k ) ( V ( k ) ) > , ( 1 ) where π ( · ) : Rn → RK is the function which maps each input to its mixture coefficients , and U ( k ) ∈ Rm×d/K , V ( k ) ∈ Rn×d/K . For example , π can be a small neural network . This introduces a small amount of extra parameters and computation . We will later discuss the details of efficient ways to implement the mixture function π . If πk , k = 1 , ... , K , is chosen to be constant ( input independent ) , it can be absorbed into either U ( k ) or V ( k ) . Thus , the proposed method reduces to the low-rank factorization . This is evidenced by rewriting W as W = [ π1U ( 1 ) , ... , πKU ( K ) ] [ V ( 1 ) , ... , V ( K ) ] > . In other words , the conventional low-rank factorization can be considered as a special case of our method . Adaptive mixing weights π ( h ) . The mixing weights can encode important information that we can use to increase the expressiveness of the projected low-dimensional space . Under our framework , the generation of the mixing weights π ( h ) is flexible . A straight-forward approach is to use a non-linear transformation of the input to the weight matrix . For example , π ( h ) = σ ( Ph ) , where σ is a non-linear transformation , such as sigmoid or hyperbolic tangent function , and P ∈ RK×n is an extra weight matrix . This adds some extra parameters and computation to the model since the linear projection that we construct is Rn → RK . To further reduce the parameter and computation in the mixing weights π , we propose the following strategies . Pooling before projection . We do not require the whole input to compute the mixture function π . Instead , we can apply pooling to the input h before projection . For example , a global average pooling can be applied if the input is a 3D tensor ( for images ) ; for a 1D vector , we can segment the vector and average each segmentations . By applying pooling , we can both save the computation and better capture the global information . Random projection . To reduce the number of parameters in the linear projection of h , we can use a random matrix Prandom in place of a fully adjustable P , i.e . π ( h ) = σ ( Prandomh ) . Although we can not control freely the information captured by a random matrix , hopefully the linear projection induced by U and V can adaptively learn features that are complementary to the mixture weights . Note that we can simply save a seed to recover the random matrix , but it still requires the same amount of memory and computation as the fully adjustable linear projection of h. Increased expressiveness . Due to the use of data-dependent mixing weights for multiple lowrank factorizations , we expect the expressiveness of the model to increase . To demonstrate this intuition , in the toy example of Figure 1 , the distribution of adopting our approach that computes the mixing weights from a random matrix P and tanh non-linearity is shown in Figure 1c . Compared to conventional linear bottleneck in Figure 1b , we see a better separation among data between the two classes , leading to improved classification accuracy ( 97 % vs. 83 % ) . This is not surprising since the mixing weights encode certain class distribution in 2D space and augments the linear projection . The original overlapped projection is now better separated . More precisely , the adaptive mixing weights introduce a non-linear transformation into the highto-low-dimensional projection that can be more expressive . Since each W ( h ) is a data-dependent low-rank matrix , there is no constant linear weight independent to the input ( even a full-rank matrix ) that can mimic the transformation W ( h ) induced by our proposed method . It is worth noting that generating the whole weight matrices can be very expensive . Our method can be seen as a swift approach to generate the weights by adaptively adjusting mixing weights for the linear bottleneck . It assigns weights into groups and dynamically controls them at the group level , as demonstrated in Figure 2 . To efficiently compute the whole linear transformation ∑K k=1 πk ( h ) U ( k ) ( V ( k ) ) > h , we use the reverse order , i.e . first computing the linear projection into low-dimensional space with mixing weights , i.e . zk = πk ( h ) ( ( V ( k ) ) Th ) , and then map to a higher dimensional space , i.e . h′ =∑ k U ( k ) zk . This reduces the FLOPs and also avoids to store different weight matrices W ( h ) for different training examples in a mini-batch . An illustration of the computation framework is presented in Figure 3 . Compared to original low-rank factorization , extra parameters and computation cost are from the mixing weights . They can be very small with techniques like pooling aforementioned . We also introduce a way in appendix to avoid breaking a bulk matrix multiplication to segments .

Summary: The paper proposes an input-dependent low rank approximations of weight matrices in neural nets with a goal of accelerating inference. Instead of decomposing a matrix (a layer) to a low rank space like previous work did, the paper proposes to use a linear combination/mixture of multiple low rank spaces. The linear combination is dynamic and data dependent, making it some sense of non-linear combination. The paper is interesting, however,