SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Manan Tomar*  
University of Alberta
Riashat Islam  
McGill University
Matthew E. Taylor  
University of Alberta
Sergey Levine  
University of California, Berkeley
Philip Bachman  
Microsoft Research Montreal
Informational parsimony provides a useful inductive bias for learning representations that achieve better generalization by being robust to noise and spurious correlations. We propose information gating as a way to learn parsimonious representations that identify the minimal information required for a task. When gating information, we can learn to reveal as little information as possible so that a task remains solvable, or hide as little information as possible so that a task becomes unsolvable. We gate information using a differentiable parameterization of the signal-to-noise ratio, which can be applied to arbitrary values in a network, e.g., erasing pixels at the input layer or activations in some intermediate layer. When gating at the input layer, our models learn which visual cues matter for a given task. When gating intermediate layers, our models learn which activations are needed for subsequent stages of computation. We call our approach InfoGating. We apply InfoGating to various objectives such as multi-step forward and inverse dynamics models, Q-learning, and behavior cloning, highlighting how InfoGating can naturally help in discarding information not relevant for control. Results show that learning to identify and use minimal information can improve generalization in downstream tasks. Policies based on InfoGating are considerably more robust to irrelevant visual features, leading to improved pretraining and finetuning of RL models.
Figure 1: InfoGating at the input layer predicts a mask conditioned on the original input and then processes the masked input to compute representations that minimize a downstream loss. The mask is encouraged to remove as much of the original input as possible without hurting performance on the downstream task. The mask is applied by taking a convex combination of the input and Gaussian noise, weighted by the mask.
Pretraining models on large, diverse datasets and transferring their representations to downstream applications is becoming common practice in deep learning. Such representations should capture useful information from available data while ignoring irrelevant features and noise [7, 8]. For instance, an object recognition model may achieve high accuracy on its training set while strongly relying on "spurious" correlations between background features and important objects. When transferred to
*Correspondence to manan.tomar@gmail.com
new data with the same objects but a different distribution of backgrounds, performance may suffer. Similarly, a policy for vision-based robot control [23, 10] may perform well in its original training environment, but fail catastrophically when transferred to a new setting [28] where all task-relevant features are the same but minor background features differ from training.
A promising principle for encouraging robustness to out-of-distribution observations is informational parsimony: learning features that contain minimal information [39] about the input while still containing enough information to solve the task(s) of interest. Such regularization is usually based on imposing an information bottleneck (IB) [38] on the network that tries to minimize the amount of information flowing from the input, through the bottleneck, to the output prediction. 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. To that end, we consider restricting the information used throughout a computation rather than just restricting the information that comes out of a computation.
In this paper, we propose learning input-conditioned functions that gate the flow of information through a model. For example, we can consider what happens when one inserts an IB near the beginning of a model's computation rather than near the end. In this case, we train the information gating functions to minimize information flow from the input into the rest of the model while still permitting the model to solve the task(s) of interest. We implement this gating using a differentiable parameterization of the signal-to-noise ratio. As the noise level goes up, the signal level goes down, and hence the model displays reduced dependencies on the input. Gating functions can be learnt in conjunction with any downstream loss corresponding to some task of interest. For example, the downstream loss could be a standard contrastive loss for self-supervised learning or an inverse dynamics loss when learning representations for reinforcement learning (RL). In both cases, we can learn gating functions that reveal minimal information used to optimize the loss or learn adversarial gating functions that will remove any information that can be used to optimize the loss. We primarily focus on InfoGating representations used for learning control policies in RL. A natural focus of the RL paradigm is on learning a mapping from observations to actions, thus discarding most information not useful for control. We bring forth such a notion of only capturing what the agent can affect through InfoGating. Using background distractors and multi-object interactions as a form of noise/irrelevant features, we see that InfoGating is able to remove almost all of the irrelevant information from the pixels, leading to better out-of-distribution generalization in the presence of noise and better in-distribution generalization in the presence of multiple task-irrelevant objects.
Our main contributions are as follows: 1. A general purpose, practical framework for informational parsimony called InfoGating that can restrict information flow throughout a computation to learn robust representations. 2. Qualitative analyses on the properties of the gating functions learned with InfoGating that show they are semantically meaningful and enhance interpretability. 3. Quantitative analyses of applying InfoGating in the context of various downstream objectives which show clear benefits in terms of improved generalization performance.
Information Bottleneck. Much prior work in learning robust representations has stemmed from the idea of imposing an information bottleneck (IB) [38, 39]. Imposing an IB involves maximizing performance on the downstream task, while removing as much information about the input as possible from a network's internal representations. Typical approaches based on Variational IB [2] achieve this by learning stochastic representations that are constrained to be close to a standard Gaussian prior. Variants of dropout [36, 19] can have a similar effect (i.e., adding noise to representations used in the network) but without explicitly maximizing noise/minimizing information in the representations. While IB is typically applied to later stages of a network, InfoGating can be seen as enforcing an IB at the input itself, which is less explored in the IB literature. Like Variational IB methods, and unlike typical dropout-based methods, InfoGating learns how much noise can be added without hurting performance on the downstream task(s). Information Dropout [1] is perhaps the most similar prior work to InfoGating. Information Dropout imposes IBs at multiple points in a network's computation and the IBs are implemented using techniques based on variational dropout [19]. A key distinction between InfoGating and Information Dropout is that InfoGating works post-hoc. With InfoGating a model can consider a computation it has already performed and predict ways in which it could produce the same result more parsimoniously.
Masked Image Modelling. There has been a lot of recent work in self-supervised learning where the input is masked and the model is tasked with reconstructing the missing bits [3, 15, 41]. Most of these works operate with a fixed, random masking scheme, instead of learning it directly through the downstream loss. The amount of information masked is also fixed a priori whereas we seek to mask as much of the input as possible, while still retaining enough information to solve the task. Some related ideas include learning adversarial augmentations for producing new views of an input during contrastive learning [37], while also learning masks that essentially segment the input space [35].
Representation Learning in Reinforcement Learning. Prior work on learning representations for RL includes temporally contrastive learning [33, 27], one-step inverse models [32], learning through next state reconstruction [14], bisimulation metrics [9], and many more. Some of these objectives have been shown to be useful in learning robust representations [40], and we see our method as complementary to these approaches since InfoGating works with any downstream loss function, without needing knowledge of specific values like reward, next state, etc. For instance, bisimulation and task-informed abstractions [11] learn to compress the observation space using reward information, while recent work on learning Denoised MDPs [42] aims at regularizing next state reconstruction using a variational objective. On the other hand, InfoGating remains agnostic to choices of downstream objectives.
This section describes the primary downstream losses we use with InfoGating in this paper, namely an InfoNCE based contrastive loss and a multi-step inverse dynamics loss.
Mutual Information Estimation via InfoNCE. Given  $\mathbf{x}$  and  $\mathbf{z}$  as two random variables, their mutual information can be defined as the decrease in uncertainty when observing  $\mathbf{x}$  given  $\mathbf{z}$ , compared to just observing  $\mathbf{x}$ :  $I(\mathbf{x},\mathbf{z}) = H(\mathbf{x}) - H(\mathbf{x}|\mathbf{z})$ , where  $H$  is the Shannon entropy. InfoNCE [31], based on Noise Contrastive Estimation [13], computes a lower bound on  $I(\mathbf{z}^1,\mathbf{z}^2)$ , where  $\mathbf{z}^1$  and  $\mathbf{z}^2$  are two representations of the input  $\mathbf{x}$  produced by some encoder  $f(\mathbf{x})$ . Specifically, the bound is optimized by discriminating "positive" and "negative" pairs:
$$
\mathcal {L} _ {\text {I n f o N C E}} = \mathbb {E} _ {Z ^ {-}} \left[ \log \frac {e ^ {\psi \left(\mathbf {z} ^ {1} , \mathbf {z} ^ {2}\right)}}{e ^ {\psi \left(\mathbf {z} ^ {1} , \mathbf {z} ^ {2}\right)} + \sum_ {\mathbf {z} ^ {-} \in Z ^ {-}} e ^ {\psi \left(\mathbf {z} ^ {1} , \mathbf {z} ^ {-}\right)}} \right], \tag {1}
$$
where  $\psi$  is a pairwise, scalar-valued function of the representations and  $Z^{-}$  is a batch of "negative samples". Typically, self-supervised contrastive learning uses data augmentation such as random cropping and color jittering to define two augmented views  $(\mathbf{z}^{1},\mathbf{z}^{2})$  of a given input as "positives" while views  $\mathbf{z}^{-}$  of other inputs are treated as "negatives". The InfoNCE objective then encourages the representations of positive views of  $\mathbf{x}$  to be close (i.e.,  $\psi (\mathbf{z}^1,\mathbf{z}^2)$  is high), while pushing apart the negatives (i.e.,  $\psi (\mathbf{z}^1,\mathbf{z}^{-})$  is low) [4, 5].
Multi-Step Inverse Dynamics Models. Multi-step inverse dynamics models [8] predict the action(s) that took an agent from some observation  $\mathbf{x}_t$  to some future observation  $\mathbf{x}_{t + k}$ , attempting to learn a useful representation of the observations. This resembles learning goal-conditioned policies through relabelling future observations as achieved goals. Although the trivial case of  $k = 1$  (a one-step model) does not capture long term dependencies, recent work has shown that multi-step models are able to capture more information that may be useful for controlling the agent [22]. Multi-step inverse models can be used to learn representations by simply predicting the actions  $(\mathbf{a}_t,\dots,\mathbf{a}_{t + k - 1})$  conditioned on  $\mathbf{x}_t$  and  $\mathbf{x}_{t + k}$ . Actions can be predicted using a standard max likelihood objective or with a contrastive reconstruction objective which maximizes the InfoNCE-based lower bound on  $I((\mathbf{x}_t,\mathbf{x}_{t + k}),(a_t,\dots,a_{t + k - 1}))$ . The learning objective in this case looks as follows:
$$
\mathcal {L} = \operatorname {I n f o N C E} \left(\left(\mathbf {z} _ {t}, \mathbf {z} _ {t + k}\right), \mathbf {y} _ {t: t + k - 1}\right),
$$
where  $\mathbf{z}_t$  and  $\mathbf{z}_{t + k}$  are representations computed from  $\mathbf{x}_t$  and  $\mathbf{x}_{t + k}$ , and  $\mathbf{y}_{t:t + k - 1}$  is a representation computed from  $(\mathbf{a}_t,\dots,\mathbf{a}_{t + k - 1})$ . The negative samples in this case come from randomly sampling action sequences of the appropriate length from the agent's collected experience. In the simplest case, the representation  $\mathbf{y}_{t:t + k - 1}$  may depend only on  $\mathbf{a}_t$ .
4 Information Gating
We present two approaches to InfoGating that we describe as cooperative and adversarial. Cooperative InfoGating involves keeping as little information as possible to obtain good performance on the downstream task. Adversarial InfoGating involves removing as little information as possible to preclude good performance on the downstream task. We include both approaches in this paper since identifying minimal sufficient information (cooperative) and identifying any useful information (adversarial) may have different affects depending on the downstream loss and the application domain.
The InfoGating approach has two major components. The first is an encoder of the input  $\mathbf{x}$ ,  $f(\mathbf{x})$ . The second is an information gating function,  $ig(\mathbf{x})$  (see Figure 1). In principle, we can gate information passing through any layer in the representation network that computes the encoding  $\mathbf{z} = f(\mathbf{x})$ . The  $ig(\mathbf{x})$  function provides continuous-valued masks (values in [0, 1]) that describe where to erase information from the computation graph for  $f(\mathbf{x})$ . The shape/size of the output of  $ig(\mathbf{x})$  will depend on where we want to gate information. For instance, if we wish to gate information in the input pixel space,  $ig(x)$  masks are the size of the image. In general, the goal of  $ig(\mathbf{x})$  is to erase as much information from  $f(\mathbf{x})$  as possible without hurting task performance.
4.1 Cooperative InfoGating
We primarily focus on cooperative InfoGating in this paper and we consider two cooperative variants: gating in the input space or in the feature space.
InfoGating in Input Space. We apply InfoGating to an input  $\mathbf{x}$  by taking a simple convex combination of  $\mathbf{x}$  and random Gaussian noise  $\epsilon$ . The combination weights are given by  $ig(\mathbf{x})$ :
$$
\mathbf {x} ^ {i g} = i g (\mathbf {x}) \odot \mathbf {x} + (1 - i g (\mathbf {x})) \odot \epsilon , \tag {2}
$$
where  $\mathbf{x}^{ig}$  denotes the info-gated input and  $\odot$  denotes element-wise multiplication. An all-zero mask corresponds to complete noise (i.e., erasing all information from the input), while an all-one mask corresponds to keeping the original input. The function  $ig(x)$  is learnt using the same downstream objective that is used to learn  $f(\mathbf{x})$ . We encourage masks to remove information from the input by minimizing their L1 norm. This tends to produce sparse masks due to properties of the L1 norm [30]. The overall objective for learning with InfoGating is
$$
\mathcal {L} _ {i g, f} = \mathcal {L} _ {\text {t a s k}} \left(f (\mathbf {x} ^ {i g})\right) + \lambda | | i g (\mathbf {x}) | | _ {1}, \tag {3}
$$
where  $\mathcal{L}_{\mathrm{task}}$  refers to any objective through which a useful representation  $\mathbf{z}$  can be learnt. Note that when doing InfoGating, we minimize the downstream loss for the masked input  $\mathbf{x}^{ig}$  (first term), instead of the original input  $\mathbf{x}$ , while masking out as much of the input as possible through the L1 penalty (second term). The  $\lambda$  coefficient is a hyperparameter that controls how much of the input is masked. In principle, any loss function can be used in conjunction with InfoGating and we consider multiple downstream objectives: contrastive learning of dynamic, Q-learning, and behavior cloning. We describe the overall procedure in Algorithm 1, which assumes contrastive multi-step inverse
dynamics modeling as the downstream objective.
InfoGating in Feature Space. As an alternative to directly gating the input pixels, we can also consider a variant where  $f(\mathbf{x})$  is masked instead of the input. Consider the following masking:
$$
\mathbf {z} ^ {i g} = i g (\mathbf {x}) \odot \mathbf {z} + (1 - i g (\mathbf {x})) \odot \epsilon , \quad \text {w h e r e} \quad \mathbf {z} = f (\mathbf {x}). \tag {4}
$$
The training objective remains the same as in the pixel-level case, i.e., minimize the downstream loss  $\mathcal{L}_{\mathrm{task}}(\mathbf{z}^{ig})$  for the masked representation  $\mathbf{z}^{ig}$ , while masking as much of  $\mathbf{z}$  as possible (Figure 2). This version is closely related to the deep variational information bottleneck (VIB), which minimizes the KL divergence between the (stochastic) representation  $\mathbf{z}$  and a prior distribution, typically chosen to be a standard Gaussian. Roughly, the values in  $ig(\mathbf{x})$  can be seen as specifying how many steps of for
Figure 2: Feature InfoGating. gates representations produced by the encoder. The gating network and the encoder are jointly trained to minimize the downstream loss, while the gating network is also encouraged to mask as much of  $\mathbf{z}$  as possible.
ward diffusion to run in a Gaussian diffusion process initiated at  $\mathbf{z}$ , where values near zero correspond
Algorithm 1 InfoGating Pseudocode
1: Input encoder  $f$ , masking net  $ig$   
2: for  $(\mathbf{x}, \mathbf{x}_{t+k}, \mathbf{a}_t) \in \text{loader do}$   
3:  $\mathbf{x}_t, \mathbf{x}_{t+k} = \text{aug}(\mathbf{x}_t)$ ,  $\text{aug}(\mathbf{x}_{t+k})$  # random crop augmentation  
4:  $\mathbf{x}_t^{ig}, \mathbf{x}_{t+k}^{ig} = ig(\mathbf{x}_t)$ ,  $ig(\mathbf{x}_{t+k})$  # get infogates  
5:  $\mathbf{x}_t^{ig} = ig(\mathbf{x}_t) \odot \mathbf{x}_t + (1 - ig(\mathbf{x}_t)) \odot \epsilon$  # infogate current state  
6:  $\mathbf{x}_{t+k}^{ig} = ig(\mathbf{x}_{t+k}) \odot \mathbf{x}_{t+k} + (1 - ig(\mathbf{x}_{t+k})) \odot \epsilon$  # infogate future state  
7:  $\mathbf{z}_t, \mathbf{z}_{t+k} = f(\mathbf{x}_t^{ig}), f(\mathbf{x}_{t+k}^{ig})$  # get encodings  
8: if Cooperative then  
9: Run Adam update on overall loss:  $\mathcal{L}_{ig,f} = \text{InfoNCE}(\mathbf{z}_t, \mathbf{z}_{t+k}, \mathbf{a}_t) + \lambda ||ig(\mathbf{x}_t)||_1 + \lambda ||ig(\mathbf{x}_{t+k})||_1$   
10: else if Adversarial then  
11: Run Adam update on masking loss:  $\mathcal{L}_{ig} = -\text{InfoNCE}(\mathbf{z}_t, \mathbf{z}_{t+k}, \mathbf{a}_t) - \lambda ||ig(\mathbf{x}_t)||_1 - \lambda ||ig(\mathbf{x}_{t+k})||_1$   
12: Run Adam update on encoder loss:  $\mathcal{L}_f = \text{InfoNCE}(\mathbf{z}_t, \mathbf{z}_{t+k}, \mathbf{a}_t)$   
13: end if  
14: end for
to running more steps of diffusion and thus sampling from a distribution that is closer to a standard Gaussian in terms of KL divergence [18]. Thus, like VIB methods, we aim to minimize  $I(\mathbf{x},\mathbf{z})$  while maximizing the task performance of  $\mathbf{z}$ . However, feature space InfoGating uses a different optimization objective and parameterization compared to existing VIB approaches. We leave similar variations of InfoGating where arbitrary layers in a computation graph are gated for future work†.
4.2 Adversarial InfoGating
The versions of InfoGating discussed above are cooperative objectives, where the gating network and the encoder work together to lower the overall loss. This leads to finding representations that capture the minimal sufficient information for a given task. We can also turn this formulation on its head, yielding an adversarial objective. In this case, the gating network is tasked with discovering masks that, when used by the encoder, lead to maximizing its loss. Instead of encouraging the masks to erase as much of the input as possible, the adversarial objective encourages the masks to remove as little information as possible while minimizing the encoder's performance on the downstream task. We write the adversarial objective for  $ig(x)$  and  $f(x)$  as:
$$
\mathcal {L} _ {i g} = - \mathcal {L} _ {\text {t a s k}} \left(f \left(\mathbf {x} ^ {i g}\right)\right) - \lambda \left| \left| i g (\mathbf {x}) \right| \right| _ {1}, \quad \mathcal {L} _ {f} = \mathcal {L} _ {\text {t a s k}} \left(f \left(\mathbf {x} ^ {i g}\right)\right) \tag {5}
$$
This gives rise to a min-max objective w.r.t. the encoder and the gating network while cooperative InfoGating corresponds to a min-min objective (see lines 11-12 in Algorithm 1).
5 Experiment Setup
To understand whether InfoGating can consistently focus on the minimal possible information required for control, we test generalization performance when using InfoGating with three different downstream control objectives: 1) contrastive dynamics models (both inverse and forward models), 2) Q-learning (TD-based critic updates), and 3) behavior cloning. We choose these three objectives since they remove irrelevant information to varying degrees by default — for example, dynamics models can capture a lot of task-irrelevant information, while behavior cloning models are meant to only contain information that is useful for predicting the optimal policy. Through our experiments, we evaluate whether all three of these objectives can benefit from InfoGating, and to what degree $\ddagger$ .
$\dagger$  Constructing differentiable upper bounds on compute cost is one possibility, similar to DARTS [24].
Since the multi-step inverse model produces the best results, we use the same for testing the feature InfoGating and adversarial InfoGating versions as well.
Table 1: Multi-step Inverse Dynamics with InfoGating. Comparing performance in the presence of noisy distractors. We report returns achieved by a policy produced by behavior cloning on top of pretrained representations in cheetah-run. Extended results are provided in Appendix C. The "inv" model is our baseline with pretraining via multi-step inverse dynamics. The "w/ Rand" model adds random info gating during pretraining and the "w/ IG" model adds learned info gating during pretraining (this is our method). The easy/medium/hard settings add different levels of distractor noise. Results are for 3 seeds each, with mean and std. dev. reported.
case	IQL	DRIML	Inv	Inv w/ Dropout	Inv w/ VIB	Inv w/ RCAD	Inv w/Rand	Inv w/ IG
expert
easy	10.7 ± 8.0	0.9 ± 0.1	42.6 ± 36.3	11.79 ± 3.5	97.1 ± 17.9	25.0 ± 3.6	21.0 ± 15.6	176.2 ± 9.1
medium	29.2 ± 28.8	1.0 ± 0.6	29.5 ± 31.9	73.3 ± 12.3	38.9 ± 16.9	19.9 ± 24.8	7.2 ± 6.12	97.0 ± 5.7
hard	8.8 ± 1.9	13.8 ± 7.9	2 ± 0.6	5.9 ± 5.7	34.0 ± 19.4	1.4 ± 0.4	4.2 ± 0.7	44.8 ± 18.4
overall	16.2	5.2	24.7	30.3	56.6	15.4	10.8	106.0
medium
easy	42.6 ± 26.4	2.6 ± 0.2	45.6 ± 23.2	28.0 ± 23.6	113.5 ± 11.8	29.6 ± 24.5	42.0 ± 31.8	133.0 ± 12.8
medium	31.4 ± 24.7	86.6 ± 55.0	2 ± 0.8	39.9 ± 39.4	44.0 ± 22.8	4.0 ± 2.2	85.4 ± 13.2	92.0 ± 35.2
hard	15.2 ± 7.7	3.1 ± 1.5	10.7 ± 6.0	18.8 ± 15.2	6.1 ± 1.8	4.2 ± 1.6	5.3 ± 1.9	3.0 ± 1.43
overall	29.7	30.7	19.4	28.9	54.5	12.6	44.2	76.0
medium-expert
easy	42.4 ± 31.4	12.9 ± 12.7	25.5 ±	45.7 ± 7.4	130.4 ± 32.0	51.8 ± 38.1	10.4 ± 9.9	158.0 ± 21.8
medium	39.6 ± 24.7	22.3 ± 13.6	14.2 ± 11.2	15.6 ± 8.1	52.2 ± 26.0	5.4 ± 3.9	24.8 ± 34.4	89.2 ± 7.0
hard	10.2 ± 4.0	5.2 ± 2.6	3.4 ± 4.8	16.2 ± 14.4	31.2 ± 38.7	9.1 ± 3.4	2.9 ± 1.2	38.8 ± 37.5
overall	30.7	13.4	14.3	25.8	71.2	22.1	12.7	95.3
We test 1) and 2) on the offline visual D4RL domain [25] and 3) on the Kitchen [12] manipulation domain. InfoGating can also be used in non-RL settings like self-supervised learning, as discussed in Appendix A. We now describe the details of each use of InfoGating and the corresponding experimental results. Hyper-parameters are listed in Appendix E.
5.1 Inverse Dynamics Models