The Variational Bandwidth Bottleneck: Stochastic Evaluation on an Information Budget

1 INTRODUCTION . A model that generalizes effectively should be able to pick up on relevant cues in the input while ignoring irrelevant distractors . For example , if one want to cross the street , one should only pay attention to the positions and velocities of the cars , disregarding their color . The information bottleneck ( Tishby et al. , 2000 ) formalizes this in terms of minimizing the mutual information between the bottleneck representation layer with the input , while maximizing its mutual information with the correct output . This type of input compression can improve generalization ( Tishby et al. , 2000 ) , and has recently been extended to deep parametric models , such as neural networks where it has been shown to improve generalization ( Achille & Soatto , 2016 ; Alemi et al. , 2016 ) . The information bottleneck is generally intractable , but can be approximated using variational inference ( Alemi et al. , 2016 ) . This variational approach parameterizes the information bottleneck model using a neural network ( i.e. , an encoder ) . While the variational bound makes it feasible to train ( approximate ) information bottleneck layers with deep neural networks , the encoder in these networks – the layer that predicts the bottleneck variable distribution conditioned on the input – must still process the full input , before it is compressed and irrelevant information is removed . The encoder itself can therefore fail to generalize , and although the information bottleneck minimizes mutual information with the input on the training data , it might not compress successfully on new inputs . To 1 Mila , University of Montreal,2 Deepmind , 3 University of California , Berkeley , : anirudhgoyal9119 @ gmail.com address this issue , we propose to divide our input into two categories : standard input and privileged input , and then we aim to design a bottleneck that does not need to access the privileged input before deciding how much information about the input is necessary . The intuition behind not accessing the privileged input is twofold : ( a ) we might want to avoid accessing the privileged input because we want to generalize with respect to it ( and therefore compress it ) ( b ) we actually would prefer not to access it ( as this input could be costly to obtain ) . The objective is to minimize the conditional mutual information between the bottleneck layer and the privileged input , given the standard input . This problem statement is more narrow than the standard information bottleneck , but encompasses many practical use cases . For example , in reinforcement learning , which is the primary subject of our experiments , the agent can be augmented with some privileged information in the form of a model based planner , or information which is the result of communication with another agent . This “ additional ” information can be seen as a privileged input because it requires the agent to do something extra to obtain it . Our work provides the following contributions . First , we propose a variational bandwidth bottleneck ( VBB ) that does not look at the privileged input before deciding whether to use it or not . At a high level , the network is trained first to examine the standard input , and then stochastically decide whether to access the privileged input or not . Second , we illustrate several applications of this approach to reinforcement learning , in order to construct agents that can stochastically determine when to evaluate costly model based computations , when to communicate with another agent , and when to access the memory . We experimentally show that the proposed model produces better generalization , as it learns when to use ( or not use ) the privileged input . For example , in the case of maze navigation , the agent learns to access information about the goal location only near natural bottlenecks , such as doorways . 2 PROBLEM FORMULATION . We aim to address the generalization issue described in the introduction for an important special case of the variational information bottleneck , which we refer to as the conditional bottleneck . The conditional bottleneck has two inputs , a standard input , and a privileged input , that are represented by random variables S and G , respectively . Hence , S , G , Y are three random variables with unknown distribution pdist ( S , G , Y ) . The information bottleneck provides us with a mechanism to determine the correct output while accessing the minimal possible amount of information about the privileged input G. In particular , we formulate a conditional variant of the information bottleneck to minimize the mutual information between the bottleneck layer and the privileged input I ( Z , G|S ) , given the standard input while avoiding unnecessary access to privileged input G. The proposed model consists of two networks ( see Fig . 1 ) : The encoder network that takes in the privileged input G as well as the standard input S and outputs a distribution over the latent variable z such that z ∼ p ( Z|G , S ) . The decoder network pdec ( Y|Z , S ) takes the standard input S and the compressed representation Z and outputs the distribution over the target variable Y . 3 VARIATIONAL BOTTLENECK ON STANDARD INPUT AND PRIVILEGED INPUT . The information bottleneck ( IB ) objective ( Tishby et al. , 2000 ) is formulated as the maximization of I ( Z ; Y ) − βI ( Z ; X ) , where X refers to the input signal , Y refers to the target signal , Z refers to the compressed representation of X , and β controls the trade-off between compression and prediction . The IB has its roots in channel coding , where a compression metric I ( Z ; X ) represents the capacity of the communication channel between Z and X . Assuming a prior distribution r ( Z ) over the random variable Z , constraining the channel capacity corresponds to limiting the information by which the posterior p ( Z|X ) is permitted to differ from the prior r ( Z ) . This difference can be measured using the Kullback-Leibler ( KL ) divergence , such that DKL ( p ( Z|X ) ‖r ( Z ) ) refers to the channel capacity . Now , we write the equations for the variational information bottleneck , where the bottleneck is learnt on both the standard input S as well as a privileged input G. The Data Processing Inequality ( DPI ) ( Cover & Thomas , 2006 ) for a Markov chain x → z → y ensures that I ( x ; z ) ≥ I ( x ; y ) . Hence for a bottleneck where the input is comprised of both the standard input as well as privileged input , we have I ( Z ; G|S ) ≥ I ( Y ; G|S ) . To obtain an upper bound on I ( Z ; G|S ) , we must first obtain an upper bound on I ( Z ; G|S = s ) , and then average over p ( s ) . We get the following result : We ask the reader to refer to the section on the conditional bottleneck in the supplementary material for the full derivation . I ( Z ; G|S ) ≤ ∑ s p ( s ) ∑ g p ( g ) DKL ( p ( Z|s , g ) ‖r ( Z ) ) ( 1 ) 4 THE VARIATIONAL BANDWIDTH BOTTLENECK . We now introduce our proposed method , the variational bandwidth bottleneck ( VBB ) . The goal of the variational bandwidth bottleneck is to avoid accessing the privileged input G if it is not required to make an informed decision about the output Y . This means that the decision about whether or not to access G must be made only on the basis of the standard input S. The standard input is used to determine a channel capacity , dcap , which controls how much information about G is available to compute Z . If dcap denotes the channel capacity , one way to satisfy this channel capacity is to access the input losslessly with probability dcap , and otherwise send no information about the input at all . In this communication strategy , we have p ( Z|S , G ) = δ ( fenc ( S , G ) ) if we choose to access the privileged input ( with probability dcap ) , where fenc ( S , G ) is a deterministic encoder , and δ denotes the Dirac delta function . The full posterior distribution p ( Z|S , G ) over the compressed representation can be written as a weighted mixture of ( a ) ( deterministically ) accessing the privileged input and standard input and ( b ) sampling from the prior ( when channel capacity is low ) , such that z is sampled using z ∼ dcap ∗ ( δ ( fenc ( S , G ) ) ) + ( 1− dcap ) ∗ r ( z ) . ( 2 ) This modified distribution p ( Z|S , G ) allows us to dynamically adjusts how much information about G is transmitted through Z . As shown in the Figure 1 , if dcap is set to zero , Z is simply sampled from the prior and contains no information about G. If it is set to one , the privileged information in G is deterministically transmitted . The amount of information about G that is transmitted is therefore determined by dcap , which will depend only on the standard input S. This means that the model must decide how much information about the privileged input is required before accessing it . Optimizing the information bottleneck objective with this type of bottleneck requires computing gradients through the term DKL ( p ( Z|S , G ) ‖r ( Z ) ) ( as in Eq . 1 ) , where z ∼ p ( Z|S , G ) is sampled as in Eq . 2 . The non-differentiable binary event , whose probability is represented by dcap , precludes us from differentiating through the channel capacity directly . In the next sections , we will first show that this mixture can be used within a variational approximation to the information bottleneck , and then describe a practical approximation that allows us to train the model with standard backpropagation . 4.1 TRACTABLE EVALUATION OF CHANNEL CAPACITY . In this section , we show how we can evaluate the channel capacity in a tractable way . We learn a deterministic function B ( S ) of the standard input S which determines channel capacity . This function outputs a scalar value for dcap ∈ ( 0,1 ) , which is treated as the probability of accessing the information about the privileged input . This deterministic function B ( S ) is parameterized as a neural network . We then access the privileged input with probability dcap = B ( S ) . Hence , the resulting distribution over Z is a weighted mixture of accessing the privileged input fenc ( S , G ) with probability dcap and sampling from the prior with probability 1− dcap . At inference time , using dcap , we sample from the Bernoulli distribution b ∼ Bernoulli ( dprob ) to decide whether to access the privileged input or not . 4.2 OPTIMIZATION OF THE KL OBJECTIVE . Here , we derive the KL divergence objective that allows for tractable optimization of DKL ( p ( Z|S , G ) ‖r ( Z ) ) ( as in Eq . 1 , 2 ) . Proposition 1 Given the standard input s , privileged input g , bottleneck variable z , and a deterministic encoder fenc ( s , g ) , we can express the DKL between the weighed mixture and the prior as DKL ( p ( z|s , g ) ‖r ( z ) ) = −dcap log dcap + ( 1− dcap ) [ log p ( f ( s , g ) ) − log ( dcap ∗ p ( f ( g , s ) ) + ( 1− dcap ) ) ] ( 3 ) The proof is given in the Appendix , section B . This equation is fully differentiable with respect to the parameters of f ( g , s ) and B ( s ) = dcap , making it feasible to use standard gradient-based optimizers . Summary : As in Eq . 2 , we approximate p ( Z|S , G ) as a weighted mixture of fenc ( S , G ) and the normal prior , such that z ∼ dcap ∗ ( fenc ( S , G ) ) + ( 1− dcap ) ∗ N ( 0,1 ) . Hence , the quantity DKL ( p ( Z|S , G ) ‖r ( Z ) ) can be seen as a bound on the information bottleneck objective . When we access the privileged input G , we pay a cost equal to I ( Z , G|S ) , which is bounded by DKL ( p ( Z|S , G ) ‖r ( Z ) ) as in Eq . 1 . Hence , optimizing this objective causes the model to avoid accessing the privileged input when it is not necessary .

Summary: This paper proposes a type of conditional Information Bottleneck (IB) that addresses the following problem: given that some features may be expensive to obtain for use in prediction, when should they be obtained such that the overall benefit outweighs the cost? A variant of the IB is proposed to model this question. However, optimization is intractable. The paper replaces a certain non-differentiable operation by a deterministic neural network which outputs the probability of seeking the expensive features. The main application here is reinforcement learning, where an agent could compute some plan or communicate with other agents at a cost, and the goal is to solve the task more efficiently while making use of this additional information. It is shown that the proposed method, VBB, makes judicious use of the limited number of costly feature acquisitions it makes, resulting in improved task performance across 3 tasks.