Drop-Bottleneck: Learning Discrete Compressed Representation for Noise-Robust Exploration

1 INTRODUCTION . Data with noise or task-irrelevant information easily harm the training of a model ; for instance , the noisy-TV problem ( Burda et al. , 2019a ) is one of well-known such phenomena in reinforcement learning . If observations from the environment are modified to contain a TV screen , which changes its channel randomly based on the agent ’ s actions , the performance of curiosity-based exploration methods dramatically degrades ( Burda et al. , 2019a ; b ; Kim et al. , 2019 ; Savinov et al. , 2019 ) . The information bottleneck ( IB ) theory ( Tishby et al. , 2000 ; Tishby & Zaslavsky , 2015 ) provides a framework for dealing with such task-irrelevant information , and has been actively adopted to exploration in reinforcement learning ( Kim et al. , 2019 ; Igl et al. , 2019 ) . For an input variable X and a target variable Y , the IB theory introduces another variable Z , which is a compressed representation of X . The IB objective trains Z to contain less information about X but more information about Y as possible , where the two are quantified by mutual information terms of I ( Z ; X ) and I ( Z ; Y ) , respectively . IB methods such as Variational Information Bottleneck ( VIB ) ( Alemi et al. , 2017 ; Chalk et al. , 2016 ) and Information Dropout ( Achille & Soatto , 2018 ) show that the compression of the input variable X can be done by neural networks . In this work , we propose a novel information bottleneck method named Drop-Bottleneck that compresses the input variable by discretely dropping a subset of its input features that are irrelevant to the target variable . Drop-Bottleneck provides some nice properties as follows : • The compression term of Drop-Bottleneck ’ s objective is simple and is optimized as a tractable solution . • Drop-Bottleneck provides a deterministic compressed representation that still maintains majority of the learned indistinguishability i.e . compression . It is useful for inference tasks that require the input representation to be consistent and stable . • Drop-Bottleneck jointly trains a feature extractor and performs feature selection , as it learns the feature-wise drop probability taking into account each feature dimension ’ s relevance to the target task . Hence , unlike the compression provided by most neural network-based IB methods , our deterministic representation reduces the feature dimensionality , which makes the following inference better efficient with less amount of data . • Compared to VIB , both of Drop-Bottleneck ’ s original ( stochastic ) and deterministic compressed representations can greatly improve the robustness to adversarial examples . Based on the newly proposed Drop-Bottleneck , we design an exploration method that is robust against noisy observations in very sparse reward environments for reinforcement learning . Our exploration maintains an episodic memory and generates intrinsic rewards based on the predictability of new observations from the compressed representations of the ones in the memory . As a result , our method achieves state-of-the-art performance on multiple environments of VizDoom ( Kempka et al. , 2016 ) and DMLab ( Beattie et al. , 2016 ) . We also show that combining our exploration method with VIB instead of Drop-Bottleneck degrades the performance by meaningful margins . Additionally , we empirically compare with VIB to show Drop-Bottleneck ’ s superior robustness to adversarial examples and ability to reduce feature dimensionality for inference with ImageNet dataset ( Russakovsky et al. , 2015 ) . We also demonstrate that Drop-Bottleneck ’ s deterministic representation can be a reasonable replacement for its original representation in terms of the learned indistinguishability , with Occluded CIFAR dataset ( Achille & Soatto , 2018 ) . 2 RELATED WORK . 2.1 INFORMATION BOTTLENECK METHODS . There have been a number of IB methods that are approximations or special forms of the original IB objective . Variational Information Bottleneck ( VIB ) ( Alemi et al. , 2017 ) approximates the original IB objective by establishing variational bounds on the compression and prediction terms . Chalk et al . ( 2016 ) propose the same variational bound on the IB objective in the context of sparse coding tasks . Conditional Entropy Bottleneck ( CEB ) and Variational Conditional Entropy Bottleneck ( VCEB ) ( Fischer , 2020 ; Fischer & Alemi , 2020 ) use an alternative form of the original IB objective derived under the Minimum Necessary Information ( MNI ) criterion to preserve only a necessary amount of information . The IB theory ( Tishby et al. , 2000 ) has been used for various problems that require restriction of information or dealing with task-irrelevant information . Information Dropout ( Achille & Soatto , 2018 ) relates the IB principle to multiple practices in deep learning , including Dropout , disentanglement and variational autoencoding . Moyer et al . ( 2018 ) obtain representations invariant to specific factors under the variational autoencoder ( VAE ) ( Kingma & Welling , 2013 ) and VIB frameworks . Amjad & Geiger ( 2019 ) discuss the use of IB theory for classification tasks from a theoretical point of view . Dai et al . ( 2018 ) employ IB theory for compressing neural networks by pruning neurons in networks . Schulz et al . ( 2020 ) propose an attribution method that determines each input feature ’ s importance by enforcing compression of the input variable via the IB framework . Similar to our goal , some previous research has proposed variants of the original IB objective . Deterministic information bottleneck ( DIB ) ( Strouse & Schwab , 2017 ) replaces the compression term with an entropy term and solves the new objective using a deterministic encoder . Nonlinear information bottleneck ( NIB ) ( Kolchinsky et al. , 2019 ) modifies the IB objective by squaring the compression term and uses a non-parametric upper bound on the compression term . While DIB is always in the deterministic form , we can flexibly choose the stochastic one for training and the deterministic one for test . Compared to NIB , which is more computationally demanding than VIB due to its non-parametric upper bound , our method is faster . 2.2 REINFORCEMENT LEARNING WITH INFORMATION BOTTLENECK METHODS . The IB theory has been applied to several reinforcement learning ( RL ) tasks . Variational discriminator bottleneck ( Peng et al. , 2019 ) regulates the discriminator ’ s accuracy using the IB objective to improve adversarial training , and use it for imitation learning . Information Bottleneck Actor Critic ( Igl et al. , 2019 ) employs VIB to make the features generalize better and encourage the compression of states as input to the actor-critic algorithm . Curiosity-Bottleneck ( Kim et al. , 2019 ) employs the VIB framework to train a compressor that compresses the representation of states , which is still informative about the value function , and uses the compressiveness as exploration signals . InfoBot ( Goyal et al. , 2019 ) proposes a conditional version of VIB to improve the transferability of a goal-conditioned policy by minimizing the policy ’ s dependence on the goal . Variational bandwidth bottleneck ( Goyal et al. , 2020 ) uses a modified , conditional version of VIB , and solves RL tasks with privileged inputs ( i.e . valuable information that comes with a cost ) . Our exploration method differs from these methods in two aspects . First , we propose a new information bottleneck method that is not limited to exploration in RL but generally applicable to the problems for which the IB theory is used . Second , our method generates exploration signals based on the noise-robust predictability i.e . the predictability between noise-robust representations of observations . 3 DROP-BOTTLENECK . 3.1 PRELIMINARIES OF INFORMATION BOTTLENECK . Given an input random variable X , the information bottleneck ( IB ) framework ( Tishby et al. , 2000 ; Tishby & Zaslavsky , 2015 ) formalizes a problem of obtaining X ’ s compressed representation Z , which still and only preserves information relevant to the target variable Y . Its objective function is minimize−I ( Z ; Y ) + βI ( Z ; X ) , ( 1 ) where β is a Lagrangian multiplier . The first and second terms are prediction and compression terms , respectively . The prediction term I ( Z ; Y ) encourages Z to preserve task-relevant information while the compression term I ( Z ; X ) compresses the input information as much as possible . As reviewed in the previous section , there have been several IB methods ( Alemi et al. , 2017 ; Chalk et al. , 2016 ; Achille & Soatto , 2018 ; Strouse & Schwab , 2017 ; Kolchinsky et al. , 2019 ) , among which the ones derived using variational inference such as Variational Information Bottleneck ( VIB ) ( Alemi et al. , 2017 ) have become dominant due to its applicability to general problems . 3.2 DROP-BOTTLENECK . We propose a novel information bottleneck method called Drop-Bottleneck ( DB ) , where the input information is compressed by discretely dropping a subset of input features . Thus , its compression objective is simple and easy to optimize . Moreover , its representation is easily convertible to a deterministic version for inference tasks ( Section 3.3 ) , and it allows joint training with a feature extractor ( Section 3.4 ) . While discrete dropping of features has been explored by prior works including Dropout ( Srivastava et al. , 2014 ) , DB differs in that its goal is to assign different drop probabilities to feature variables based on their relevance to the target variable . For an input variableX = [ X1 , . . . , Xd ] and a drop probability p = [ p1 , . . . , pd ] ∈ [ 0 , 1 ] d , we define its compressed representation as Z = Cp ( X ) = [ c ( X1 , p1 ) , . . . , c ( Xd , pd ) ] such that c ( Xi , pi ) = b · Bernoulli ( 1− pi ) ·Xi , where b = d d−∑k pk , ( 2 ) for i = 1 , . . . , d. That is , the compression procedure drops the i-th input feature ( i.e . replaced by zero ) with probability pi . Since the drop probability is to be learned , we use a scaling factor b to keep the scale of Z constant . We use a single scaling factor for all feature dimensions in order to preserve the relative scales between the features . With the definition in Equation ( 2 ) , we derive the compression term of DB to minimize as I ( Z ; X ) = d∑ i=1 I ( Zi ; X1 , . . . , Xd|Z1 , . . . , Zi−1 ) ( 3 ) = d∑ i=1 [ I ( Zi ; Xi|Z1 , . . . , Zi−1 ) + I ( Zi ; X1 , . . . , Xd \Xi|Z1 , . . . , Zi−1 , Xi ) ] ( 4 ) = d∑ i=1 I ( Zi ; Xi|Z1 , . . . , Zi−1 ) ≤ d∑ i=1 I ( Zi ; Xi ) = Î ( Z ; X ) ( 5 ) using that Zi⊥X1 , . . . , Xi−1 , Xi+1 , . . . , Xd|Z1 , . . . , Zi−1 , Xi and Zi⊥Z1 , . . . , Zi−1|Xi . Note that Î ( Z ; X ) − I ( Z ; X ) = ( ∑d i=1H ( Zi ) ) − H ( Z1 , . . . , Zd ) = TC ( Z ) where TC ( Z ) is the total correlation of Z and H ( · ) denotes the entropy , and Î ( Z ; X ) = I ( Z ; X ) if X1 , . . . , Xd are independent . To provide a rough view on the gap , due to the joint optimization with the compression term Î ( Z ; X ) and the prediction term I ( Z ; Y ) , Z becomes likely to preserve less redundant and less correlated features , and TC ( Z ) could decrease as the optimization progresses . Finally , DB ’ s new compression term , Î ( Z ; X ) , is computed as Î ( Z ; X ) = d∑ i=1 I ( Zi ; Xi ) = d∑ i=1 ( H ( Xi ) −H ( Xi|Zi ) ) ( 6 ) = d∑ i=1 ( H ( Xi ) − pi ·H ( Xi|Zi = 0 ) − ( 1− pi ) ·H ( Xi|Zi = bXi ) ) ( 7 ) ≈ d∑ i=1 ( H ( Xi ) − pi ·H ( Xi ) − ( 1− pi ) · 0 ) = d∑ i=1 H ( Xi ) ( 1− pi ) . ( 8 ) From Equation ( 7 ) to Equation ( 8 ) , we use the two ideas : ( i ) H ( Xi|Zi = 0 ) = H ( Xi ) because Zi = 0 means it contains no information about Xi , and ( ii ) H ( Xi|Zi = bXi ) = 0 because Zi = bXi means Zi preserves the feature ( i.e . Zi fully identifies Xi ) and thus their conditional entropy becomes zero . Importantly , DB ’ s compression term is computed as the simple tractable expression in Equation ( 8 ) . As the goal of the compression term is to penalize I ( Z ; X ) not H ( X ) , the drop probability p is the only parameter optimized with our compression term . Thus , each H ( Xi ) can be computed with any entropy estimation method such as the binning-based estimation , which involves quantization for continuous Xi , since the computation has no need to be differentiable . However , one issue of Equation ( 8 ) is that Z is not differentiable with respect to p as Bernoulli distributions are not differentiable . We thus use the Concrete relaxation ( Maddison et al. , 2017 ; Jang et al. , 2016 ) of the Bernoulli distribution to update p via gradients from Z : Bernoulli ( p ) ≈ σ ( 1 λ ( log p− log ( 1− p ) + log u− log ( 1− u ) ) ) , ( 9 ) where u ∼ Uniform ( 0 , 1 ) and λ is a temperature for the Concrete distribution . Intuitively , p is trained to assign a high drop probability to the feature that is irrelevant to or redundant for predicting the target variable Y .

Summary: This paper proposes an information bottleneck method, Drop-Bottleneck, that allows the input to be compressed by dropping each input feature with probability p_i. The model then learns the drop probability vector p = [p_1, ... , p_n], where dropping "redundant" features will reduce the "compression penalty" term I(XZ). The approach is demonstrated in experiments in (1) robust exploration setting for RL, (2) adversarial attacks on ImageNet, and (3) an experiment showing that their approach is able to maintain performance on ImageNet with reduced dimensionality.