Simple Augmentation Goes a Long Way: ADRL for DNN Quantization

1 INTRODUCTION . By reducing the number of bits needed to represent a model parameter of Deep Neural Networks ( DNN ) , quantization ( Lin et al. , 2016 ; Park et al. , 2017 ; Han et al. , 2015 ; Zhou et al. , 2018 ; Zhu et al. , 2016 ; Hwang & Sung , 2014 ; Wu et al. , 2016 ; Zhang et al. , 2018 ; Köster et al. , 2017 ; Ullrich et al. , 2017 ; Hou & Kwok , 2018 ; Jacob et al. , 2018 ) is an important way to reduce the size and improve the energy efficiency and speed of DNN . Mixed precision quantization selects a proper bit-width for each layer of a DNN , offering more flexibility than fixed precision quantization . A major challenge to mixed precision quantization ( Micikevicius et al. , 2017 ; Cheng et al. , 2018 ) is the configuration search problem , that is , how to find the appropriate bit-width for each DNN layer efficiently . The search space grows exponentially as the number of layers increases , and assessing each candidate configuration requires a long time of training and evaluation of the DNN . Research efforts have been drawn to mitigate the issue for help better tap into the power of mixed precision quantization . Prior methods mainly fall into two categories : ( i ) automatic methods , such as reinforcement learning ( RL ) ( Lou et al. , 2019 ; Gong et al. , 2019 ; Wang et al. , 2018 ; Yazdanbakhsh et al. , 2018 ; Cai et al. , 2020 ) and neural architecture search ( NAS ) ( Wu et al. , 2018 ; Li et al. , 2020 ) , to learn from feedback signals and automatically determine the quantization configurations ; ( ii ) heuristic methods to reduce the search space under the guidance of metrics such as weight loss or Hessian spectrum ( Dong et al. , 2019 ; Wu et al. , 2018 ; Zhou et al. , 2018 ; Park et al. , 2017 ) of each layer . Comparing to the heuristic method , automatic methods , especially Deep Reinforcement Learning ( DRL ) , require little human effort and give the state-of-the-art performance ( e.g. , via actor-critic set∗This work was done when Lin Ning was an intern at Alibaba Group . †Now at Google Research . ting ( Whiteson et al. , 2011 ; Zhang et al. , 2016 ; Henderson et al. , 2018 ; Wang et al. , 2018 ) ) . It however suffers from overestimation bias , high variance of the estimated value , and hence slow convergence and suboptimal results . The problem is fundamentally due to the poor function approximations given out by the DRL agent , especially during the early stage of the DRL learning process ( Thrun & Schwartz , 1993 ; Anschel et al. , 2017 ; Fujimoto et al. , 2018 ) when the neural networks used in the DRL are of low quality . The issue prevents DRL from serving as a scalable solution to DNN quantization as DNN becomes deeper and more complex . This paper reports that simple augmentations can bring some surprising improvements to DRL for DNN quantization . We introduce augmented DRL ( ADRL ) as a principled way to significantly magnify the potential of DRL for DNN quantization . The principle of ADRL is to augment the neural networks in DRL with a complementary scheme ( called augmentation scheme ) to complement the weakness of DRL policy approximator . Analytically , we prove the effects of such a method in reducing the variance and improving the convergence rates of DRL . Empirically , we exemplify ADRL with two example augmentation schemes and test on four popular DNNs . Comparisons with four prior DRL methods show that ADRL can shorten the quantization process by 4.5-64× while improving the model accuracy substantially . It is worth mentioning that there is some prior work trying to increase the scalability of DRL . Dulac-Arnold et al . ( 2016 ) , for instance , addresses large discrete action spaces by embedding them into continuous spaces and leverages nearest neighbor to find closest actions . Our focus is different , aiming to enhance the learning speed of DRL by augmenting the weak policy approximator with complementary schemes . 2 BACKGROUND . Deep Deterministic Policy Gradient ( DDPG ) A standard reinforcement learning system consists of an agent interacting with an environment E . At each time step t , the agent receives an observation xt , takes an action at and then receives an award rt . Modeled with Markov decision process ( MDP ) with a state space S and an action space A , an agent ’ s behavior is defined by a policy π : S → A . A state is defined as a sequence of actions and observations st = ( x1 , a2 , · · · , at−1 , xt ) when the environment is ( partially ) observed . For DNN quantization , the environment is assumed to be fully observable ( st = xt ) . The return from a state s at time t is defined as the future discounted return Rt = ∑T i=t γ i−tr ( si , ai ) with a discount factor γ . The goal of the agent is to learn a policy that maximizes the expected return from the start state J ( π ) = E [ R1|π ] . An RL agent in continuous action spaces can be trained through the actor-critic algorithm and the deep deterministic policy gradient ( DDPG ) . The parameterized actor function µ ( s|θµ ) specifies the current policy and deterministically maps a state s to an action a . The critic networkQ ( s , a ) is a neural network function for estimating the action-value E [ Rt|st = s , at = a , π ] ; it is parameterized with θQ and is learned using the Bellman equation as Q-learning . The critic is updated by minimizing the loss L ( θQ ) = E [ ( yt −Q ( st , at|θQ ) ) 2 ] , where yt = r ( st , at ) + γQ ( st+1 , µ ( st+1|θµ ) |θQ ) . ( 1 ) The actor is updated by applying the chain rule to the expected return J with respect to its parameters : ∇θµJ ≈ E [ ∇aQ ( s , a|θQ ) |s=st , a=µ ( st ) ∇θµµ ( s|θ µ ) |s=st ] . ( 2 ) DDPG for Mixed Precision Quantization To apply the DRL to mixed precision quantization , previous work , represented by HAQ ( Wang et al. , 2018 ) , uses DDPG as the agent learning policy . The environment is assumed to be fully observed so that st = xt , where the observation xt is defined as xt = ( l , cin , cout , skernel , sstride , sfeat , nparams , idw , iw/a , at−1 ) for convolution layers and xt = ( l , hin , hout , 1 , 0 , sfeat , nparams , 0 , iw/a , at−1 ) for fully connected layers . Here , l denotes the layer index , cin and cout are the numbers of input and output channels for the convolution layer , skernel and sstride are the kernel size and stride for the convolution layer , hin and hout are the numbers of input and output hidden units for the fully connected layer , nparams is the number of parameters , idw and iw/a are binary indicators for depth-wise convolution and weight/activation , and at−1 is the action given by the agent from the previous step . In time step t− 1 , the agent gives an action at−1 for layer l − 1 , leading to an observation xt . Then the agent gives the action at for layer l at time step t given xt . The agent updates the actor and critic networks after one episode following DDPG , which is a full pass of all the layers in the target neural network for quantization . The time step t and layer index l are interchangeable in this scenario . The systems use a continuous action space for precision prediction . At each time step t , the continuous action at is mapped to the discrete bit value bk for layer k using bk = round ( bmin − 0.5 + ak × ( bmax − bmin + 1 ) ) . The reward function is computed usingR = λ× ( accquant − accorigin ) . 3 AUGMENTED DEEP REINFORCEMENT LEARNING ( ADRL ) . This section explains our proposal , ADRL . It builds upon the default actor-critic algorithm and is trained with DDPG . In the original actor-critic algorithm , the policy approximator is the actor networks , which generates one action and feeds it into the critic networks . The essence of ADRL is to augment the policy approximator with a supplementary scheme . The scheme may be in an arbitrary form , constructed with domain knowledge or other mechanisms . It should boost the weak policy approximator in accuracy especially when it is not yet well trained , and hence help efficiently reduce the variance and accelerate the convergence of the DRL algorithm . 3.1 DESIGN . Components of ADRL Figure 1 illustrates the general flow of ADRL that employs post augmentation . The augmented policy approximator consists of two main components : an expanded actor network and a refinement function . The expansion of the actor network makes it output multiple rather than one candidate action ; the refinement function derives the most promising candidate by processing those actions and feeds it to the critic networks . The two components can be formally described as follows . Actions generation The expanded actor function , µ̂θµ̂ : S → Rk×n , maps a state from the state space Rm to k ( k > 1 ) actions in a continuous actions space Rn : Ak = µ̂ ( s|θµ̂ ) where Ak = [ aT1 , · · · , aTk ] T . The expansion can be done by modifying the last layer of the actor network ( Sec 3.3 ) . The outputs of this actor function serve as the candidate actions to the refinement component . Action refinement . The action refinement function , g : Rk×n → Rn , derives a promising action from the candidate actions . A simple form of derivation is selection , that is , a∗ = arg maxa=Ak , i Q ( s , a ) Depending on how well the critic network is trained , the critic may not be able to consistently give a good estimation of Q ( s , a ) , especially at the early training stage . The augmented policy may use a Q-value indicator Q̃E ( a ) whose output depends on only the action and the property of the environment . Hence , g ( Ak ) = arg maxa=Ak , i Q̃E ( a ) The choice of Q̃E ( a ) also depends on the specific tasks the ADRL is trying to solve . Combining the generation function and the refinement function , the augmented policy estimator can be represented as πθµ̂ ( s ) = g ◦ µ̂θµ̂ . Training with DDPG We train parameters for the actor and the critic networks using DDPG . Although the augmented policy πθµ̂ consists of an actor network µ̂θµ̂ and a refinement function g , the training of the full policy follows the policy gradient of µ̂θµ̂ because the effects of g are deterministic aspects of the environment E ( Dulac-Arnold et al. , 2015 ) . The actor and the critic network are trained with variations of formulae ( 1 ) and ( 2 ) : L ( θQ ) = E [ ( yt −Q ( st , µ̂ ( st|θµ̂ ) |θQ ) ) 2 ] , where yt = r ( st , πθµ̂ ( st ) ) + γQ ( st+1 , µ̂ ( st+1|θµ̂ ) |θQ ) Algorithm 1 Augmented Policy 1 : State s previously received from the environment E 2 : Ak = µ ( s|θµ ) ( generating k candidate actors ) 3 : a = g ( Ak ) ( a ) ( refines the choice of the action with g ( Ak ) = arg maxa=Ak , i Q̃E ( a ) ) 4 : Apply a to environment ; receive r , s′ ∇θµ̂J ≈ E [ ∇AkQ ( s , Ak|θQ ) |s=st , Ak=µ̂ ( st ) · ∇θµ̂ µ̂ ( s|θ µ̂ ) |s=st ] . ( 3 )

Summary: This paper describes an improved way to determine weight quantization bit lengths using reinforcement learning, by injecting model evaluation directly into action selection.  Building upon a DRL setup where the action at each timestep corresponds to selecting a bit value for each layer, the method adds a "Q-value indicator" function Q~ that selects among a set of candidate actions, and filters based on the model performance quantizing the layer to each level.  This seems to form a hybrid between DRL and greedy search, using a greedy criteria Q~ to filter proposals made by the DRL agent.  Experiments show very good performance, with similar or better quantization levels and accuracy as other DRL-based methods and much faster runtime.