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: Inspired by the information bottleneck principle, the authors propose to learn minimal representations by independently training multiple predictive heads on top of a shared representation, such that the resultant predictor is more robust to spurious correlations. To test this proposal, the authors provide extensive set of experiments in RL from pixels and vision benchmark tasks. The ablation studies help make the case too.