Contrastive Mutual Information Maximization for Binary Neural Networks

1 INTRODUCTION . Although deep neural networks ( DNNs ) [ 1 ] have achieved remarkable success in various computer vision tasks such as image classification [ 2 ] and semantic image segmentation [ 3 ] , their over-parametrization problem makes them computationally expensive and storage excessive . To advance the development of deep learning towards resource-constrained edge devices , researchers proposed several neural network compression paradigms , such as network pruning [ 4 , 5 ] , knowledge distillation [ 6 ] and network quantization [ 7 , 8 ] . Among the network quantization methods , the network binarization method [ 7 ] stands out for quantizing weights and activations ( i.e . intermediate feature maps ) to ±1 , compressing the full-precision counterpart 32× , and replacing time-consuming inner-product in full-precision networks with efficient xnor-bitcount operation in the BNNs . However , severe accuracy drops always exist between full-precision models and their binary counterparts . To tackle this problem , previous works mainly focus on reducing the quantization error induced by weights binarization [ 9 , 10 ] , and elaborately approximating binarization function to relieve the gradient mismatch in the backward propagation [ 11 , 8 ] . Indeed , they achieve the state-of-the-art performance . Yet with those two paradigms developing , narrowing down the quantization error and enhancing the gradient transmission reach their bottlenecks [ 12 , 13 ] , since the 1W32A ( only quantizing the weights into 1-bit , remaining the activations 32-bit ) models are capable of performing as well as the full-precision models , implying that the activations binarization becomes the main issue for further performance improvement . To address the accuracy degradation caused by the activations binarization , a few studies propose to regulate the distributions of the binary activations , e.g . researchers in [ 14 ] design a distribution loss to explicitly regularize the activation flow ; researchers in [ 13 ] propose to shift the thresholds of binary activation functions to make the distribution of binary activation unbalanced . They heuristically design low-level patterns to analyze the distributions of binary activations such as minimum of the activations and the balanced property of distributions . Nevertheless , they neglect the high-level indicators of the distribution and the unique characteristics of BNN , where the binary activations and latent full-precision activations exist in the same forward pass . Thus , we argue that the high-level properties of distributions , such as correlations and dependencies between binary and full-precision activations should be captured and utilized . In this work , we explore introducing mutual information for BNNs , in which the mutual information acts as a fundamental quantity to measure the information amount shared by the binary and latent real-valued activations in BNNs . In contrast to the aforementioned works focusing on learning the distribution of binary activations , mutual information naturally captures non-linear statistical dependencies between variables , and thus can be used as a measure of true dependence [ 15 ] . Based on this metric , we propose a novel method , termed as Contrastive Mutual Information Maximization for Binary Neural Networks ( CMIM-BNN ) . Specifically , we design a highly effective optimization strategy using contrastive estimation for the mutual information maximization . As illustrated in Figure 1 , we replace the data transformation module in contrastive learning with the exclusive structure in BNNs , where full-precision and binary activations are in the same forward pass . In this way , contrastive learning contributes to inter-class decorrelation of binary activations , and avoids collapse solutions . In other words , our method is built upon a contrastive learning framework to learn representative binary activations , in which we pull the binary activation closer to the full-precision activation and push the binary activation further away from other binary activations in the contrastive space . Moreover , by utilizing an additional MLP module to extract representations of activations , our method can explicitly capture higher-order dependencies in the contrastive space . To the best of our knowledge , it is the first work aiming at maximizing the mutual information of the activations in BNNs within a contrastive learning framework . Overall , the contributions of this paper are three-fold : • Considering the distributions of activations , we propose a novel CMIM framework to optimize BNNs , by maximizing the mutual information between the binary activation and its latent realvalued counterpart ; • We develop an effective contrastive learning strategy to achieve the goal of mutual information maximization for BNNs , and benefited from it , the representation ability of BNNs is clearly strengthened for not only the classification task but also downstream CV tasks ; • Experimental results show that our method can significantly improve the existing SOTA methods over the classification task on CIFAR-10/100 and ImageNet , e.g . 6.4 % on CIFAR-100 and 3.0 % on ImageNet . Besides , we also demonstrate good generalization ability of the proposed CMIM on other challenging CV tasks such as depth estimation and semantic segmentation . 2 MUTUAL INFORMATION MAXIMIZATION FOR TRAINING BNNS . 2.1 PRELIMINARIES . We first define a Multi-Layer Perceptron ( MLP ) with K layers . For simplification of derivation , we discard the bias term of the network . Then the MLP f ( x ) can be denoted as : f ( W1 , · · · , WK ; x ) = ( WK · σ ·WK−1 · · · · · σ ·W1 ) ( x ) , ( 1 ) where x is the input sample and Wk : Rdk−1 7−→ Rdk ( k = 1 , ... , K ) stands for the weight matrix connecting the ( k− 1 ) -th and the k-th layer , with dk−1 and dk representing the sizes of the input and output of the k-th network layer , respectively . The σ ( · ) function performs element-wise activation for the feature maps . Notably , for a convolution layer with the input map of m channels and the output map of n channels , and the size of the kernel w× h , it results in m× n×w× h parameters . We can re-arrange the parameters to a weight matrix of size n × ( m × h × w ) , such that this convolution layer can also operate in the same way as the other fully-connected layers . Hence , it is sufficient to consider networks with the fully-connected layers . Based on those predefined notions , the sectional MLP fk ( x ) with the front k layers of the f ( x ) can be represented as : fk ( W1 , · · · , Wk ; x ) = ( Wk · σ · · ·σ ·W1 ) ( x ) . ( 2 ) And the MLP f can be seen as a special case in the function sequence { fk } ( k ∈ { 1 , · · · , K } ) , i.e . f = fK , when k = K. Binary Neural Networks . Here , we revisit the general binarization method in [ 16 , 7 ] , which maintains full-precision latent variables WF for gradient updates , and the k-th weight matrix WkF is binarized into ±1 , obatining the binary weight matrix WkB by a binarize function ( normally sgn ( · ) ) , i.e . WkB = sgn ( WkF ) . Then the intermediate activation map ( full-precision ) of the k-th layer is produced by AkF = W k BA k−1 B , then the same quantization method is used to binarize the full-precision activation map as AkB = sgn ( A k F ) , and a whole forward pass of binarization is performed by iterating this process for L times . Mutual Information . For two discrete variables X and Y , their mutual information can be defined as [ 17 ] : I ( X , Y ) = ∑ x , y PXY ( x , y ) log PXY ( x , y ) PX ( x ) PY ( y ) , ( 3 ) where PXY ( x , y ) is the joint distribution , PX ( x ) = ∑ y PXY ( x , y ) and PY ( y ) = ∑ x PXY ( x , y ) are the marginals of X and Y , respectively . Mutual information quantifies the amount of information obtained about one random variable by observing the other random variable . It is a dimensionless quantity with ( generally ) units of bits , and can be thought as the reduction in uncertainty about one random variable given knowledge of another . High mutual information indicates a large reduction in uncertainty ; low mutual information indicates a small reduction ; and zero mutual information between two random variables means the variables are independent . In the content of binarization , considering the binary and full-precision activations as random variables , we would like them share as much information as possible , since the binary activations are proceeded from their corresponding full-precision activations . Theoretically , the mutual information between those two variables should be maximized . 2.2 CONTRASTIVE MUTUAL INFORMATION MAXIMIZATION . In the following section , we formalize the idea of constructing a Noise-Contrastive Estimation ( NCE ) loss to maximize the mutual information between the binary and the full-precision activations . Particularly , we derive a novel CMIM loss for training BNNs , where NCE is introduced to avoid the direct mutual information computation by estimating it with its lower bound in Eq . 7 . For binary network fB and its latent full-precision counterpart fF in the same training iteration , the series of their activations { akB } and { akF } ( k ∈ { 1 , · · · , K } ) , where AkB = ( a k,1 B , · · · , a k , N B ) and AkF = ( a k,1 F , · · · , a k , N F ) can be considered as a series of variables . The corresponding variables ( akB , a k F ) should share more information , i.e . the mutual information of the same layer ’ s output activations I ( akB , a k F ) ( k ∈ { 1 , · · · , K } ) should be maximized to enforce them mutually dependent . To this end , we introduce the contrastive learning framework into our targeted binarization task . The basic idea of contrastive learning is to compare different views of the data ( usually under different data augmentations ) to calculate similarity scores [ 18 , 19 , 20 , 21 , 22 ] . This framework is suitable for our case , since the binary and full-precision activations can be seen as two different views . For a training batch with N samples , the samples can be denoted as : { xi } ( i ∈ { 1 , · · · , N } ) . We feed a batch of samples to the BNN and obtain KN2 pairs of activations ( ak , iB , a k , j F ) , which augments the data for the auxiliary task . We define a pair containing two activations from the same sample as positive pair , i.e . if i = j , ( ak , iB , a k , j F ) + and vice versa . With the Bayes ’ theorem , the posterior probability of two activations from the positive pair can be formalized as : P ( i = j | ak , iB , a k , j F ) = P ( ak , iB , a k , j F | i = j ) 1 N P ( ak , iB , a k , j F | i = j ) 1 N + P ( a k , i B , a k , j F | i 6= j ) N−1 N . ( 4 ) And the probability of activations from negative pair is : P ( i 6= j | ak , iB , a k , j F ) = 1 − P ( i = j | ak , iB , a k , j F ) . To simplify the NCE derivative , several works [ 23 , 24 , 25 ] build assumption about the dependence of the variables , we also use this assumption that the activations from positive pairs are dependent and the ones from negative pairs are independent , i.e . P ( ak , iB , a k , j F | i = j ) = P ( a k , i B , a k , j F ) and P ( ak , iB , a k , j F | i 6= j ) = P ( a k , i B ) P ( a k , j F ) . Hence , the above equation can be simplified as : P ( i = j | ak , iB , a k , j F ) = P ( ak , iB , a k , j F ) P ( ak , iB , a k , j F ) + P ( a k , i B ) P ( a k , j F ) ( N − 1 ) . ( 5 ) Performing logarithm to Eq . 5 and arranging the terms , we can achieve logP ( i = j | ak , iB , a k , j F ) = − log [ 1 + ( N − 1 ) P ( ak , iB ) P ( a k , j F ) P ( ak , iB , a k , j F ) ] ≤ log P ( ak , iB , a k , j F ) P ( ak , iB ) P ( a k , j F ) − log ( N − 1 ) . ( 6 ) Taking expectation on both sides with respect to P ( ak , iB , a k , j F ) , and combining the definition of mutual information in Eq . 3 , we can derive the form of mutual information as : targeted mutual information︷ ︸︸ ︷ I ( akB , a k F ) = ∑ i ∑ j P ( ak , iB , a k , j F ) log P ( ak , iB , a k , j F ) P ( ak , iB ) P ( a k , j F ) ≥ ∑ i ∑ j P ( ak , iB , a k , j F | i = j ) [ logP ( i = j | ak , iB , a k , j F ) + log ( N − 1 ) ] = optimized lower bound︷ ︸︸ ︷ EP ( ak , iB , ak , jF |i=j ) [ logP ( i = j | ak , iB , a k , j F ) ] + log ( N − 1 ) , ( 7 ) where I ( akB , a k F ) is the mutual information between the binary and full-precision distributions , our targeted object . Instead of directly maximizing the mutual information , we choose to maximize its lower bound in the Eq . 7 . However , the distribution P ( i = j | ak , iB , a k , j F ) is hard to estimate . We take advantage of the idea of contrastive learning , and introduce a critic function h to approximate the targeted distribution [ 18 , 19 , 20 ] . In practice , we use the following : h ( ak , iB , a k , j F ) = exp ( τ ( ak , iB ) > ak , jF ) ∑ j exp ( τ ( a k , i B ) > ak , jF ) ( 8 ) Algorithm 1 Forward and Backward Propagation of CMIM Require : A minibatch of data samples ( X , Y ) , current binary weight WkB , latent fullprecision weights WkF , and learning rate η . Ensure : Update weights WkF ′ . 1 : Forward Propagation : 2 : for k = 1 to K − 1 do 3 : Binarize latent weights : WkB ←− sgn ( WkF ) ; 4 : Perform binary operation with the activations of next layer : AkF ←− XnorDotProduct ( WkB , A k−1 B ) ; 5 : Perform Batch Normalization : AkF ←− BatchNorm ( AkF ) ; 6 : Binarize full-precision activations and obtain binary activations : AkB ←− sgn ( AkF ) ; 7 : end for 8 : For k = 1 , · · · , K , pair { ak , iB } and { ak , jB } as negative and positive pairs , then use Eq . 9 layer by layer to compute the NCE loss LkNCE between AkB and AkF for contrastive learning ; 9 : Combine a series of NCE loss { LkNCE } with the classification loss L into the CMIM loss LCMIM , with Eq . 11 ; 10 : Backward Propagation : compute the gradient of the overall loss function , i.e . ∂L∂WB , using the straight through estimator to tackle the sign function ; 11 : Parameter Update : update the fullprecision weights : WiF ′ ←−WkF − η ∂L∂WkB . where τ is a temperature parameter that controls the concentration level of the distribution [ 6 ] . τ is important for supervised feature learning , and also necessary for tuning the concentration of ak , iB and ak , jF on our contrastive space . Loss Function . We define the contrastive loss function LkNCE between the k-th layer ’ s activations AkB and A k F as : LkNCE = EP ( ak , iB , ak , jF |i=j ) [ log h ( ak , iB , a k , j F ) ] +NEP ( ak , iB , ak , jF |i6=j ) [ log ( 1− h ( ak , iB , a k , j F ) ) ] . ( 9 ) We would like to comment on the above loss function from the perspective of contrastive learning . The first term of positive pairs is optimized for capturing more intra-class correlations and the second term of negative pairs is for inter-class decorrelation . Because the pair construction is instance-wise , the number of negative samples theoretically can be the size of the entire training set , e.g . 1.2 million for ImageNet . With those additional hand-craft designed contrastive pairs for the proxy optimization problem in Eq . 9 , the representation capacity of BNNs can be further improved , as many contrastive learning methods demonstrated [ 22 , 18 , 19 , 20 ] . Moreover , the optimal ĥ = argmaxh LkNCE can approximate the targeted distribution , i.e . ĥ ( ak , iB , a k , j F ) = P ( i = j | a k , i B , a k , j F ) , ( 10 ) where the detailed proof is shown in the supplementary material . Thus , with Eq . 7-10 we can derive that minimizing the NCE loss LkNCE is equivalent to maximizing the targeted mutual information between the binary and full-precision activations , I ( akB , a k F ) . Combining the series of NCE loss from different layers { LkNCE } , ( k = 1 , · · · , K ) , the overall loss L can be defined as : L = λ K∑ k=1 LkNCE βK−1−k + Lcls , ( 11 ) where Lcls is the classification loss respect to the ground truth , λ is used to control the degree of NCE loss , β is a coefficient greater than 1 , and we denote the CMIM loss as LCMIM = ∑K k=1 LkNCE βK−1−k . Hence , the βK−1−k decreases with k increasing and consequently the L k NCE βK−1−k increases . In this way , the activations of latter layer can be substantially retained , which leads to better performance in practice . The complete training process of CMIM is presented in Algorithm 1 . Discussion on the CMIM Loss . Besides the theoretical formulation from the perspective of mutual information maximization , we also provide an intuitive explanation about CMIM . As illustrated in Figure 2 , we strengthen the representation ability of binary activations ( Figure 3 ) via designing a proxy task with the contrastive learning framework . By embedding the activations to the contrastive space and pull-and-push the paired embeddings , the BNNs can learn better representations from this difficult yet effective auxiliary contrastive learning task . Note that even though we only pick up two images to formulate Figure 2 , the actual number of negative samples can be huge in practice ( e.g . 16,384 for training ResNet-18 on ImageNet ) , benefit from the MemoryBank [ 24 ] technique . With this property , we speculate that the contrastive pairing works as the data augmentation , which contributes to our method . This additional pairing provides more information for training the BNNs , thus our CMIM loss can be treated as an overfitting-mitigated module . We also conduct experiments in the Section 3.2 and 3.3 to validate our speculation . Difference with other contrastive learning methods . The key idea of contrastive learning is to pull representations close in positive pairs and push representations apart in negative pairs in a contrastive space . Several self-supervised learning methods are rooted in well-established idea of the mutual information maximization , such as Deep InfoMax [ 19 ] , Contrastive Predictive Coding [ 18 ] , MemoryBank [ 24 ] , Augmented Multiscale DIM [ 20 ] , MoCo [ 21 ] and SimSaim [ 22 ] . These are based on NCE [ 23 ] and InfoNCE [ 19 ] which can be seen as a lower bound on mutual information [ 27 ] . The formulation of our CMIM-BNN is similar to the classic contrastive learning methods , where we all are inspired by NCE . However , our approach has several differences from those methods . Firstly , the training process of BNNs is different from regular network training , where binary and latent full-precision activations exist in the same forward pass . We seamlessly integrate this mixedactivation property with NCE , and thus the targeted lower bound formulated for optimization is different . Secondly , the binary and full-precision weights are both optimized by the NCE loss ( i.e . the two view augmentation networks are optimized simultaneously in contrastive learning ) , yet most aforementioned contrastive learning methods optimize their view augmentation networks separately .

Summary: This paper proposes to use contrastive distillation to train a binary network by maximizing the mutual information between itself (student network) and the full-precision network (teacher network). Empirical results show that this new objective further improves the binarization performance on top of several recent binary networks on image classification tasks. The authors also empirically show that models trained with the proposed contrastive objective have good transfer performance.