Ignorance is Bliss: Robust Control via Information Gating

Summary

This paper empirically investigates a few approaches to modulating the amount of information used by a neural network learner in a variety of control-adjacent learning problems. The proposed approach, InfoGating, learns an input-conditioned continuous-valued mask that is applied to the same input or features thereof. The learner is trained to optimize a task loss and minimize the information allowed through by the mask. Experiments on representation learning via multi-step inverse dynamics modeling demonstrate the relative benefit of InfoGating over other approaches to information parsimony such as the variational information bottleneck. Experiments on Q-learning and fine-tuning pre-trained visual representations for behavior cloning demonstrate the benefit of InfoGating over naive baselines.

Strengths

### Originality The particular formulation of InfoGating as using input-conditioned masking amortized via a separate neural network is, to my knowledge, original. The authors remark that InfoGating is closely related to previously proposed ideas in the space of information-based regularization of neural networks. ### Quality The scope of the experiments chosen to validate the proposed idea is perhaps the main strength of this work. Three distinct experimental testbeds (dynamics modelling + probing, Q-learning, fine-tuning pre-trained representations for behavior cloning) are considered, as are a few variants of InfoGating (input gating vs. feature gating, cooperative masking vs. adversarial masking). ### Clarity The proposed methods are simple and clearly described. However, while some experimental details (architectures and hyperparameters) are presented in the supplementary material, this seems insufficiently detailed to guarantee reproducibility, so I would encourage the authors to release code. ### Significance The considered problem is important and the proposed methods are simple and intuitively reasonable.

Weaknesses

I list my perceived weaknesses in rough decreasing order of importance. - (W1) The experimental design leaves a lot to be desired. First and foremost, for the main experiment (multi-step inverse dynamics modeling + behavior cloning probe on half-cheetah), the authors report performance on "a noise-free observation space" (line 221), which I take to mean the original DeepMind Control Suite background or similar. I would instead expect evaluation to use held-out noise (distractions from the Distracting Control Suite) to properly assess the generalization capabilities of the models. As-is, contribution claim 3 (line 62) is unsubstantiated, since, to my knowledge, we typically do not use "generalization" to mean noisy train -> noiseless test (I would instead characterize this as denoising), and neither of the other testbeds include comparisons to other information-based regularization techniques. - (W2) The design choice of operationalizing information regularization via a learned input-conditioned mask pushes the burden of generalization to the masking network. This seems to be the main distinguishing aspect of InfoGating over other forms of information regularization, yet the impact of this design is neither evaluated in isolation nor in toto. - (W3) There are several key statements that are unclear or confusing. - "Although IB approaches have been beneficial in some cases, adding an information bottleneck at the penultimate step of computation does little to prevent overfitting in the preceding steps of computation, which typically comprise the overwhelming bulk of the model" (line 34). It's not clear to me that IB approaches predominantly bottleneck just before computation output. In particular, methods that use informational/Gaussian/variational dropout seem to contradict this. I would appreciate citations supporting the quoted statements. - "Roughly, the values in $ig(x)$ can be seen as specifying how many steps of forward diffusion to run in a Gaussian diffusion process initiated at $z$, where values near zero correspond to running more steps of diffusion and thus sampling from a distribution that is closer to a standard Gaussian in terms of KL divergence" (line 174). This sounds like an interesting connection, but it seems to be only true if the mask specifies a constant value. Otherwise, different pixels are noised to different extents, which doesn't seem to correspond correctly to the standard modelling assumptions in diffusion models. - (W4) The material in the half-page Background section is completely orthogonal to the proposed idea. I would move a condensed version to the experimental section since the content really is just context for the experiments. Relatedly, Algorithm 1 is InfoGating applied specifically to multi-step inverse dynamics modeling, yet one of the contribution claims is that InfoGating is general purpose. Addressing W1 and showing good performance on more meaningful out-of-distribution evaluation would cause me to increase my score substantially.

Questions

- (Q1) What do the learned masks look like for i) inputs whose underlying state is trained on, but with held-out distractors; ii) inputs whose underlying state is held-out, but with trained-on distractors; and iii) inputs whose state and distractors are both held-out? This would go a long way in explaining the benefits of (or diagnosing weaknesses with) the proposed approach. - (Q2) Why does baseline inverse dynamics modeling in Table 1 do so poorly, given the protocol of train-on-noisy, test-on-noise-free? Concretely, why doesn't adding the noise to the training data confer benefits a la data augmentation, since it specifies input transformations that the output (actions) should be invariant to? Is it because the considered noise (frames from continuously playing videos) has temporal consistency?

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

2 fair

Contribution

3 good

Limitations

Yes.

Summary

The authors hypothesise that gating information propagation in neural networks will lead to better generalisation. To achieve this they propose a system that performs gates information using a multiplicative differentiable operation, which they call InfoGating. To show the validity of their approach they compare to related models on several tasks such as contrastive dynamics control, Q-Learning and behaviour cloning.

Strengths

1. The method is well motivated. 2. The tasks are relevant to the issue the model is seeking to solve. 3. The comparison models appear to be relevant and well justified. 4. They perform relevant ablations of their design choices.

Weaknesses

While the experimental results show improved generalisation, it is not clear how general this approach will prove to be. So, while I agree that masking irrelevant parts of the input so as to limit distractor information is a good strategy (and one that humans definitely use when interacting with the world), I am less clear that this specific approach fully captures how this process works and will generalise effectively beyond to more complex datasets. To the author’s credit, they do point out limitations of their approach.

Questions

1. What is the relation between lambda and the overall performance? This is the main hyperparmeter of the system yet there is no graph showing how manipulating it affects performance. 2. If the authors have access to ground truth masks for some or all of the datsets, could they compare the masks produced by their approach with the ones produced by InfoGating?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

2 fair

Limitations

Main technical issue I see is that the reliance on the lambda parameter is very similar to the reliance of a $\beta$ parameter in the $\beta$-VAE literature, which creates a tradeoff that is not always optimal. It would have been interesting to test the information gating approach with other more principled techniques like object-centric learning which already segment the images into relevant parts (or at least they try to, eg [1]). Then info gating would be only responsible for selecting the appropriate objects for the downstream task. This approach is also very similar to work in meta-learning which draws on the parallel between gating and neuromodulation in Neuroscience [2]. It would have been interesting to have both a conceptual and experimental comparison between the two. ### References [1] Locatello, F., Weissenborn, D., Unterthiner, T., Mahendran, A., Heigold, G., Uszkoreit, J., ... & Kipf, T. (2020). Object-centric learning with slot attention. *Advances in Neural Information Processing Systems*, *33*, 11525-11538. [2] Beaulieu, S., Frati, L., Miconi, T., Lehman, J., Stanley, K. O., Clune, J., & Cheney, N. (2020). Learning to continually learn. *arXiv preprint arXiv:2002.09571*.

Summary

The authors proposed a mutual information-based encoder that generates masks to gate inputs in order to either pass on minimal information for downstream tasks, or to remove any useful information that could be used to optimize the downstream loss in an adversarial setting. This is a general model that can be applied to a variety of tasks, and the authors focus on the application to reinforcement learning tasks.

Strengths

* Through experiments, the authors show that the idea of directly removing visual information in a learnable way through dynamic masking improves the downstream policy learning in terms of performance and training stability when there are visual distractors. * InfoGating is very general and could be applied to input or any intermediate representations, with good interpretability

Weaknesses

* training a minimax objective is usually not very stable, and the authors are not very clear in algorithm 1 or the appendix about the details of their minimax training. See questions. * The idea of masking inputs directly is reasonable, and experiments show that it improves performance and training stability. But it seems not so obviously advantageous on intermediate layers. The authors discussed in Section 6 and 7, but there's no clear statement or experiments towards this line.

Questions

For the minimax objective, are there any nested loops when updating the two sets of parameters? Is there code available?

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

The authors appropriately addresss ethical or social impacts of their research.

Summary

This paper introduces a novel masking technique called InfoGating for learning masks in contrastive loss settings with InfoNCE. The proposed approach is simple, well-motivated, and is evaluated in several RL setting including inverse dynamics models, Q-learning, and behavior cloning. Some ablation studies on regularization of the mask, as well as variants of the objective are also considered.

Strengths

* The proposed approach is a novel way for learning masking and is applied for RL tasks, whereas traditional approaches are typically evaluated for classification, or SSL tasks. Evaluation of the proposed approach covers multiple RL frameworks including Q learning, behavior cloning, and inverse dynamics. * The proposed method is stated very clearly including all hyperparameters for reproducibility, clear pseudocode for inverse dynamics, and a clear description of the approach with intuition.

Weaknesses

* My primary concern with this paper is that the proposed approach does not have sufficient evaluation not situate itself well with prior works. There are two notabl method comparisons, VIB (2017), and RCAD (2022). Neither of these approaches, however are masking-based approaches which is the main line of this work from my point of view. These methods are also not included in comparisons on behavioral cloning and Q-learning experiments. I would like to see additional comparisons on Q-learning and behavioral cloning, and perhaps some experiments comparing other masking techniques which from my understanding have been used for RL such as Q-learning: https://proceedings.neurips.cc/paper_files/paper/2022/hash/a0709efe5139939ab69902884ecad9c1-Abstract-Conference.html, https://proceedings.neurips.cc/paper_files/paper/2022/hash/802a4350ca4fced76b13b8b320af1543-Abstract-Conference.html. Although these approaches approach masking from a slightly different angle, my understanding is that they are still aimed at better observation learning. There are also works orthogonal to this that focus on sparse convolutions: https://openaccess.thecvf.com/content_CVPR_2020/html/Verelst_Dynamic_Convolutions_Exploiting_Spatial_Sparsity_for_Faster_Inference_CVPR_2020_paper.html. I think contrasting with these approaches may also distinguish the paper. * Additionally, cheetah run is only one of many tasks within the continuous control suite. It would be good for the authors to clarify why this task in particular, and if IG performs well across multiple additional tasks. Further, the reported cheetah run numbers are a fair bit lower than those traditionally reported for the return (usually 300+). Is there commentary on why the numbers are much lower in this setting, and why learning curves are not provided? * There are some followup questions about whether or not information gating without controlling for consistency in all layers is good. In particular, whether this may violate properties of the IB principle which say that the information will only decrease throughout the network.

Questions

* How does infogating satisfy or violate the IB principle in deep learning? Is it possible for the model to result in higher information in later layersdepending on how much gating happens in features of the network? Does this violate any benefits of the IB principle, for example better generalization? * How does IG perform compared with other methods for masking/information bottleneck on Q-learning and behavior cloning settings? * What is the connection between IG and adversarial masking?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

1: Your assessment is an educated guess. The submission is not in your area or the submission was difficult to understand. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

2 fair

Limitations

The proposed approach is not compared with other masking approaches for traditional augmentation and image classification settings.

Summary

In this paper, the authors propose InfoGating as a way to learn parsimonious representation that could achieve better generalization by being robust to noise and spurious correlations. Representations that identify the minimal information required for certain task. Those representations attempt to be robust to out-of-distribution observations. In contrast with previous works that applied a similar concept, i.e., information bottleneck, in later stage of the computation, InfoGating is applied to the input layer, which allows models to learn which information is key to solve the task. For example, what pixels matter for a given task. There are two approaches proposed for InfoGating namely cooperative and adversarial. The former aims to identifying minimal sufficient information. The latter to identifying any useful information.

Strengths

S1. The method is simple and well explained, easy to read. It sounds correct. S2. I found the experimental setup very interesting and well choose. I have seen methods motivated by similar concepts in self-supervised learning and efficient transformer [ATS: Adaptive token sampling for efficient vision transformers. Fayyaz et al. ECCV-22], but not for RL tasks. S3. Ablation results on feature space vs pixel InfoGating shows the importance of removing distractors early in the pipeline. Supporting the motivation of this work. S4. Appendix, Visualization and additional results are good and provide good insides to readers.

Weaknesses

W1. Since the InfoGaiting is presented as a general method to obtain more robust representations that can deal with out-of-distribution observations, I was expecting to see experiments in more dedicated settings like NICO++ [NICO++: Towards Better Benchmarking for Domain Generalization]

Questions

Q1. Can you please explain further what is the metric used on the experiments? Q2. In the experiments with inverse dynamic models, How consistent is the masking between two consecutive frames?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

2 fair

Limitations

I found the limitations are well covered.

Rating Update

The authors have addressed my main concerns with the experiments. They have added evaluations on held-out distracting backgrounds, showing that their method outperforms prior methods (though not for the "hard" level of noise). They have also shown that a more performant way to use their model at test time is to simply discard the masking network, removing their method's dependence on the masking network's generalization to out-of-distribution inputs. I have increased my rating.

Response to Author Rebuttal

Thank you for the response, and addressing many of my concerns with the extra experiments. * Given experiments on other masking approaches are completed, it would be nice to see the final results in the paper (or appendix) as I believe this is an important comparison. * Experiments with the Walker look convincing, and demonstrate the approach generalizes across environments. * My concern with the data processing inequality is with masking. If I denote input at layer $\ell$ as $X_{\ell}$, then data processing inequality says $MI(X_{\ell}; Y) \geq MI(X_{\ell+1}; Y)$. My question is whether this holds when applying masking i.e. $MI(M(X_{\ell}); Y) \geq MI(M(X_{\ell+1}); Y)$ for the mask function $M$. A potential violation could happen if there is no constraint on $M$, for example $M$ could mask everything at $\ell$ and nothing at $\ell+1$. Nonetheless, I will update my score to reflect positive results from the updated experiments.