Year: 2021

Abstract:
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.

text: 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.

heading: 1 INTRODUCTION
text: 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.

heading: 2 PRELIMINARIES

heading: 2.1 INFORMATION BOTTLENECK PRINCIPLE
text: 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).

heading: 2.2 SELF-SUPERVISED LEARNING IN RL
text: 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.

heading: 2.3 SELF-SUPERVISED LEARNING IN VISION
text: 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.

heading: 3 METHOD
text: 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.

heading: 4 EXPERIMENTS

heading: 4.1 REINFORCEMENT LEARNING
text: For RL over pixels, we build over a standard actor-critic setup and use MODINV as part of an auxiliary task that predicts the next latent state (see Figure 2). This auxiliary task is used to train the representation which supports the actor-critic heads. Specifically, the input is the observation xt and the representation fθ encodes it into the latent state st. We then concatenate the latent state and the action and pass it through a linear MLP to generate a common state-action representation (predicting the next state st+1 requires information about both the state and action). Since the true latent state is not known, predicting only the next state st+1 can result in a representation collapse*. Therefore, we ground the state prediction loss by also predicting the reward, which does have a valid grounding and thus avoids collapse. Finally, we deploy MODINV by using multiple predictor heads for the state prediction task, with each head being coupled with a corresponding reward prediction head.

heading: 4.1.1 IMPLEMENTATION DETAILS
text: We implement SAC (Haarnoja et al., 2018) as the base agent with the actor and critic sharing a common representation. Our actor-critic setup is similar to the one used in SAC-AE (Yarats et al., 2019), except the reward and state prediction task. Our transition or state prediction model is deterministic in nature, involving a 3 layer encoder MLP, followed by a 3 layer decoder MLP. The reward prediction network is a 2 layer MLP. We do not tune for MLP sizes and widths at all in reporting our results. Both the reward and transition losses train the representation alongside the critic. We use three MODINV heads, each with a slightly different learning rate of 3e− 3, 3e− 4, and 3e− 5 respectively.
*Recently it was shown that collapse can be prevented when using a cosine loss and an architecture similar to BYOL instead of a squared error loss.

heading: 4.1.2 RESULTS
text: We test our method on the DMC Suite with distractors. Particularly, we test on the six popular domains of Walker Walk, Reacher Easy, Hopper Hop, Finger Spin, Cheetah Run and Cartpole Swingup. We compare with other methods that learn over pixels including DREAMER, CURL, DBC, and RAD. Dreamer (Hafner et al., 2020) is a model-based method that performs pixel reconstruction and reward prediction to learn a model in the latent state space and then performs model-based updates to a base SAC agent. CURL (Laskin et al., 2020a) uses a contrastive loss similar to that in CPC (Oord et al., 2018) to train the representation network, using data augmentations to generate positive and negative samples. RAD (Laskin et al., 2020b) is another method that simply augments the samples in the replay buffer without adding any other loss function to the base SAC setup. Finally, note that both RAD and DrQ (Kostrikov et al., 2020) are similar in performance as they both use data augmentations in the replay buffer samples. Since data augmentations explicitly remove the actual distractions in the background, a direct comparison with data augmentations is not fair. Therefore, we also compare our method when augmentations are added alongside MODINV (denoted by MODINV+AUG). We evaluate all methods at both 100K and 500K environment steps, which is the standard number of steps in this benchmark.
Our results (Table 1) show that MODINV clearly outperforms all 4 methods (DREAMER, CURL, DBC) at the 500K step mark. Note that since all methods use the same number of gradient steps, MODINV is at a slight disadvantage as the total number of gradient updates are divided amongst the K predictor heads.
4.2 VISION
For vision, we adopt the base framework of SIMSIAM, a non-contrastive self-supervised method for learning representations. SIMSIAM has three model components, the backbone network or the representaion, a projection MLP and a prediction MLP. We deploy MODINV in the projection MLP, with a single backbone representation and a single prediction MLP (see Figure 3). At each training iteration, the same MODINV projection head is sampled for both the online and target networks and trained using the original SIMSIAM loss function. Unlike the reinforcement learning case, where any of the modules after the representation (state and reward prediction
heads) are not used actively in the algorithm, in the vision case the projector and predictor both are used to align the backbone representation (Tian et al., 2021). Therefore, ensuring that each projection head in MODINV is trained sufficiently (as in the base case) is important.

heading: 4.2.1 IMPLEMENTATION DETAILS
text: We test our approach on the CIFAR-10 and STL-10 datasets. As data augmentations, we use random crop, color jitter, horizontal flip, random grayscale for both datasets and also use blur for STL-10. Furthermore, we deploy three MODINV heads for the projection MLP, with the learning rates of 0.03, 0.3, 0.003 for each head respectively. The learning rate of all other components in the model is set to 0.03, with a weight decay of 0.0004, momentum 0.9. We run two sets of experiments, one with a linear prediction MLP and one with the standard two layer prediction MLP. The rest of the architecture is the same as that of SIMSIAM, with the output dimension of the projector and the predictor MLP set to 2048. We use a batch size of 512 for CIFAR-10 and 128 for STL-10.
Linear vs Non-linear Predictor . For the SSL setting we deploy MODINV in two different configurations. The first involves a linear predictor while the second involves a standard 2-layer predictor. The 2-layer predictor has been speculated to be using a lot of ‘lucky’ initialization that lead to alignment between the backbone representation and the prediction MLP (Tian et al., 2021). However, the linear predictor does not have this property and thus testing in this setting offers a much more

heading: 100K STEP SCORES DREAMER CURL DBC MODINV RAD MODINV+AUG

heading: 500K STEP SCORES
text: robust evaluation option. Unless specified otherwise, our ablations for vision in Section 5 use the linear predictor setting for MODINV.

heading: 4.2.2 RESULTS
text: We see that with a linear predictor, the base SIMSIAM version fails for CIFAR-10 and only achieves 41.3%† accuracy while the MODINV version achieves an accuracy of 82.5% at 330 epochs, similar to that of the 2-layer predictor case. A similar gap in performance is observed for earlier training epochs as well (Table 2). A careful reader might attribute this gap to the use of different learning rates in MODINV, since a relatively slower learning rate for the non-predictor layers can correspond to improved performance, especially when EMA is not used in the target network (Tian et al., 2021). To test this, we train the base SIMSIAM model individually with the different projector learning rates used in the MODINV version. We see that all three learning rates fail to the same performance of around 40% accuracy. This clearly shows that MODINV contributes much more to better generalization than using slower or faster learning rates.
†Note that this is not complete collapse since the accuracy still increases very gradually after 200 training epochs.

heading: 5 ABLATIONS

heading: 5.1 ABLATION ON NUMBER OF HEADS
text: In our experiments, we observe that only K = 3 heads are sufficient to provide good performance boosts compared to when not using MODINV (i.e. K = 1) across the RL and SSL settings. Furthermore, with everything else the same, adding an extra head after K = 3 leads to diminishing performance gains. This can be most clearly seen for SIMSIAM + MODINV, where we see increasing gains with increasing K (see Table 3). Note that although relative gains saturate after K = 3, it is possible that further diversification in predictor training can lead to even better results when K > 3.

heading: 5.2 CORRELATION IN DIMENSIONS
text: As mentioned, our initial motivation stems from sparsity in concepts. This refers to low mutual information between any two dimensions in the representation fθ. A useful metric to check for such a characteristic is therefore the mean correlation between any two dimensions. We plot this metric during training with and without MODINV. Ideally, we would hope to get lower mean correlation when using MODINV as opposed to when not. Our experiments show that this is indeed the case, as the mean correlation decreases much more steadily when we use MODINV. Moreover, the decrease in this metric has a strong correlation with the performance of the agent.
L = ∑ b ( P̂(zb)− P(zb) )2 + λ ∑ i ∑ j 6=i Cij2 where Cij = ∑ b zb,i zb,j√∑ b (zb,i) 2 √∑ b (zb,j) 2 (1)
Based on this observation, we also run an experiment where we take the baseline architecture (i.e. without MODINV) and add the mean correlation between dimensions of fθ as an auxiliary loss. Equation 1 describes the exact loss, where the summation is over the batch b. The first term is the state prediction error from Section 4.1 while the second term is the exact loss from Zbontar et al. (2021). Our results show that optimizing directly for this loss leads to much better performance in the Reacher Easy task, while much worse performance in the Walker Walk (Table 4). We choose to test on these two domains since they are hard to solve for most methods that we compare to. This is an interesting observation since similar losses that have been used successfully in vision recently (Bardes et al., 2021) might offer strong performance improvements in RL as well.

heading: 5.3 IMPORTANCE OF DIFFERENT LEARNING RATES
text: When using multiple heads in MODINV, it is important to diversify the learning of heads so they capture different spurious correlations. Using different learning rates is a vital component for achieving this, besides different model initialization and independent training. We test this by evaluating performance with multiple sets of learning rates for K = 3 predictor heads. We see that performance improves when the learning rates are diverse as compared to when all heads use the same learning rate (Table 5). Moving on to diversifying learning of predictor heads through different optimizers is left for future work (see Section 7).

heading: 5.4 ABLATION ON ONLY PREDICTING REWARD
text: The RL application of MODINV involves both state prediction and reward prediction. We now test how the performance might vary if we only have the reward prediction task. It is worth noting that MODINV can be applied to the actor-critic heads as well, however optimization in that case becomes difficult as the critic learning loss does not have a stable ground truth. In Table 6 we see that performance deteriorate in general for the three tasks, especially for Cheetah Run. This may suggest that state prediction provides a slightly more stable ground truth to optimize the representation.

heading: 5.5 ABLATION WITH DATA AUGMENTATIONS
text: Our results show that MODINV acts complimentarily to augmentations. In our experiments in the RL setting, we observe that adding data augmentations over MODINV leads to better performance than when only using MODINV. ‘Crop’ is the only augmentation used here, which is standard for RL environments. Essentially, cropping removes a lot of background information which is irrelevant to the task (top-down approach), thus aiding in better or more robust
estimates for next state and critic losses. However, not all augmentations have a similar effect (e.g. flip), and thus end up with similar performance as the baseline. On the other hand, MODINV avoids spurious correlations by enforcing minimality, resulting in a bottom-up approach to robust reward and critic loss optimization.

heading: 6 RELATED WORK AND DISCUSSION

heading: 6.1 DECODABLE INFORMATION BOTTLENECK
text: The IB principle does not take into account the classifier/predictor family which is attached over the representation and thus can be considered too restrictive. Moreover, the minimality term is hard to approximate for practical losses. To remedy this, the decodable information bottleneck (Dubois et al., 2020) recently introduced notions of sufficiency and minimality for a given predictor/classifier family. They also provided a practical method which has a similar structure to the MODINV approach. In particular, multiple predictor heads are attached to a common representation, where one head (sufficiency head) minimizes the standard empirical risk while all other heads (minimality heads) are learnt such that they cannot predict arbitrary relabellings of the same data. This is done by providing different random relabellings and then reversing the gradients for all the minimality heads. This ensures that the common representation contains information required to predict only the correct labels and no other. The idea in MODINV is similar but we do not deploy reverse gradients or relabel data, which in practise can lead to instability issues. Therefore, MODINV can be seen as a specific instantiation of the DBC framework, following similar theoretical guarantees while being much easier to implement and resulting in better performance.

heading: 6.2 DECORRELATION IN NON-CONTRASTIVE SSL
text: Recently, BARLOW TWINS (Zbontar et al., 2021) was introduced as a new non-contrastive selfsupervised learning method which does not use an asymmetric architecture while still avoids collapse in representation. This was achieved by enforcing the off-diagonal elements of the cross-correlation
matrix of the outputs of the representation to be zero. Essentially, such a loss ensures that each output dimension of the representation is decorrelated with each other while minimizing the standard loss of enforcing invariance to augmentations. Interestingly, the BARLOW TWINS loss can be traced back as an approximation to the information bottleneck principle again, hence connecting decorrelation in representation dimensions as an indicator for better generalization (as noted similarly in the introduction). We showed that MODINV also leads to more decorrelation in the representation dimensions, thus supporting this observation. However, we also showed that directly optimizing for this loss might not be a good idea always, with such a variant leading to worse empirical performance than MODINV in the RL setting. Nevertheless, it is important to note that the MODINV and BARLOW TWINS objectives are primarily complimentary in nature. This is since MODINV always requires a sufficiency objective to start with, which could be provided by any SSL loss which avoids collapse, be it SIMSIAM with the use of a predictor and stop gradient or be it BARLOW TWINS with the use of the decorrelation loss. In similar flavor to BARLOW TWINS, VICREG (Bardes et al., 2021) also uses the decorrelation loss from BARLOW TWINS and adds a variance regularization loss to prevent collapse. Finally, Hua et al. (2021) introduced a technique called decorrelated batch normalization (DBN) and also showed that collapse can be avoided with no predictor.

heading: 6.3 AUXILIARY VALUE PREDICTION IN RL
text: To improve representation learning in RL, a prominent idea has been that of predicting multiple value functions from a common representation. Sutton et al. (2011) use random cumulants to define the different value functions. The value improvement path (Dabney et al., 2020) idea uses a mixture of value functions of past policies. Although these works have a similar looking structure to MODINV, there are a couple of critical differences. First, the multiple heads in MODINV are predictors of the same object while the above papers use different objects (different value functions). Second, the training of each head is not independent. These two differences dissociate the above works from the IB idea, thus making them complimentary to MODINV.

heading: 7 FUTURE WORK
text: MODINV shows promising improvements in both the RL and SSL settings, thus hinting at the benefits offered by the IB principle in practice. Nevertheless, MODINV is still one instantiation of this idea. There remain quite a few directions to explore in future. In particular, analysing how further diversification in the K predictors can be achieved is a promising idea. For example, we can introduce different optimizers (Adam, SGD) for each predictor, while also varying parameters such as momentum, weight decay etc. Another possibility could be to update a subset of predictors for a given sample, instead of only updating one predictor at each training iteration. Finally, since data augmentations change the input data, each predictor could be trained on a different augmentation scheme, thus ensuring further diversification in optimizing different predictors.
On a separate axis, analysing how similar MODINV for vision is with other SSL methods such as VICREG is interesting. For the RL case, MODINV can also potentially be applied to the Q functions directly (each predictor is a separate Q head over a common representation) instead of the reward and transition modules. This could turn out to be hard to optimize for certain environments though, since Q functions can be less developed for sparse reward settings.

heading: 8 CONCLUSION
text: We introduce a general representation learning method motivated from the ideas of sparsity and minimal-sufficient representations. Our method, MODINV, uses multiple predictors over a common representation, diversifying the training of each predictor such that the common representation acts as an implicit invariance mechanism to spurious correlations. We show instances of this idea in both reinforcement learning and vision settings. Overall, MODINV leads to more robust generalization across different evaluations, leading us to believe that with further analysis, we can get further gains in representation generalization. MODINV can be viewed as a specific instantiation of the information bottleneck principle, one that is simple in implementation and in concept.

heading: 1 HYPERPARAMETERS

Review summaries:
Motivated from the idea that minimal representations — that encode information relevant to a task and nothing more — are likely to generalize well, the paper proposes an approach, titled Model Invariance (ModInv) to learn representations. The underlying idea is to use multiple diversified predictors solving auxiliary objectives on top of a shared representation/encoder. Different predictors are diversified in terms of their training dynamics — different initializations and learning rates. The approach relies on the hope that a shared representation that can lead to optimal performance for a diverse set of predictors is likely encoding some notion of implicit invariance to changes in predictors, thereby, introducing some notion of minimality. The authors conduct experiments on reinforcement learning and self-supervised (in computer vision) settings and claim that ModInv provides strong performance improvements in both. Based on further ablations,  the authors demonstrate how sensitive ModInv is to the number of predictors and the diversity in terms of training dynamics (in terms of learning rate), and whether ModInv is complementary to recent augmentation strategies in reinforcement learning.

The authors propose MODINV, a new algorithm to learn minimal and compressed representations. MODINV uses several predictors with different initializations and learning but all built on top of the same learnt representation.  After explaining the algorithm and its relation to minimal representation the authors evaluate it on vision and RL tasks.  On vision tasks the authors show improvement on CIFAR-10 by adapting their method on top of SimSiam in the specific case of a linear predictor. On RL they evaluate on control tasks from pixel inputs.  The authors also provide an experiment to show that adding a loss to discourage correlation within features helps to improve RL performance.

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.

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.