Learning Minimal Representations with Model Invariance

Sparsity has been identified as an important characteristic in learning neural networks that generalize well , forming the key idea in constructing minimal representations . Minimal representations are ones that only encode information required to predict well on a task and nothing more . In this paper we present a powerful approach to learning minimal representations . Our method , called MODINV or model invariance , argues for learning using multiple predictors and a single representation , creating a bottleneck architecture . Predictors ’ learning landscapes are diversified by training independently and with different learning rates . The common representation acts as a implicit invariance objective to avoid the different spurious correlations captured by individual predictors . This in turn leads to better generalization performance . We test MODINV on both reinforcement learning and vision settings , showcasing strong performance boosts in both . It is extremely simple to implement , does not lead to increased wall clock time while training , and can be applied across different problem settings . 1 INTRODUCTION . Learning efficient representations that generalize well is a long standing problem of machine learning , and particularly of deep learning ( Doersch & Zisserman , 2017 ; Noroozi & Favaro , 2016 ; Oord et al. , 2018 ) . Algorithms that exploit structure in the real world through effective inductive biases are key to solving this problem . Several inductive biases have been successfully used in the past , from early work using translational invariance for developing CNNs ( LeCun et al. , 2004 ; 2010 ) , to recent data augmentations for developing multiple self-supervised algorithms ( Chen et al. , 2020 ; Grill et al. , 2020 ; Chen & He , 2021 ) . A key inductive bias in multiple works is sparsity of concepts ( Hoefler et al. , 2021 ) . Essentially , it refers to how objects in the real world tend to interact and affect the dynamics of other objects only in a small neighbourhood . Consider for example picking up a pen . This simple task only changes the dynamics of a very small part of the world , namely the pen and the hand of the person grasping it . Exploiting such structure allows learning representations that can generalize better in the real world . In this paper , we develop a method that exploits sparsity by learning minimal representations . The well known saying ‘ neurons that fire together wire together ’ points to the phenomenon that that neurons that have similar output for a given input also have strong weights between them ( Hebb , 2005 ) . Ideally , given the same set of neurons to model certain concepts , we would want neurons that relate to a particular concept to have strengthened weights while those relating to another concept to have weakened weights ( Sun et al. , 2016 ) . Therefore , low correlation in the outputs of the representation would lead to less redundancy in modelling concepts and thus better generalization . This is precisely the motivation behind minimality leading to sparse representations and improved generalization . Minimal representations ( Tishby et al. , 2000 ; Shamir et al. , 2010 ) refer to ones that have sufficient but minimal information w.r.t . the task at hand . Therefore , it must be only possible to solve the task in hand and no other task , thus leading to better generalization for the given task . Another viewpoint is that minimal representations also lead to a reduction in spurious correlations . Since there is minimal noise present in the representation to cause spurious correlations , this again results in better generalization . Our approach , which we call model invariance ( MODINV ) , learns minimal representations by incorporating multiple predictors over a single , common representation for a given task ( see Figure 1 ) . Each predictor head is then trained independently ( a given sample only trains one head ) and with a learning rate different to other heads . Theoretically , this can be looked at as a manifestation of the information bottleneck principle , which balances between sufficiency and minimality . Intuitively , this can be looked at as providing robustness to spurious correlations . Since each head learns a different prediction ( based on its own optimization and learning on different parts of data set ) , each head is susceptible to different spurious correlations . However , since there is only one common representation to support all the predictor heads , this representation must be robust to any spurious correlations as training progresses . We evaluate our method on two separate settings ; 1 ) Reinforcement Learning : The DeepMind Control Suite ( DMC ) ( Tassa et al. , 2018 ) benchmark with distractors , which involves natural video playing in the background while the agent is controlled in the foreground and 2 ) Vision : The CIFAR10 ( Krizhevsky et al. , 2009 ) and STL-10 ( Coates et al. , 2011 ) datasets , where no labels are used for learning the representations , followed by the standard linear probing evaluation protocol to test generalization . Contributions : We introduce MODINV , a simple method motivated by incorporating sparsity in representations , while showing connections to the information bottleneck principle . MODINV is conceptually simple , easy to implement , and improves representation learning performance in both reinforcement learning ( from pixels ) and vision settings . In reinforcement learning , MODINV achieves stronger than state-of-the-art performance on the DMC Suite with distractors benchmark , both with augmentations and without . In vision , MODINV leads to better performance than the baseline method of SIMSIAM when using a linear predictor on both CIFAR-10 and STL-10 datasets . 2 PRELIMINARIES . 2.1 INFORMATION BOTTLENECK PRINCIPLE . The well known information bottleneck ( IB ) principle ( Tishby et al. , 2000 ; Shamir et al. , 2010 ) formally describes the ideas of sufficiency and minimality for learning representations . Specifically , it states that representations should contain sufficient information to do well on the downstream task but also contain minimal information about the inputs , leading to better , more robust generalization . This is in contrast to minimizing the empirical risk only , which is a sound technique in the infinite sample case . However , in the finite sample case , minimizing only empirical risk can lead to poor generalization . A similar observation has been noted in works that build over the invariant risk minimization principle ( IRM ) ( Arjovsky et al. , 2019 ) . 2.2 SELF-SUPERVISED LEARNING IN RL . RL considers the agent ’ s interaction with the environment as a discrete time γ-discounted Markov Decision Process ( MDP ) ( Puterman , 2014 ) M = ( X , A , P , R , γ , µ0 ) , where X denotes a finite state space and A is the action space ; P ≡ P ( x′|x , a ) is the transition kernel ; R ≡ r ( x , a ) is the reward function ; γ ∈ [ 0 , 1 ) is the discount factor ; and µ0 is the initial state distribution . The objective is to find a policy π : X → ∆A , where ∆A is the set of probability distributions on A such that the value function of a policy π at a state x ∈ X , V π ( x ) ≡ E [ ∑ t≥0 γ tr ( xt , at ) |x0 = x , π ] is maximized . Recently a lot of progress has been made in making RL over pixels as sample efficient as when learning over true states . Most approaches can be divided into two classes . The first involves using auxiliary losses over the representation so as to inject as much information about the downstream task as possible . These auxiliary objectives include reconstructing observations ( Hafner et al. , 2020 ; Ha & Schmidhuber , 2018 ) , or predicting next observation , reward ( Schaul et al. , 2015 ) or even encoded states ( Gelada et al. , 2019 ) . The second class borrows ideas from self-supervised learning methods and augments the observations ( Yarats et al. , 2021 ; Laskin et al. , 2020b ) with techniques like random crop , cutout , color jitter etc . Since randomly augmented samples from the same observation have the same Q values , an implicit invariance to augmentations is enforced in the representation pipeline . In the case of hard exploration tasks , auxiliary objectives such as estimating value functions of random cumulants has shown to be useful as well . A common baseline architecture used in a lot of methods is an actor-critic setup where a common convolutional encoder is used to compute the encoded latent state from the raw pixel observations . For continuous control tasks , soft actor-critic ( SAC , Haarnoja et al . ( 2018 ) ) is used as the actor-critic algorithm and the gradients from the actor are stopped from updating the encoded state . Only the critic trains the representation network fθ . This is since allowing gradients from both actor and critic leads to noisy estimates and the eventual divergence . Moreover , some methods also add a reward prediction loss from the encoded state , a transition prediction loss , or a pixel reconstruction loss . 2.3 SELF-SUPERVISED LEARNING IN VISION . An important property of self-supervised methods in vision is avoiding collapse in representation , or trivial solutions where the representation simply outputs a constant . These involve the use of two similar or identical networks , an online and a target network . Particular approaches largely fall in two categories , that of contrastive ( Chen et al. , 2020 ) and non-contrastive methods ( Grill et al. , 2020 ; Chen & He , 2021 ) , depending on how they avoid representation collapse . Contrastive methods create positive and negative pairs from random augmentations of data ( e.g. , cutout , crop , flip ) , and enforce an objective such that positive samples are brought closer and negative samples are pulled away in representation space . On the other hand , non-contrastive methods do away with creating positive and negative pairs and instead use the likes of assymetric architectures , stop gradient , and momentum encoders for the target network to learn non-trivial representations . One method we use to build over later in the paper is SIMSIAM ( Chen & He , 2021 ) , which is a non-contrastive method and only incorporates the most essential components required for preventing collapse . In particular , it uses two networks , an online and a target network , where both use a backbone network ( usually a ResNet 18 or ResNet 50 ( He et al. , 2016 ) ) and a projection MLP ( 1 hidden layer ) . Moreover , only the online network has a prediction MLP attached to it , which makes the overall architecture assymetric . The target network has a stop gradient applied to it and the overall objective is a cosine similarity loss between the outputs of the online and target networks . It has been noted that stop-gradient is enough to avoid collapse and that slowly updating the target network weights to match those of the online network , i.e . exponential moving average ( EMA ) , is not a necessary requirement . Adding the EMA to the SIMSIAM architecture results in the BYOL ( Grill et al. , 2020 ) architecture . 3 METHOD . Consider a task Y , which can correspond to a variety of objectives . For instance , it could be an auxiliary task of predicting the next state when interacting with an RL environment , or simply a classification task on a vision benchmark . We are interested in training a model that can do well on task Y , where the model consists of two components , a backbone representation fθ and a predictor/classifier hφ which is attached over the representation . In a lot of tasks that involve pixel input , spurious correlations can arise due to various reasons such as lack of non-iid data , irrelevant information , confounders etc . that can result in poor generalization at test time . An attractive property for better generalization has been the idea of sparse representations , in that the mutual information of any two dimensions of the representation must be low . However , in order to avoid learning trivial solutions , the representation as a whole should encode enough information about the downstream task as well . A combination of both of these objectives then leads to a model with better generalization capability . In this paper , we balance the above two objectives using the idea of model invariance ( MODINV ) . MODINV deploys multiple predictors over a single representation , while training each predictor independently , each capturing different spurious correlations . The common representation acts as a implicit invariance loss which ensures that only the optimal representation remains at convergence . Intuitively , each predictor can be looked at as an augmented version of the optimal model , where the augmentation is over the model space and refers to particular spurious correlations arising in the model that differ it from the optimal model . Note that each predictor head is the same in architecture , and we only diversify the learning procedure of each head ( through different initialization , independent training , different learning rates ) . This is so that we eventually converge to an optimal representation for a particular predictor/classifier and not for all predictor/classifier families , which is a much more hard to optimize in practise . For all experiments , we use a random routing to decide which data sample is used to train which predictor head . Furthermore , the learning rates of the different predictor heads are chosen intuitively , i.e . if the base rate is 3e-3 then we choose one slightly higher rate ( 5e-3 ) and one slightly smaller rate ( 1e-3 ) for K=3 predictors .

Summary: This paper proposes using an ensemble of predictors or heads on top of a shared representation to improve performance. The ensemble of predictors/heads are made "diverse" by using different learning rates. Some performance gains are shown in RL settings, while some stability gains are shown in a particular vision setting.