How many weights are enough : can tensor factorization learn efficient policies ?

Deep reinforcement learning requires a heavy price in terms of sample efficiency and overparameterization in the neural networks used for function approximation . In this work , we employ tensor factorization in order to learn more compact representations for reinforcement learning policies . We show empirically that in the low-data regime , it is possible to learn online policies with 2 to 10 times less total coefficients , with little to no loss of performance . We also leverage progress in second order optimization , and use the theory of wavelet scattering to further reduce the number of learned coefficients , by foregoing learning the topmost convolutional layer filters altogether . We evaluate our results on the Atari suite against recent baseline algorithms that represent the state-of-the-art in data efficiency , and get comparable results with an order of magnitude gain in weight parsimony . 1 INTRODUCTION . The successes of reinforcement learning ( thereafter ’ RL ’ ) , and specifically deep RL , come at a heavy computational price . It is well known that achieving human-level performance in domains such as Atari ( Sutton & Barto , 2018 ; Mnih et al. , 2013 ; Hessel et al. , 2017 ) requires hundreds of millions of frames of environment interaction . As such , the problem of sample efficiency in RL is of critical importance . Several tracks of concurrent research are being investigated , and have reduced by orders of magnitude the number of environment interactions required for good performance beyond the previous benchmark of biologically-inspired episodic control methods ( Blundell et al. , 2016 ; Pritzel et al. , 2017 ) to a couple hours of human gameplay time ( van Hasselt et al. , 2019 ; Kaiser et al. , 2019 ) . However , while the data-efficiency of RL methods has seen recent drastic performance gains , the function approximators they use still require millions of learned weights , potentially still leaving them heavily overparameterized . Independently motivated by biological facts like the behavioural readiness of newborn animals , several authors ( Gaier & Ha , 2019 ; Cuccu et al. , 2018 ; Wang et al. , 2019 ) have recently looked at doing away with learning so many weights for RL tasks . Smaller networks not only train faster , but may yet offer another avenue for gains in the form of better generalization ( Zhang et al. , 2016 ) . Recent work from Gaier & Ha ( 2019 ) studies the effect of inductive bias of neural architectures in RL ; they forego training altogether , but transfer networks that only obtain ’ better than chance performance on MNIST ’ . In similar fashion , Wang et al . ( 2019 ) investigate the effect of random projections in the restricted setting of imitation learning . Finally , Cuccu et al . ( 2018 ) manage human-level performance on the Atari suite using a separate dictionary learning procedure for their features , bypassing the usual end-to-end learning paradigm . The perspective of neural architecture search applied to RL appears difficult , if not computationally inextricable . Concurrently , the study of biologically-inspired models of learning has exhibited two mathematical characterizations that might be critical in explaining how biological learning takes place so efficiently . First , the low-rank properties of learned perceptual manifolds ( Chung et al. , 2018 ; 2016 ) are giving rise to a rich theory borrowing from statistical physics . Second , another well known line of work has identified Gabor filters , and more generally wavelet filter-like structures , in the actual visual cortex of animals ( Jones & Palmer , 1987 ) , and linked those to sparsity-promoting methods and dictionary learning ( Olshausen & Field , 1996 ; 1997 ; Hyvärinen & Hoyer , 2001 ) . But these breakthroughs have not , so far , been reflected as inductive priors in the shape of modifications in deep RL neural networks architectures , which remain fairly fixed on the Atari domain . Therefore the following questions remain : how parsimonious do function approximators in RL need to be , in order to maintain good performance ? And can we be at once sample-efficient and weightefficient ? In this work , we turn to the mathematical theories of tensor factorization ( Cichocki et al. , 2009 ) , second-order optimization ( Amari , 1998 ; Martens & Grosse , 2015 ) and wavelet scattering ( Mallat , 2011 ) to answer this question positively and empirically , in a model-free setting . We propose to use these methods in order to save weights and therefore favour convergence of policies : • We replace dense , fully-connected layers with tensor regression layers . • Optionally , we replace the topmost layer in the convolutional architecture with a scattering layer ; the deeper convolutional layers are left untouched . • The ( positive ) impact of second-order optimization is also evaluated . To the best of our knowledge , this is the first time those fields have been combined together in this context , and that tensor factorization is applied to deep RL . 2 BACKGROUND & RELATED WORK . 2.1 DEEP REINFORCEMENT LEARNING . We consider the standard Markov Decision Process framework as in Sutton & Barto ( 2018 ) . This setting is characterised by a tuple 〈S , A , T , R , γ〉 , where S is a set of states , A a set of actions , R a reward function that is the immediate , intrinsic desirability of a certain state , T a transition dynamics and γ ∈ [ 0 , 1 ] a discount factor . The purpose of the RL problem is to to find a policy π , which represents a mapping from states to a probability distribution over actions , that is optimal , i.e. , that maximizes the expected cumulative discounted return ∑∞ k=0 γ kRt+k+1 at each state st ∈ S. In Q-learning , the policy is given implicitly by acting greedily or -greedily with respect to learned action-value functions qπ ( s , a ) , that are learned following the Bellman equation . In deep Q-learning , qθ becomes parameterized by the weights θ of a neural network and one minimizes the expected Bellman loss : E ( Rt+1 + γt+1 max a′ qθ ( St+1 , a ′ ) − qθ ( St , At ) ) 2 In practice , this is implemented stochastically via uniform sampling of transitions in an experience replay buffer , as is done in the seminal paper Mnih et al . ( 2013 ) . Several algorithmic refinements to that approach exist . First , Double Q-learning ( van Hasselt et al. , 2015 ) proposes to decouple learning between two networks in order to alleviate the Q-value overestimation problem . Second , dueling Q-networks ( Wang et al. , 2015 ) explicitly decompose the learning of an action-value function qθ ( s , a ) as the sum of an action-independent state-value , much like what is traditionally done in policy gradient methods ( Sutton & Barto , 2018 ) , implemented via a two-headed neural network architecture . Finally , prioritized RL ( Schaul et al. , 2015 ) proposes to replace the uniform sampling of transitions in the experience replay buffer with importance sampling , by prioritizing those transitions that present the most Bellman error ( those transitions that are deemed the most ’ surprising ’ by the agent ) . Fortunato et al . ( 2017 ) use extra weights to learn the variance of the exploration noise in a granular fashion , while Bellemare et al . ( 2017 ) propose to learn a full distribution of action-values for each action and state . Combined , those methods form the basis of the Rainbow algorithm in Hessel et al . ( 2017 ) . 2.2 TENSOR FACTORIZATION . Here we introduce notations and concepts from the tensor factorization literature . An intuition is that the two main decompositions below , CP and Tucker decompositions , can be understood as multilinear algebra analogues of SVD or eigendecomposition . CP decomposition . A tensor X ∈ RI1×I2×···×IN , can be decomposed into a sum of R rank-1 tensors , known as the Canonical-Polyadic decomposition , where R is the rank of the decomposition . Its purpose is to find vectors u ( 1 ) k , u ( 2 ) k , · · · , u ( N ) k , for k = [ 1 . . . R ] , as well as a vector of weights λ ∈ RR such that : X = R∑ k=1 λku ( 1 ) k ◦ u ( 2 ) k ◦ · · · ◦ u ( N ) k︸ ︷︷ ︸ rank-1 components Tucker decomposition . A tensor X ∈ RI1×I2×···×IN , can be decomposed into a low rank approximation including a core G ∈ RR1×R2×···×RN and a set of projection factors ( U ( 0 ) , · · · , U ( N−1 ) ) , with U ( k ) ∈ RRk , Îk , k ∈ ( 0 , · · · , N − 1 ) that , when projected along the corresponding dimension of the core , reconstruct the full tensor X . The tensor in its decomposed form can be written : X = G ×1 U ( 1 ) ×2 U ( 2 ) × · · · ×N U ( N ) = î G ; U ( 1 ) , · · · , U ( N ) ó Tensor regression layer . For two tensors X ∈ RK1×···×Kx×I1×···×IN and Y ∈ RI1×···×IN×L1×···×Ly , we denote by 〈X , Y〉N ∈ RK1×···×Kx×L1×···×Ly the contraction of X by Y along their N last modes ; their generalized inner product is 〈X , Y〉N = I1∑ i1=1 I2∑ i2=1 · · · IN∑ in=1 X ... , i1 , i2 , ... , inYi1 , i2 , ... , in , ... This makes it possible to define a tensor regression layer ( Kossaifi et al. , 2017b ) that is differentiable and learnable end-to-end by gradient descent . Let us denote by X ∈ RI1×I2×···×IN the input activation tensor for a sample and y ∈ RIN the label vector . A tensor regression layer estimates the regression weight tensorW ∈ RI1×I2×···×IN under a low-rank decomposition . In the case of a Tucker decomposition ( as per our experiments ) with ranks ( R1 , · · · , RN ) , we have : y = 〈X , W〉N + b withW = G ×1 U ( 1 ) ×2 U ( 2 ) · · · ×N U ( N ) as G ∈ RR1×···×RN , U ( k ) ∈ RIk×Rk for each k in [ 1 . . . N ] and U ( N ) ∈ R1×RN . 2.3 WAVELET SCATTERING . The wavelet scattering transform was originally introduced by Mallat ( 2011 ) and Bruna & Mallat ( 2012 ) as a non-linear extension to the classical wavelet filter bank decomposition . Its principle is as follows . Denoting by x ~ y [ n ] the 2-dimensional , circular convolution of two signals x [ n ] and y [ n ] , let us assume that we have pre-defined two wavelet filter banks available ¶ ψ ( 1 ) λ1 [ n ] © λ1∈Λ1¶ ψ ( 2 ) λ2 [ n ] © λ2∈Λ2 , with λ1 and λ2 two frequency indices . These wavelet filters correspond to high frequencies , so we also give ourselves the data of a lowpass filter φJ [ n ] . Finally , and by opposition to traditional linear wavelet transforms , we also assume a given nonlinearity ρ ( t ) . Then the scattering transform is given by coefficients of order 0,1 , and 2 , respectively : S0x [ n ] = x~ φJ [ n ] S1x [ n , λ1 ] = ρ Ä x~ ψ ( 1 ) λ1 ä ~ φJ [ n ] λ1 ∈ Λ1 S2x [ n , λ1 , λ2 ] = ρ Ä ρ Ä x~ ψ ( 1 ) λ1 ä ~ ψ ( 2 ) λ2 ä ~ φJ [ n ] λ1 ∈ Λ1 , λ2 ∈ Λ2 ( λ1 ) This can effectively be understood and implemented as a two-layer convolutional neural network whose weights are not learned but rather frozen and given by the coefficients of wavelets ψ and φ ( with Gabor filters as a special case ( Mallat , 1998 ) ) . The difference with traditional filter banks comes from the iterated modulus/nonlinear activation function applied at each stage , much like in traditional deep learning convolutional neural networks . The generic mathematical definition involves order n iterated scatterings , in the vein of Six above , but sometimes restricts nonlinearity ρ to be a modulus function | · | . In practice , the potential of scattering transforms to accelerate learning by providing ready-made convolutional layers has been investigated in Oyallon et al . ( 2013 ) and Oyallon et al . ( 2018 ) and is a subject of active ongoing research . Scattering will be the second of our weight-saving methods .

Summary: This paper investigates a list of methods to reduce the number of weights for deep RL architecture under the low-data regime. These methods include tensor regression, wavelet scattering, as well as second-order optimization (K-FAC). The experiments on the Atari games shows that by using tensor regression to replace the dense layer of the neural nets and using K-FAC for the optimization, one can reduce around 10 times of parameters without losing too much of performance.