summaries: This paper proposes an extension to DQN, more generally applicable to value-based deep RL systems, that encodes the return using a thermometer encoding with exponentially-sized bins. This enables returns of vastly differing magnitudes to be learned without hurting performance. The authors propose an algorithm for learning these encode returns, including the use of a variance scaling term to speed up learning.

summaries: This work describes and addresses the issue of _reward progressivity_ in reinforcement learning, where as the task progresses the scale of the reward changes. The authors argue that reward progressivity harms Q-learning when training signals arising from large rewards interfere with those arising from smaller rewards. They propose a form of reward decomposition with an analogous modification to the Q-network output, which together help to ensure that training losses from small and large rewards are similarly scaled. The authors present a handful of experimental results demonstrating that their proposed method outperforms two other reward re-scaling baselines when reward progressivity is an issue and maintains good performance in more standard tasks.

summaries: In this paper the authors propose a new RL method, spectral DQN, in which rewards are decomposed into different frequencies. This decomposition allow for the training loss to better balanced on certain tasks - in particular those with progressive rewards. The new method is shown to perform well on specially constructed tasks with extreme reward progressively, as well as on a selection of standard Atari tasks.

abstractText: In this paper we consider reinforcement learning tasks with progressive rewards; that is, tasks where the rewards tend to increase in magnitude over time. We hypothesise that this property may be problematic for value-based deep reinforcement learning agents, particularly if the agent must first succeed in relatively unrewarding regions of the task in order to reach more rewarding regions. To address this issue, we propose Spectral DQN, which decomposes the reward into frequencies such that the high frequencies only activate when large rewards are found. This allows the training loss to be balanced so that it gives more even weighting across small and large reward regions. In two domains with extreme reward progressivity, where standard value-based methods struggle significantly, Spectral DQN is able to make much farther progress. Moreover, when evaluated on a set of six standard Atari games that do not overtly favour the approach, Spectral DQN remains more than competitive: While it underperforms one of the benchmarks in a single game, it comfortably surpasses the benchmarks in three games. These results demonstrate that the approach is not overfit to its target problem, and suggest that Spectral DQN may have advantages beyond addressing reward progressivity.

authors:
Adrià Puigdomènech Badia
Pablo Sprechmann
Alex Vitvitskyi
Daniel Guo
Bilal Piot
Steven Kapturowski
Olivier Tieleman
Martı́n Arjovsky
Alexander Pritzel
Andew Bolt

title: Never Give Up: Learning Directed Exploration Strategies

venue: In Proceedings of the 2020 International Conference on Learning Representations (ICLR),

authors:
Marc G. Bellemare
Will Dabney
Rémi Munos

title: A distributional perspective on reinforcement learning

venue: In Proceedings of the 34th International Conference on Machine Learning,

authors:
Will Dabney
Georg Ostrovski
David Silver
Remi Munos

title: Implicit Quantile Networks for Distributional Reinforcement Learning

venue: In Proceedings of the 35th International Conference on Machine Learning,

authors:
Will Dabney
Mark Rowland
Marc G Bellemare
Rémi Munos

title: Distributional Reinforcement Learning with Quantile Regression

venue: In Thirty-Second AAAI Conference on Artificial Intelligence,

authors:
William Fedus
Carles Gelada
Yoshua Bengio
Marc G Bellemare
Hugo Larochelle

title: Hyperbolic Discounting and Learning over Multiple Horizons

authors:
Max Jaderberg
Volodymyr Mnih
Wojciech Marian Czarnecki
Tom Schaul
Joel Z Leibo
David Silver
Koray Kavukcuoglu

title: Reinforcement Learning with Unsupervised Auxiliary Tasks

venue: arXiv preprint arXiv:1611.05397,

authors:
Steven Kapturowski
Georg Ostrovski
John Quan
Remi Munos
Will Dabney

title: Recurrent Experience Replay in Distributed Reinforcement Learning

venue: In Proceedings of the 2019 International Conference on Learning Representations (ICLR),

authors:
Clare Lyle
Marc G Bellemare
Pablo Samuel Castro

title: A Comparative Analysis of Expected and Distributional Reinforcement Learning

venue: In Proceedings of the AAAI Conference on Artificial Intelligence,

authors:
Volodymyr Mnih
Koray Kavukcuoglu
David Silver
Andrei A Rusu
Joel Veness
Marc G Bellemare
Alex Graves
Martin Riedmiller
Andreas K Fidjeland
Georg Ostrovski

title: Human-Level Control Through Deep Reinforcement Learning

authors:
Joshua Romoff
Peter Henderson
Ahmed Touati
Emma Brunskill
Joelle Pineau
Yann Ollivier

title: Separating value functions across time-scales

venue: Proceedings of Machine Learning Research,

authors:
Stuart J Russell
Andrew Zimdars

title: Q-Decomposition for Reinforcement Learning Agents

venue: In Proceedings of the 20th International Conference on Machine Learning

authors:
Nathan Sprague
Dana Ballard

title: Multiple-Goal Reinforcement Learning with Modular Sarsa(0)

authors:
Hado P van Hasselt
Arthur Guez
Matteo Hessel
Volodymyr Mnih
David Silver

title: Learning values across many orders of magnitude

venue: In Advances in Neural Information Processing Systems

authors:
Harm Van Seijen
Mehdi Fatemi
Joshua Romoff
Romain Laroche
Tavian Barnes
Jeffrey Tsang

title: Hybrid Reward Architecture for Reinforcement Learning

authors:
Harm Van Seijen
Mehdi Fatemi
Arash Tavakoli

title: Using a Logarithmic Mapping to Enable Lower Discount Factors in Reinforcement Learning

venue: Advances in Neural Information Processing Systems,

title: Learning from Delayed Rewards

venue: PhD thesis, King’s College,

heading: 1 INTRODUCTION

text:
In decision making tasks involving compounding returns, such as stock market investing, it is common for the rewards received by the agent to increase in magnitude over time. An investor that starts out with little capital will achieve relatively small profits or losses early on, but if they are successful, their capital – and hence their potential profits and losses – will gradually increase. This property also arises in many games settings. For example, in the Atari game Video Pinball, the player can increase a “bonus multiplier” that increases the score paid per bumper hit. Since the bonus multiplier does not reset until the player dies, the rewards typically increase with time. In this paper, we refer to any task exhibiting this phenomenon as a progressive reward task.
We hypothesise that reward progressivity may be problematic for value-based deep reinforcement learning agents. Our rationale is that the temporal difference errors that arise under such algorithms typically scale with the magnitude of the training targets. Accordingly, experience collected from states with large expected returns may dominate training, causing the agent’s performance to degrade in other states. While this is rational in one sense – generally, it is more important to make accurate decisions in high-stakes situations – it is potentially harmful on progressive reward tasks. This is because, on progressive reward tasks, it may be necessary for the agent to perform well in relatively unrewarding regions before it can reach more rewarding regions. For example, in stock market investing, an investor can only reach states where they have large capital if they first perform well with small capital. Note too that the rewards need not be strictly progressive for this problem to arise; all that is needed is for the rewards to be progressive over some extended period.
Algorithms that clip the reward, such as DQN (Mnih et al., 2015), are less susceptible to this problem, because they do not perceive increases in reward magnitude beyond the clipping point. However, it is straightforward to construct examples where reward clipping masks the optimal solution. For example, in Bowling, clipping the rewards to [-1, 1] makes bowling a strike appear no better than hitting a single pin (Pohlen et al., 2018). While subsequent approaches make learning feasible without reward clipping (van Hasselt et al., 2016; Pohlen et al., 2018), we show in Section 3 that
these methods are negatively impacted by strong reward progressivity. Motivated by this, we seek an algorithm that is capable of learning from unclipped rewards, while mitigating the negative impact that large rewards have on the agent’s ability to learn in less rewarding regions of the task.
To meet this aim, we propose Spectral DQN. Under this approach, rewards are decomposed into a multidimensional spectral reward, where each component is bound to [-1, 1]. The upper frequencies of the spectrum activate only when large magnitude rewards are received. For example, in Video Pinball, the lowest frequency activates on any score, while the upper frequencies only activate when the player has accumulated a large bonus multiplier. The bounded nature of the spectral rewards allows the agent to learn expected spectral returns without instability arising from large training targets. Meanwhile, the full, unclipped expected return can be recovered by summing across the return spectrum. (The full, technical definitions of these terms are provided in Section 4). In terms of addressing reward progressivity, a key advantage of this approach is that it allows us to balance the training loss and prevent the upper frequencies from receiving undue weight.
To test whether this approach helps mitigate the impact of reward progressivity, we perform two set of experiments: First, we apply Spectral DQN to two domains that we specifically designed to exhibit strong reward progressivity. While previous methods struggle on these tasks, failing to learn almost anything at all in the more extreme domain, Spectral DQN performs markedly better, making considerable progress on both tasks. Next, we apply our approach to a less constructed set of tasks; namely, 6 standard Atari games. Only some of these games exhibit reward progressivity, and none of them nearly so strongly as the extreme domains from earlier. While Spectral DQN does not outperform the benchmarks as noticeably in these domains (it is beaten by one of the benchmarks in a single game, but comfortably outperforms both benchmarks in three games), the results clearly show that Spectral DQN is not overfit to constructed domains. Moreover, since its outperformance in some of the games cannot be attributed solely to reward progressivity, it appears that the approach may offer additional advantages, as explained later in our analysis of the results.

heading: 2 PREVIOUS APPROACHES TO HANDLING REWARD VARIABILITY

text:
In the Atari domain, where the DQN algorithm was first evaluated, the agent receives a reward equal to the score increase at each frame. However, since score magnitudes vary greatly across games, DQN clips all rewards to [−1,+1]. This constrains the size of the training targets so that it is easier to find a step size that performs well across the entire suite of games. While this heuristic turns out to perform well in Atari, a clear downside of the approach is that the agent becomes unable to distinguish large rewards from small rewards.
The Pop-Art algorithm (van Hasselt et al., 2016) offers a more principled solution. It trains from the unclipped reward, while normalising the training targets to have zero mean and unit variance. To ensure that the network’s predictions are preserved when the normalisation parameters are updated, the output of the final layer is scaled and shifted by an offsetting amount.
Pohlen et al. (2018) propose an alternative approach to handling unclipped rewards that reduces the variance of the training targets by applying a squashing function, h. The agent learns squashed action-values, Q̃(s, a) = h(Q(s, a)), via the transformed Bellman backup:
Q̃(st, at)← Q̃(st, at) + α [ h ( r + γmax
a′ h−1(Q̃(st+1, a
′)) ) − Q̃(st, at) ] (1)
The authors prove that the transformed Bellman operator remains a contraction provided that both h and h−1 are Lipschitz continuous and γ < 1/LhLh−1 (where Lh and Lh−1 are the respective Lipschitz constants). In their experiments, they use the following squashing function:
h(x) = sign(x)( √ |x|+ 1− 1) + x (2)
where > 0 ensures that h−1 is Lipschitz continuous.
For the remainder of this paper, we refer to this approach as target compression. In the Atari domain, Pohlen et al. found target compression to perform significantly better than Pop-Art. It has since become the dominant method for handling unclipped rewards in Atari, with several recent largescale agents (Kapturowski et al., 2019; Badia et al., 2020a;b), including the state-of-the-art Agent57, favouring target compression over Pop-Art.

heading: 3 DEMONSTRATING THE PROBLEM OF REWARD PROGRESSIVITY

text:
While both Pop-Art and target compression address reward variance across tasks, making it easier to devise a single learning configuration that performs well across a task suite, they were not designed to address intra-task reward variability. As such, we posit that they are ill-suited to addressing reward progressivity, since it is an intra-task property. To test this claim, we introduce two domains that are designed to exhibit extreme reward progressivity, and thus exacerbate the issue: Ponglantis and Exponential Pong.
Ponglantis. In Ponglantis, the game starts out exactly the same as the Atari game Pong, where valued-based deep reinforcement learning methods are known to perform well (Mnih et al., 2015). The reward scale in Pong is small relative to most Atari games; the player receives a +1 reward for hitting the ball past the computer paddle, and a -1 reward if they concede a point. In Ponglantis, however, if the player manages to score 10 points, the task immediately switches to Atlantis, where the reward scale is roughly 1000 times larger. The Atlantis phase ends upon the loss of a single life.
Exponential Pong. The game Exponential Pong is identical to regular Pong, except that the rewards are multiplied by 2n, where n is the player’s current score.
We applied two agents to these games: Pop-Art and DQN+TC, where the latter is a version of DQN that we modified to use target compression and learn from the unclipped reward. The results of this experiment, averaged over 5 seeds, are shown in Figure 1. (For full experimental details, see the Appendix.) To simplify analysis of the Exponential Pong results, we have plotted the unexponentiated scores, i.e. the scores that the agents would have achieved under standard scoring in Pong. For context, the best possible score on standard Pong is +21, and the worst possible score is -21.
In Exponential Pong, which exhibits the more extreme reward progressivity (since the reward multiplier can potentially reach 220 ≈ 1 million), both agents essentially fail, providing strong support for our hypothesis that reward progressivity can be damaging, at least in extreme cases.
In the more moderate case of Ponglantis, Pop-Art fails to make meaningful progress beyond the low-scoring Pong phase of the task. The likely reason for this is that as soon as the agent starts reaching Atlantis and receiving large rewards, the target scaling mechanism employed by Pop-Art starts dividing the training targets by a large factor, heavily dampening the much smaller learning signals arising from the Pong phase of the task. On the other hand, DQN+TC is eventually able to make some progress into the Atlantis phase. The fact that DQN+TC outperforms Pop-Art is consistent with our hypothesis; while Pop-Art sees the reward scale jump by a factor of≈1000 upon reaching Atlantis, the training targets of DQN+TC jump by a smaller factor, owing to the square root squashing function (Equation 2). In Appendix B.2, we perform some additional experiments in Ponglantis to find where the “cutoff points” for Pop-Art and DQN+TC lie, i.e. how big the jump in reward magnitude has to be for the agents’ performances to start deteriorating.
Given the above discussion, a natural question to ponder is whether DQN+TC might achieve better performance in Exponential Pong by employing a logarithmic squashing function instead of a square root. Unfortunately though, as Pohlen et al. (2018) point out to a reviewer of their work (see https://openreview.net/forum?id=BkfPnoActQ), a logarithmic transform does not have a Lipschitz continuous inverse, and thus the transformed Bellman is no longer guaranteed to be a contraction.

heading: 4 SPECTRAL REINFORCEMENT LEARNING

text:
Based on the findings of the previous section, it seems that what we require is an algorithm where: (1) The training targets are constrained in magnitude, no matter how large the reward scale grows. (2) Regions of the task with small rewards yield meaningfully large temporal difference errors, even once much larger rewards have been discovered. Note that Pop-Art fails to meet the second criterion, while target compression fails to meet both, albeit it takes larger shifts in reward magnitude for the problems associated with reward progressivity to become apparent.
In this section, we introduce a reinforcement learning approach that meets both of the above criteria by leveraging the idea of spectral decomposition. This algorithm does not in and of itself address reward progressivity; in fact, we show that it behaves identically to the standard Q-learning algorithm (Watkins, 1989), given certain caveats. However, the properties of this algorithm facilitate adaption to progressive rewards, and in Section 5 we present a deep learning version of the approach that is capable of learning in domains with extreme reward progressivity, such as Exponential Pong.

heading: 4.1 SPECTRAL DECOMPOSITION

text:
Spectral decomposition is a broad term that arises in many different fields, but in this paper it refers to a specific method for decomposing real numbers into a sum of exponentially weighted components.
The intuition behind our method is illustrated in Figure 2. In this example, the number being spectrally decomposed is 6.5. The method begins by dividing the number line into buckets of width bi, where b > 0 is the base of the decomposition and i ∈ Z≥0. Since the base here is 2, bucket 0 has width 1 and spans from 0 to 1, bucket 1 has width 2 and spans from 1 to 3, bucket 2 has width 4 and spans from 3 to 7, etc. The idea is that each bucket corresponds to a different frequency in the spectrum. The spectral decomposition is found by calculating the proportion of each bucket that is filled. Since the number 6.5 completely fills the first two buckets, then 87.5% of the next bucket, then none of the remaining buckets, the spectral decomposition of 6.5 to base 2 is (1, 1, 0.875, 0, 0, . . . ).
More generally, let r be a non-negative real number, and let ri+b denote the ith element of r’s spectral decomposition to base b. The formula for calculating each element is:
ri+b = clamp ( (r − b i−1 b−1 )/b i, 0, 1 )
(3)
(See the Appendix for the derivation of this formula.) The value r can be recovered from the spectral decomposition by multiplying the bucket widths by the bucket fill proportions and summing:
r = ∑∞ i=0 b iri+b (4)
This method can be extended to negative numbers by negating the sign on each element of the representation. Given any real number r, the formula for calculating each element becomes:
rib = sign(r)× clamp ( (|r| − b i−1 b−1 )/b i, 0, 1 )
(5)

heading: 4.2 DECOMPOSING THE RETURN

text:
The core idea behind our reinforcement learning approach is to apply spectral decomposition to the reward stream, and thus learn a spectral decomposition of the expected return. Before delving into the technical details, see Table 1 for a worked example that highlights the intuition. The discount
factor is γ = 0.99, and the return over the first five time steps is 2.254. The usual way of calculating the return is, of course, to multiply the rewards by γt then sum. The alternative, spectral approach is to first compute the spectral decomposition of the rewards according to Equation 5, then calculate the returns arising from each reward frequency. We refer to the decomposed rewards as spectral rewards, and to the decomposed returns as spectral returns. The overall return of 2.254 can be found by multiplying the spectral returns by their corresponding weights of bi then summing.
To appreciate why we go to this extra effort in calculating the return, note that the spectral rewards are all bound to [-1, 1]. This in turn constrains the magnitude of the spectral returns to be no more than 1 / (1 - γ), which, again, is useful from a function approximation standpoint. Moreover, note that the relative sparsity of large rewards is made more apparent by this approach: As we progress from r02 to r 3 2 , the magnitudes of the spectral reward components decrease, with only the two largest magnitude rewards triggering a response in r32 .

heading: 4.3 SPECTRAL Q-LEARNING

text:
To translate the above intuition into a concrete algorithm, we begin by defining the spectral actionvalue function as:
Qπ(s, a, i) = Eπ[ ∑∞ k=t γ k−trib,t | at = a, st = s] (6)
where rib,t denotes the ith element of the spectral reward at time t. As shown in the Appendix, it is straightforward to prove via the linearity of expectation that the standard action-value function can be recovered as the weighted sum:
Qπ(s, a) = ∑∞ i=0 b iQπ(s, a, i) (7)
This leads to a natural extension of the Q-learning algorithm (Watkins, 1989), whereby the greedy action is found by maximising this expression with respect to the available actions, and Qπ(s, a, i) is trained via the standard Bellman backup, except that the reward is replaced by the ith element of the spectral reward. For clarity, the full pseudocode for this algorithm, which we term Spectral Q-learning, is provided in the Appendix (see Algorithm 1).
An important practical limitation of this algorithm is that it is impossible to train over full, infinite spectrum; that is, one must choose a finite upper bound for i. (This would not be an issue if we knew the maximum reward magnitude in advance, but for generality we do not wish to assume this.) Regardless, for base b, observe that the first N + 1 frequencies capture all rewards up to magnitude (bN+1 − 1)/(b− 1), since the bucket widths form a geometric series. Given the exponential nature of this term, it is possible to capture large magnitude rewards with a relatively small value of N . For example, with b = 2 and N = 20, the algorithm fully captures all rewards of magnitude up to 1 million, which is more than sufficient for Atari. Taking this logic one step further, we arrive at the following proposition, which is proved in the Appendix:
Proposition 1 In the tabular setting, Spectral Q-learning behaves identically to standard Qlearning, provided that the rewards are bound in magnitude to no more than (bN+1 − 1)/(b− 1).
Now that we have presented the full algorithm, it is worth reflecting on why we use a spectral representation for the rewards instead of, say, a binary encoding. Put simply, very small differences in reward should not matter much to a rational agent. For example, it should not matter much whether a particular reward’s magnitude is 63 or 64. However, under a binary encoding, 63 is represented as 0111111, while 64 is represented as 1000000. Since a reward of 63 activates more frequencies than a reward of 64, it would thus trigger a larger TD error at the start of training. The thermometer-style encoding introduced earlier addresses this problem.

heading: 5 SPECTRAL DQN

text:
In order to create a deep learning version of the Spectral Q-learning algorithm, we modify the DQN algorithm in an analogous way to which we modified the Q-learning algorithm in the previous section. We name the resultant algorithm Spectral DQN.
To parameterise all N + 1 frequencies of the spectral action-value function via a single network, we multiply the number of outputs on the standard DQN network by the same count. One crucial detail here is that we initialise the weights and biases of the final layer to zero. This allows us to set N to a conservatively large value, since it ensures that superfluous heads have no effect on training. Moreover, it ensures that the action-value calculations do not blow out due to the exponential weights in Equation 7 being applied to random noise.
For the training loss, we use a weighted sum of the temporal difference errors for each frequency: L(θ) = 12 ∑N i=0 wi [ rib + γmaxa′ Q̃(s ′, a′, i; θ−)− Q̃(s, a, i; θ) ]2 (8)
where wi is the error weighting given to the ith frequency, s is a state sampled uniformly from the replay memory, a is the sampled action, s′ is the next state, r is the reward (with spectral components rib for i ∈ {0, . . . , N}), θ are the parameters of the network being trained, and θ− are the parameters of the periodically refreshed target network.
An important detail here is the approach we take to weighting the loss for each frequency. Given that the ith frequency is weighted by bi when calculating the overall action-values, it is tempting to set wi = bi. However, this goes completely against one of the aims of our approach, which is to prevent the large reward frequencies from dominating the loss. In fact, we show in the next section that weighting by bi has a disastrous impact in the two extreme domains introduced earlier.
In light of the above, rather than weighting by bi, we strive to give all frequencies equal weighting. However, we do not simply set wi = 1, because doing so would effectively give too much emphasis to the small reward frequencies. As pointed out by van Hasselt et al. (2016), the gradients propagated to the Q-network’s hidden layers depend quadratically on the scale of the targets (see Section 2.3 of their work). Since the lower frequencies have denser rewards, and hence larger training targets, they will dominate the gradients at the hidden layers if this effect is not accounted for. To address this, we set wi = 1/σ2i , where σi is the standard deviation of the training targets for frequency i, calculated as a running average as in the Pop-Art algorithm. Finally, to ensure that this correction only affects the hidden layers, we unscale the gradients at the final layer to train as if wi = 1, as in standard DQN. (Implementing this is not especially difficult; the precise details can be found in our source code.) In Appendix B.1, we provide an ablation where we do not correct the hidden layer gradients, and we find that this has a significant adverse effect on performance.

heading: 6 RETURNING TO THE EXTREME DOMAINS

text:
To test how well Spectral DQN meets our aims, we first applied it to the two extreme domains from Section 3 – Ponglantis and Exponential Pong – where reward progressivity is clearly a key issue. In addition, to probe the importance of the frequency loss weighting, we included a variant of Spectral DQN with exponential loss weighting, i.e. wi = bi, which we argued against in the previous section. The results of this experiment, including the baseline results from Section 3, are shown in Figure 3, with each curve averaged over 5 seeds.
There are two clear findings here: (1) Spectral DQN is much better at handling extreme reward progressivity than the baselines. It scores much higher in Ponglantis, and makes far more meaningful progress in Exponential Pong, eventually learning to beat the computer player comfortably on
average. (2) The loss weighting is of crucial importance, with the agent using exponential weights performing worse than even DQN+TC. The latter finding matches our expectations; as we noted in Section 3, target compression does at least somewhat mitigate the overemphasis given to large reward regions of the task, while Spectral DQN with wi = bi gives heavy emphasis to such regions.
In Figure 4, we examine how the agents’ training errors evolve as they play Exponential Pong. The errors are broken down by the player’s current score ∈ [0, 20], with bluer lines corresponding to states with a low player score and redder lines corresponding to states with high player scores. To make the graphs comparable, we plot percentage TD errors, i.e.:
mean(|Qpredicted −Qtarget|) / mean(|Qtarget|) (9) where the means are running averages, calculated from states sampled from the replay buffer. In all cases we report the true TD error, i.e. for DQN+TC we report the error on the uncompressed action-values, and for Spectral DQN we report the error on the full, summed action-values.
The instability of Pop-Art is quite apparent here: After a very brief period of stable learning, the TD errors start spiking as soon as the agent reaches even middling scores. Consistent with our earlier analysis, DQN+TC appears a lot more stable. For states with low player scores, the TD percentage errors hover at around 5%, although the agent never really learns accurate action-values for states where the player score exceeds ≈7 points. At first this might seem strange – since the reward scale increases exponentially with the player’s score, states with a high player score ought to be effectively prioritised by the loss function. However, there is an additional effect to consider here: As soon the agent’s predictions start to deteriorate for the lower-score states, its policy for these states likewise deteriorates and it starts struggling to reach the higher-score states. As such, the agent seems to get stuck in a middle ground, where it does not encounter the higher-score states often enough to learn accurate action-values there, nor for the estimates in the lower-score states to truly deteriorate. It appears that a TD percentage error of around 5% is a crucial cutoff line, where the agent is right on the cusp of being able to play well, but playing better, conversely, leads to deterioration.
In contrast to both Pop-Art and DQN+TC, when Spectral DQN reaches a new high score, it quickly learns accurate action-value estimates for the associated states, as indicated by the sharply dropping red lines in Figure 4. Moreover, the TD errors for the lower-score states remain stable as the agent improves, almost certainly because Spectral DQN’s loss does not give so much weight to large magnitude rewards. Finally, note that the agent’s sudden rise in performance at around 2.5 × 107 frames in Figure 3 seems to coincide with the errors reducing below the 5% line, adding further credence to the analysis above.

heading: 7 EXPERIMENTS IN STANDARD ATARI GAMES

text:
To validate that Spectral DQN’s success in extreme domains does not come at the expense of performance in less constructed tasks, we also applied it to 6 standard Atari games. Moreover, we included 3 games where reward clipping is known to be detrimental, as per van Hasselt et al. (2016): Bowling, Centipede and Ms. Pacman. This was done to validate that the agent’s earlier success was not attributable to it somehow performing an implicit form of reward clipping. The results of these experiments are shown in Figure 5, with all curves averaged over 5 seeds.
While Spectral DQN’s performance relative to the baselines is not as exceptional here as in the extreme domains, it remains the strongest overall. Though the agent underperforms DQN+TC in Centipede, it noticeably outperforms both baselines in Bowling, Chopper Command and Video Pinball. On the other games, it performs at least on par with the benchmarks. Moreover, the fact that Spectral DQN’s performance compares favourably to Pop-Art in Bowling, Centipede and Ms. Pacman shows that it is truly responding to the unclipped reward.
Since Video Pinball exhibits reward progressivity, as explained in the introduction, Spectral DQN’s strong performance there fits well with our hypothesis. On the other hand, while Ms. Pacman also has progressive rewards — the player receives an increasing bonus for eating multiple ghosts with a single power pellet — these rewards come so close together that they are within the agent’s temporal horizon at the start of an episode. That is, there is not a small-expected-return phase preceding a large-expected-return phase, which we hypothesised to be a central component of the issue. (We return to this point and discuss further in the next section.)
Besides addressing reward progressivity, another possible advantage of Spectral DQN in high scoring games is simply that its training targets are more constrained. In games such as Video Pinball and Chopper Command, the expected returns can reach into the thousands, so even with target compression this may be problematic for DQN+TC. On the other hand, the only game where DQN+TC clearly outperforms Spectral DQN is Centipede, which is also high scoring. One possible explanation for this is that the reward structure in Centipede “confuses” Spectral DQN: Since large rewards trigger both low and high frequencies, it is likely that Spectral DQN learns a positive association between the various components of the return. However, in Centipede, losing a life triggers a stream of consecutive small rewards. Therefore, associating small and large rewards may cause Spectral DQN to predict an inflated return from dying.
There has been some conjecture that an advantage of combining distributional RL (Bellemare et al., 2017; Lyle et al., 2019) with function approximation is that it forces the agent to learn a richer representation. The same argument may also apply to our approach. In particular, in some tasks, small rewards are of more informational value than their face value represents. For example, in Video
Pinball, where rewards range into the thousands, the play field includes “spinners” that pay only one point per hit. Despite their negligible face value, hitting the spinners can dramatically change the velocity of the ball, so giving the lower frequencies more weight in the training loss may promote a stronger representation. Similarly, in Bowling, certain frequencies only activate after bowling strikes and spares. Since the game does not have a strongly progressive reward structure, improved representation learning is a more plausible explanation of Spectral DQN’s outperformance.

heading: 8 DISCUSSION AND CONCLUSION

text:
In this paper, we brought attention to the idea of reward progressivity, a task property that poses a potential issue for value-based reinforcement learning agents. Specifically, an agent’s performance in relatively unrewarding regions of a task may degrade once it discovers larger rewards. To address this issue, we proposed Spectral DQN, which decomposes the expected return magnitudinally, allowing the losses arising across different regions of the task to be balanced. Our experiments showed this approach to be particularly effective in domains exhibiting extreme reward progressivity, while still yielding more than competitive performance in less constructed domains.
Our discussion of the Ms. Pacman results from the previous section raised an interesting point: There is clearly a link between reward progressivity and discounting. If it were somehow possible to train reliably with γ = 1, a Ponglantis agent would be able to see that reaching the large, distant rewards in the Atlantis phase is contingent on first performing well in Pong. However, despite some recent progress, DQN-style agents typically destabilise with discounts very close to 1. Agent57 (Badia et al., 2020a) trains multiple value functions across a range of discounts, ranging as high as γ = 0.9999. However, the authors note that γ = 0.9999 generally yields very unstable learning, and they rely on a meta-controller to select smaller discounts in games where learning destabilises. As such, it does not offer a general solution to the problem. The idea of setting γ = 1 also raises the issue of reward sparsity: If the agent is cognizant of large, distant rewards, then smaller, more immediate rewards may appear so negligible that the task effectively becomes a sparse reward problem.
We noted in Section 3 that another idea for addressing reward progressivity is to apply target compression, but with logarithmic squashing instead of a square root. However, as Pohlen et al. (2018) note, the inverse of a logarithmic mapping is not Lipschitz continuous, and thus the transformed Bellman is no longer guaranteed to be a contraction. On the other hand, one way to get around poor convergence in practice is to use a shaper discount factor. Van Seijen et al. (2019) show that it is possible to achieve performance comparable to DQN via a logarithmic squash, provided the discount is small (γ = 0.96). However, they clip the rewards, which greatly constrains the gradient of the inverse mapping h−1, and thus limits the degree to which their implementation transgresses the theory. Moreover, the performance of their method degrades significantly with γ = 0.99, consistent with convergence issues (see Appendix C.2 in their work). Unfortunately, it is simple to provide examples where γ = 0.96 will lead to poor performance. For example, in the game Bowling, there is a delay of about 100 frames between the ball being released and the subsequent reward. However, if the player bowls a spare, the reward is delayed by a further 100 frames. Thus, under a discount of 0.96, it is better to consistently bowl 9s than to bowl spares.
Our work fits into a broader line of research whereby reinforcement learning agents seek to predict quantities other than just the mean expected return (Jaderberg et al., 2016; Bellemare et al., 2017; Fedus et al., 2019). Of this related work, one of the most similar approaches to ours is that of Romoff et al. (2019), who learn the value function in a modular way by decomposing based on the discount factor. Distributional reinforcement learning (Bellemare et al., 2017; Dabney et al., 2018b;a) is another closely related approach. Finally, an older line of work in this vein is decomposed RL (Russell & Zimdars, 2003; Sprague & Ballard, 2003; Van Seijen et al., 2017), whereby an agent is responsible for predicting only a portion of the action-value function. While some of these works bear many surface-level similarities to our own — the distributional C51 agent (Bellemare et al., 2017) in particular bears quite a similar network architecture to ours — on a deeper level, they are not addressing the same problem. Leveraging a probabilistic loss, as in C51, could foreseeably help to address the large training errors that arise in tasks like Exponential Pong. However, the method assumes that returns are constrained to [-10, 10], and it is unclear how this could be extended to unclipped rewards without requiring an intractable number of atoms in the distribution, or greatly reducing the precision of the atoms.

heading: A EXPERIMENTAL SETUP

text:
The source code for our experiments is available at https://github.com/mchldann/ SpectralDQN.
We injected stochasticity into the Atari environment via the no-ops regime used in the Nature DQN paper (Mnih et al., 2015). Rather than conducting periodic evaluations of the agents with a small exploration constant, we simply report their training performance, since this arguably provides a more reliable measure of progress. All curves in the paper were averaged over 5 seeds. Shaded regions represent bootstrapped single standard deviation confidence intervals.
Hyperparameters specific to Spectral DQN are listed in Table 2, while hyperparameters common to all agents are listed in Table 3. Spectral DQN used an identical network architecture to Nature DQN Mnih et al. (2015), except that we multiplied the number of outputs by N + 1.

heading: B ADDITIONAL EXPERIMENTS

text:
In this section, we provide the details of some additional experiments that had to be omitted from the main text due to space constraints.
B.1 ABLATION: NO ERROR WEIGHTING
Recall from Section 5 that we weight the frequencies’ individual losses by wi = 1/σ2i , where σi is the standard deviation of the training targets for frequency i. In Figure 6, we show the effect of ablating this adjustment and simply settingwi = 1. The results reveal that the adjustment is a crucial component, with the ablation proving disastrous in 4 of the 6 games. (Note: It unsurprising that the Fishing results are unaffected, because only one reward frequency is active in that game.)
B.2 EASIER PONGLANTIS
In Section 3, we applied Pop-Art and DQN+TC to the game Ponglantis. The results suggested that DQN+TC was somewhat more robust to reward progressivity. Here, we provide some further experiments to determine just how strong the reward progressivity has to be for these agents to start deteriorating.
To hone in on the agents’ breaking points, we created two modified versions of the game: Ponglantis Easier and Ponglantis Even Easier. In Ponglantis Easier, the agent’s perceived rewards in the Atlantis phase of the task are scaled down by a factor 10. In Ponglantis Even Easier, they are scaled down by a factor of 100. The agents’ results in these games are shown in Figure 7. For comparability,
the plots show the true episodic return, rather than the perceived episodic return. We also include results for Spectral DQN as a reference point.
Relative to its performance in the original Ponglantis task, DQN+TC performs markedly better in both easier tasks. Since the rewards jump by a factor of around 1000 in the original task, this suggests that DQN+TC can tolerate a jump in reward magnitude of around 100x. Further, this implies that DQN+TC’s underperformance relative to Spectral DQN in the main Atari experiments is attributable to more than just reward progressivity, as none of the standard Atari games exhibit this level of jump. Our alternate hypotheses about representation learning and the relative ease with which the spectral architecture can represent large expected returns thus appear more likely.
Pop-Art, on the other hand, still performs poorly when the perceived rewards in Atlantis are scaled down by a factor of 10, only learning reliably when they are scaled down 100 times. In this case, it is plausible that the algorithm struggles with reward progressivity in standard Atari games.
B.3 REVERSE PONGLANTIS
To test whether it is truly reward progressivity that is the key issue in Ponglantis, as opposed to just reward variability, we created one further task, Reverse Ponglantis, where the Pong and Atlantis phases are reversed. The player progresses to the Pong phase upon losing a life in Atlantis.
The agents’ results in this game are plotted in Figure 8. While all agents improve steadily, which confirms that progressivity is the main issue for Pop-Art, it is notable DQN+TC’s performance is actually worse here than in Ponglantis Even Easier. This is despite the fact that the agents now start in the high-scoring Atlantis phase, so there is no need to play Pong well. From this result, it is clear that a key issue for DQN+TC is simply the reward scale in Atlantis. This adds further weight to our
hypothesis from Section 7 that one of the main advantages of Spectral DQN in standard Atari games is the constrained nature of its training targets.

heading: C SPECTRAL Q-LEARNING ALGORITHM

text:
Pseudocode for the Spectral Q-Learning algorithm is provided in Algorithm 1.
Algorithm 1 Spectral Q-Learning 1: var: the base of the spectral decomposition, b 2: var: the number of spectral components trained, N 3: var: the learning rate, α 4: 5: Initialise Q(s, a, i) 6: for each episode do 7: Sample s from start distribution 8: while s is not terminal do 9: Choose a according to -greedy policy derived from Q(s, a) = ∑N i=0 b
iQ(s, a, i) 10: Take action a, observe r, s′
11: a′∗ = argmaxa′ ∑N i=0 b
iQ(s′, a′, i) 12: for i ∈ {0, . . . , N} do 13: Q(s, a, i)← Q(s, a, i) + α[rib + γQ(s′, a′∗, i)−Q(s, a, i)] 14: s← s′

heading: D DERIVATION OF THE SPECTRAL DECOMPOSITION FORMULA

text:
Recall from Section 4.1 that the spectral representation of a number is found by laying buckets of exponentially increasing width on the number line, then calculating the proportion of each bucket that is filled.
To calculate the proportion of the ith bucket that is filled by a non-negative number r, note that the sum of all the preceding buckets’ widths is equal to:∑i−1
k=0 b k = b i−1 b−1 (10)
Therefore, the amount leftover to fill the ith bucket is:
r − b i−1 b−1 (11)
To calculate the proportion of the ith bucket that is filled, we must divide this quantity by the width of the ith bucket (which is equal to bi) then clamp the result to the allowable range of values:
Filled proportion = clamp ( (r − b
i−1 b−1 )/b
i, 0, 1 )
(12)
Finally, to convert this into an odd function so that the representation of −r is equal to the negated representation of r, we replace r with |r| and multiply the result by sign(r). This yields the formula for the ith element of the spectral representation:
rib = sign(r)× clamp ( (|r| − b i−1 b−1 )/b i, 0, 1 )
(13)

heading: E PROOFS

text:
E.1 PROOF OF EQUATION 7
Starting with the standard definition of the action-value function, then replacing the reward at each time step with its spectral decomposition, we have:
Qπ(s, a) = Eπ[ ∑∞ k=t γ k−trk | at = a, st = s] (14)
= Eπ[ ∑∞ k=t γ k−t∑∞ i=0 b irib,k | at = a, st = s] (15)
Since there is no interaction between the summation indices, we can rearrange (15) into: Qπ(s, a) = Eπ[ ∑∞ i=0 b i ∑∞ k=t γ k−trib,k | at = a, st = s] (16)
Finally, using the linearity of expectation, we can take the expectation operator inside and write: Qπ(s, a) = ∑∞ i=0 b iEπ[ ∑∞ k=t γ k−trib,k | at = a, st = s] (17)
= ∑∞ i=0 b iQπ(s, a, i) (18)
E.2 PROOF OF PROPOSITION 1
Proposition 1 In the tabular setting, Spectral Q-learning behaves identically to standard Qlearning, provided that the rewards are bound in magnitude to no more than (bN+1 − 1)/(b− 1).
We prove this statement via induction. The first point to observe is that, provided the action-value estimates of the two algorithms agree at a given point in time, their -greedy policies at that time will also be equivalent. Therefore, assuming that the action-value functions agree at t = 0 (e.g. they are both initialised to zero), all we need to prove is the inductive step, i.e. that agreement in action-values at time t implies agreement in action-values at time t+ 1.
To distinguish between the action-value estimates of the two algorithms at different time steps, we use the following notation:
• Let Q̂(s, a; t) denote the action-value estimates under standard Q-learning after t updates. • LetQ(s, a; t) = ∑N i=0 b
iQ(s, a, i; t) denote the action-values estimates under Spectral Q-learning after t updates.
For the induction step, we need to show that Q(s, a; t + 1) = Q̂(s, a; t + 1) ∀s ∈ S, a ∈ A given that Q(s, a; t) = Q̂(s, a; t) ∀s ∈ S, a ∈ A. To begin, note that the standard Q-learning backup expressed under the above notation is:
Q̂(s, a; t+ 1)← Q̂(s, a; t) + α [ r + γQ̂(s′, a′∗; t)− Q̂(s, a; t) ] (19)
where a′∗ = argmaxa′ Q̂(s ′, a′∗; t) is the greedy action for the next state.
Since we assume equal action-values at time t, Spectral Q-learning will calculate the same greedy action for s′. Hence, the backup for each frequency of the spectral action-value function is:
Q(s, a, i; t+ 1)← Q(s, a, i; t) + α [ rib + γQ(s ′, a′∗, i; t)−Q(s, a, i; t) ]
(20)
After performing this update on each frequency, the function at time t+ 1 becomes:
Q(s, a; t+ 1) = ∑N i=0 b i [ Q(s, a, i; t) + α [ rib + γQ(s ′, a′∗, i; t)−Q(s, a, i; t) ]]
(21) = ∑N i=0 b iQ(s, a, i; t) + ∑N i=0 αb i [ rib + γQ(s ′, a′∗, i; t)−Q(s, a, i; t) ] (22)
= Q(s, a; t) + ∑N i=0 αb i [ rib + γQ(s ′, a′∗, i; t)−Q(s, a, i; t) ]
(23)
Since the induction base case assumes that Q̂(s, a; t) = Q(s, a; t) ∀s ∈ S, a ∈ A, we can rewrite this as:
Q(s, a, i; t+ 1) = Q̂(s, a; t) + α ∑N i=0 b i [ rib + γQ(s ′, a′∗, i; t)−Q(s, a, i; t) ]
(24) = Q̂(s, a; t) + α [∑N
i=0 b irib + γ ∑N i=0 b iQ(s′, a′∗, i; t)− ∑N i=0 b iQ(s, a, i; t) ] (25)
= Q̂(s, a; t) + α [∑N
i=0 b irib + γQ(s
′, a′∗; t)−Q(s, a; t) ]
(26) = Q̂(s, a; t) + α [∑N
i=0 b irib + γQ̂(s
′, a′∗; t)− Q̂(s, a; t) ]
(27)
Now, all that remains to establish equality with the standard Q-learning update in (19) is to show that ∑N i=0 b
irib = r. That is, we need to show that the partial sum over indices 0 to N fully captures the reward, or, equivalently, that rib = 0 for all i >= N + 1.
From the spectral decomposition formula (13) and the proposition’s statement that the rewards are bound to no more than (bN+1 − 1)/(b− 1), we have:
|rib| = clamp ( (|r| − b i−1 b−1 )/b i, 0, 1 )
(28) ≤ clamp ( ( b
N+1−1 b−1 − bi−1 b−1 )/b
i, 0, 1 )
(29) = clamp ( ( b
i(bN+1−i−1) b−1 )/b
i, 0, 1 )
(30)
Now, if i >= N + 1, it implies that (bN+1−i − 1) ≤ 0, so the whole expression in (30) becomes 0. Thus ∑N i=0 b irib = ∑∞ i=0 b irib = r and the proof is complete.

title: SPECTRAL REINFORCEMENT LEARNING


summaries: Broadly, this work is an attempt to understand how neural networks can form generalizable representations while being severely overparameterized. This work proposes an information theoretic measure, called the "usable information", and use it to quantify the amount of relevant information in different layers of a neural network during training. The key idea is that, in order for the information represented in one layer to be "usable" by the next layer, it should be decodable by a simple transformation (affine + element-wise nonlinearity).

summaries: The paper studies how initialization and the implicit regularization of SGD affect the training dynamics of neural networks in terms of minimality and sufficiency of learned representations. The main findings are that 1) SGD with random initialization learns almost minimal and sufficient representations and 2) SGD with an initialization that contains information about irrelevant factors fails to converge to minimal representations, increasing the chance of overfitting. These findings are interesting, useful for understanding neural networks, relevant to the ICLR community, but lack evidence of generality.

summaries: The authors contribute to the recent research on whether neural network training (in particular, SGD) favors minimal representations, in which irrelevant information is not represented by deeper layers. They do so by implementing a simple neuroscience-inspired task, in which the network is asked to make a decision by combining color and target information. Importantly, the network's output is conditionally independent of the color information, given the direction decision, so the color information is in some sense irrelevant at the later stages. Using this, the authors quantify the 'relevant' and 'irrelevant' information in different layers of the neural network during training. Interestingly, the authors show that minimal representation are uncovered only if the network is started with random initial weights. Information is quantified using a simple decoder network.

abstractText: We introduce a notion of usable information contained in the representation learned by a deep network, and use it to study how optimal representations for the task emerge during training. We show that the implicit regularization coming from training with Stochastic Gradient Descent with a high learning-rate and small batch size plays an important role in learning minimal sufficient representations for the task. In the process of arriving at a minimal sufficient representation, we find that the content of the representation changes dynamically during training. In particular, we find that semantically meaningful but ultimately irrelevant information is encoded in the early transient dynamics of training, before being later discarded. In addition, we evaluate how perturbing the initial part of training impacts the learning dynamics and the resulting representations. We show these effects on both perceptual decision-making tasks inspired by neuroscience literature, as well as on standard image classification tasks.

authors:
Alessandro Achille
Stefano Soatto

title: Emergence of invariance and disentanglement in deep representations

venue: Journal of Machine Learning Research,

authors:
Alessandro Achille
Matteo Rovere
Stefano Soatto

title: Critical learning periods in deep networks

venue: In International Conference on Learning Representations,

authors:
David Barber
Felix Agakov

title: The im algorithm: A variational approach to information maximization

venue: In Proceedings of the 16th International Conference on Neural Information Processing Systems,

authors:
Chandramouli Chandrasekaran
Diogo Peixoto
William T. Newsome
Krishna V. Shenoy

title: Laminar differences in decision-related neural activity in dorsal premotor cortex

venue: Nature Communications,

authors:
Yann Dubois
Douwe Kiela
David J. Schwab
Ramakrishna Vedantam

title: Learning optimal representations with the decodable information

authors:
Aditya Sharad Golatkar
Alessandro Achille
Stefano Soatto

title: Time matters in regularizing deep networks: Weight decay and data augmentation affect early learning dynamics, matter little near convergence

authors:
Ziv Goldfeld
Ewout van den Berg
Kristjan H. Greenewald
Igor Melnyk
Nam Nguyen
Brian Kingsbury
Yury Polyanskiy

title: Estimating information flow in neural networks

authors:
K. He
X. Zhang
S. Ren
J. Sun

title: Deep residual learning for image recognition

venue: In 2016 IEEE Conference on Computer Vision and Pattern Recognition (CVPR),

authors:
Michael Kleinman
Chandramouli Chandrasekaran
Jonathan C. Kao

title: Recurrent neural network models of multi-area computation underlying decision-making. bioRxiv, 2019

authors:
Alexander Kraskov
Harald Stögbauer
Peter Grassberger

title: Estimating mutual information

venue: Physical Review E,

title: Estimation of entropy and mutual information

venue: Neural Comput.,

authors:
Ben Poole
Sherjil Ozair
Aaron Van Den Oord
Alex Alemi
George Tucker

title: On variational bounds of mutual information

venue: Proceedings of the 36th International Conference on Machine Learning,

authors:
Andrew Michael Saxe
Yamini Bansal
Joel Dapello
Madhu Advani
Artemy Kolchinsky
Brendan Daniel Tracey
David Daniel Cox

title: On the information bottleneck theory of deep learning

authors:
Ravid Shwartz-Ziv
Naftali Tishby

title: Opening the black box of deep neural networks via information

authors:
Jost Tobias Springenberg
Alexey Dosovitskiy
Thomas Brox
Martin Riedmiller

title: Striving for simplicity: The all convolutional net

authors:
Yilun Xu
Shengjia Zhao
Jiaming Song
Russell Stewart
Stefano Ermon

title: A theory of usable information under computational constraints

venue: In 8th International Conference on Learning Representations,

heading: 1 INTRODUCTION

text:
An important open question for the theory of deep learning is why highly over-parametrized neural networks learn solutions that generalize well even though the model can in principle memorize the entire training set. Some have speculated that neural networks learn minimal but sufficient representations of the input through implicit regularization of Stochastic Gradient Descent (SGD) (Shwartz-Ziv & Tishby, 2017; Achille & Soatto, 2018), and that the minimality of the representations relates to generalizability. Follow-up work has disputed the validity of some of these claims when using deterministic deep networks (Saxe et al., 2018), leading to an ongoing debate on the notion of optimality of representations and how they are learned during training.
Part of the disagreement stems from the use of information-theoretic quantities: most previous studies in deep learning have analyzed the amount of information that the learned representation contains about the inputs using Shannon’s mutual information. However, when the mapping from input to representation is deterministic, the mutual information between the representation and input is degenerate (Saxe et al., 2018; Goldfeld et al., 2018). Rather than study the mutual information in a neural network, here we instead define and study the “usable information” in the network, which measures the amount of information that can be extracted from the representation by a learned decoder, and is scalable to high dimensional realistic tasks. We use this notion to quantify how relevant and irrelevant information is represented across layers of the network throughout the training process, and how this is affected by the optimization algorithms and the network pretraining.
In particular, we propose to study a simple task inspired by decision-making tasks in neuroscience, where inputs and outputs are carefully designed to probe specific information processing phenomena. We then extend our findings to standard image classification tasks trained with state-of-the-art models. Our neuroscience-inspired task is the checkerboard (CB) task (Chandrasekaran et al., 2017; Kleinman et al., 2019). In the CB task, one discerns the dominant color of a checkerboard filled with red and green squares. The subject then makes a reach to a left or right target whose color matches the dominant color in the checkerboard (Fig 1a). This task therefore involves making two binary choices: a color decision (i.e., reach to the red or green target) and a direction decision (i.e., reach to left or right). Critically, the color of the targets (red left, green right; or green left, red right) is
random on every trial. The direction decision output is conditionally independent of the color decision, as detailed further in Fig 1b and Section B.6, even though the color information needs to be used to solve the task. This task allows us to evaluate how both of these components of information are represented through training and across layers.
We used this task and extensions to study the evolution of minimal representations during training. If a representation is sufficient and minimal, we refer to this representation as optimal (Achille & Soatto, 2018). Our contributions are the following. (1) We introduce a notion of usable information for studying representations and training dynamics in deep networks (Section 3). (2) We used this notion to characterize the transient training dynamics in deep networks by studying the amount of usable relevant and irrelevant information in deep network layers and across training epochs. We first use the CB task to gain intuition of the training dynamics in a simplified setting. We find that training with SGD is critical to bias the network toward learning minimal representations in intermediate layers (Section 4.1). This adds to the literature suggesting that SGD results in minimal representations of input information (Achille & Soatto, 2018; Shwartz-Ziv & Tishby, 2017) while avoiding some of the pitfalls. (3) We used the intuition gained from the simple task, evaluating our findings on CIFAR-10 and CIFAR-100 task using modern architectures. Remarkably, we find that the networks increased usable information about an irrelevant component of information early in training and discarded it later on in training to arrive at a minimal sufficient solution, consistent with a proposed (Shwartz-Ziv & Tishby, 2017) though controversial theory (Saxe et al., 2018).

heading: 2 RELATED WORK

text:
Some efforts to understand why neural networks generalize focus on representation learning, that is, how deep networks learn optimal (i.e., minimal and sufficient) representations of inputs in order to solve a task. Typically, representation learning is focused on studying the properties of the asymptotic representations after training (Achille & Soatto, 2018). Recent work suggests that these asymptotic representations contain minimal but sufficient input information for performing a task (Achille & Soatto, 2018; Shwartz-Ziv & Tishby, 2017). Implicit regularization coming from SGD, and in particular from the use of large learning rates and small batch sizes, is believed to play an important role in forming these minimal sufficient representations.
How does the training process lead to these minimal but sufficient asymptotic representations? Shwartz-Ziv & Tishby (2017) propose that there are two distinct phases of training: an empirical risk minimization phase where the network minimizes the loss on the training set, and a “compression” phase where the network discards information about the inputs that do not need to be represented to solve the task. Recently, Saxe et al. (2018) challenged this view, arguing that the observed compression was dependent on the activation function and the mutual information estimator used in Shwartz-Ziv & Tishby (2017). These works highlight the challenges of estimating mutual information to study how representations emerge through training.
In general, estimating mutual information from samples is challenging for high-dimensional random variables (Paninski, 2003). The primary difficulty in estimating mutual information is estimating a high-dimensional probability distribution from the samples, since generally the number of samples required scales exponentially with the dimension. This is impractical for realistic deep learning tasks where the representations are high dimensional. To estimate the mutual information, Shwartz-Ziv & Tishby (2017) used a binning approach, discretizing the activations into a finite number of bins. While this approximation is exact in the limit of infinitesimally small bins, in practice, the size of the bin affects the estimator (Saxe et al., 2018; Goldfeld et al., 2018). In contrast to binning, other approaches to estimate mutual information include entropic-based estimators (e.g., Goldfeld et al. (2018)) and a nearest neighbours approach (Kraskov et al., 2004). Although mutual information is difficult to estimate, it is an appealing quantity to summarily characterize key aspects of the transient neural network training behavior because of its invariance to smooth and invertible transformations. In this work, rather than estimate the mutual information directly, we instead define and study the “usable information” in the network, which corresponds to a variational approximation of the mutual information (Barber & Agakov, 2003; Poole et al., 2019) (see Sections 3 and A.1). Recently, such variational approximations to mutual information have been viewed as a meaningful characterization of representations in deep networks, and the theoretical underpinnings of this approach are beginning to be investigated (Xu et al., 2020; Dubois et al., 2020).
Research into the training dynamics of deep networks, and how they represent relevant and irrelevant task information, is nascent. A related study by Achille et al. (2019) found that early periods of training were critical for determining the asymptotic network behavior. Additionally, it was found that the timing of regularization was important for determining asymptotic performance (Golatkar et al., 2019), with regularization during this “critical period” having the most influential effect. Notably, both of these studies found an initial increase in the amount of information that weights encode about the dataset (as measured by the Fisher information), that coincides with the critical period of learning. This phase is followed later in training by a “forgetting” phase where the network discards unnecessary information. This suggests that a similar dynamic to the one we study can be observed in weight space instead of representation space.

heading: 3 USABLE INFORMATION IN A REPRESENTATION

text:
A deep neural network consists of a set of ` layers, with each layer forming a successive representation of the input. A representation Z` may store information in a variety of ways. It may be that a complex transformation is required to read out the information, or it may be that a simple linear decoder could read out the information. In both cases, from an information-theoretic perspective, the same information is contained in the representation, however, there is an important distinction regarding how “usable” this information is. Information is usable if later layers, which comprise affine transformations and element-wise nonlinearities, can easily extract it to solve the task. Equivalently, usable information should be decodable by a separate neural network also employing affine transformations and element-wise nonlinearities.
Formally, we define the usable information that a representation Z contains about a quantity Y , which may refer to the output or a component of the input, as:
Iu(Z;Y ) = H(Y )− LCE(p(y|z), q(y|z)). (1)
Here, H(Y ) is the entropy, or uncertainty, of Y , and LCE is the cross-entropy loss on the test set of a discriminator network q(y|z) trained to approximate the true distribution p(y|z). Our definition is motivated in the following manner. The test set cross-entropy loss approximates how much uncertainty there is in the output Y given Z and the discriminator. A low loss implies that there is low uncertainty in Y given Z, or that the discriminator can extract a lot of “information” about Y from Z. If the logarithm in the cross-entropy loss is in base 2, it is measured in bits. If the value of Y were approximately the same for any Z, there would be little uncertainty in Y to begin with, so it is important to know the amount of uncertainty in Y given Z with respect to the initial uncertainty in Y . What is most relevant is the amount of remaining uncertainty in Y given Z. Thus we use the
difference in uncertainty H(Y )−LCE as the amount of “usable information” that Z contains about Y , as shown in our definition in Equation 1.
This definition is appealing to study representations, in part, because it can be computed from samples of Z and Y , and is a quantity that is comparable across network training. We estimate LCE using a small neural network that learns a distribution q(y|z). To train the network, we sample activations Z and the quantity Y and learn q(y|z) by minimizing the cross-entropy loss on a training set. We then evaluate the LCE on the test set (Equation 1). We provide details about the neural network and the training we used for decoding in Appendix B.3 and C.2. We also show in the Appendix that the usable information is a lower bound on the mutual information (Appendix A.1). Importantly, usable information also is not constrained by the data processing inequality; that is, the information can be made more “usable” by transformation to later layers, consistent with the representation learning view that later layers are forming improved representations of the inputs (Xu et al., 2020).

heading: 4 EXPERIMENTS

text:
Our goal was to characterize how optimal representations are formed through SGD training. We trained multiple network architectures on tasks and assessed the usable information in representations across layers and training epochs. For a given architecture and task, all hyper-parameters were kept constant throughout experiments, unless explicitly stated.
To develop intuition, we initially investigate how small fully connected networks represent the relevant and irrelevant information in the CB Task. We trained two different network architectures, ‘Small FC’: 5 layers, with 10−7−5−4−3 units in each layer, ‘Medium FC’: 100−20−20−20. Small FC was a network used in prior literature (Shwartz-Ziv & Tishby, 2017; Saxe et al., 2018). Our networks were fully-connected and used ReLU activation. We trained the networks using SGD with a constant learning rate to perform the CB task, described in detail in Appendix B.4. The hyper-parameters used for the CB experiments are listed in Appendix B.5.
In our CB task experiments, we quantified the usable color and direction information in the hidden representation, Z`. In the n = 2 CB task, the color information represents half of the input information. We emphasize that, unless otherwise specified, the network was only trained to output the correct direction choice, so given a representation of the direction, representing the color choice is irrelevant. Therefore, a minimal representation should not include information about the color choice, since it is not necessary to represent given a representation of the direction decision. To make the task more complex, we also generalized the CB task to have n = 10 and n = 20 targets.
We then use this framework to examine how relevant and irrelevant information are represented in more realistic tasks and architectures, and how hyper-parameters affect the learning dynamics. We define a coarse labelling of task labels and study how the network represents the fine and coarse labelling through training, using a ResNet-18 (He et al., 2016) and All-CNN (Springenberg et al., 2015) on CIFAR-10 and CIFAR-100.

heading: 4.1 SGD WITH RANDOM INITIALIZATION RESULTS IN MINIMAL SUFFICIENT

text:
REPRESENTATIONS IN THE CB TASK
We first assessed the optimality of the network representations by training Small FC networks on the CB task using n = 2 colors (Fig 2a) using a random initialization for the weights. In particular, the initial weights do not contain information about the dataset. We computed the usable color and direction information across layers of the neural network and epochs of training. In our plots, later layers are denoted by darker shades. In deeper layers, there was a decrease in usable color information, corresponding to more minimal representations. After training, the asymptotic representation in the last layer contained zero usable color information and 1 bit of usable direction information. To visualize this minimal sufficient representation, we plotted the activations of the 3 units in the last layer of the Small FC network for different inputs. These visualizations are labeled by the correct color (red and green) and direction (cross or circle). In the asymptotic representation, representation of the input color is overlapping (red and green), while the representation of the direction output is separable (crosses and circles), forming a minimal sufficient representation.
(a)
(b)
(c)
To test if this observed minimality was a result of our simple task, we extended the CB task to a variant with n input checkerboard colors, with n corresponding output direction classes. We trained networks using a larger architecture (Medium FC). We show results for n = 10 and n = 20 classes in Fig 2b,c. We observed similar phenomena to the n = 2 case: there was decreasing usable color information in deeper layers, and nearly zero color information in the last layer’s representation. In contrast, there was significant usable direction information across all layers in the asymptotic representation, with usable information about the direction increasing for deeper layers. We validated our results using different random initializations (Figures 9, 10, 11).
These results show that, for a simple task with SGD and random initialization, minimal sufficient representations emerge through training. Asymptotic representations were sufficient to perform the task, but contained less usable color information in deeper layers, approaching zero color information in the last layer. In this simple task, we observed that it was possible for the network to solve the task with nearly zero usable color information in its last layer across training (Fig 2b,c).
We also examined how changing the initialization by pretraining the network to output the color choice affected the resulting representations. We found that the resulting representations were not

heading: Resnet

text: minimal for the n = 2 checkerboard case (Fig 6a), retaining some structure from the initialization (Fig 6d). This result also held for the CB task with n = 10 and n = 20 (Fig 6b,c). Furthermore, we found that pretraining on the color choice led to worse generalization performance (Fig 7).

heading: 4.2 ACQUISITION AND FORGETTING OF USABLE INFORMATION IN MODERN DEEP NETWORKS

text:
Using a similar approach as we did for the CB task to characterize relevant and irrelevant information, we next investigated how modern deep neural networks trained with SGD learned task representations. To study learning dynamics, we investigated (1) how networks learned and represented task information as well as information about a representative semantically meaningful variable, and (2) how this information was represented across training epochs. To this end, we defined coarse labels corresponding to groups of classes in the CIFAR-10 and CIFAR-100 datasets. The CIFAR-100 dataset defines fine labels corresponding to each of the 100 classes, as well as 20 coarse labels corresponding to meaningful groupings of 5 from the 100 classes. In the CIFAR-10 case, we defined two coarse labels arbitrarily, corresponding to even and odd class labels. Thus, when training the network to output the coarse label, we can investigate the network’s representation of the semantically meaningful fine label description, which serves as a proxy for the computation and representations that the network is learning. We note that, when trained to output the coarse label, a minimal representation should contain no additional information about the fine label.
We trained a ResNet-18 (He et al., 2016) to output the coarse label of CIFAR-10, using an initial learning rate of 0.1 with exponential annealing (0.97), momentum (0.9), and a batch size of 128. We investigated the usable information in the last layer of the ResNet-18, which has a dimension of 512. We found that while training the network to predict the coarse-grained class, the network acquired information about the coarse-grained class, evidenced by an increase in usable information during training (orange curve) while validation accuracy (green dashed line; scale on the right hand side of plot) was increasing (Fig 3). Strikingly, while the validation accuracy and usable coarse-grained class information increased, the information about the fine labels first increased and then decreased (around epoch 100). It then decreased to minimality, storing no additional usable information about the fine labels than was contained in the coarse labels. These learning dynamics were proposed Shwartz-Ziv & Tishby (2017), but due to controversies of their information estimation and experimental setup, have been widely debated (Saxe et al., 2018). We emphasize that even though we did
(a)
0 20 40 60 80 100 120 140 160 Epoch
0.0
0.5
1.0
1.5
2.0
2.5 3.0 Us ab le in fo rm at io n (b its )
fine coarse
70
75
80
85
90
95
Va lid
at io
n ac
cu ra
cy
Batch size: 64
(b)
0 20 40 60 80 100 120 140 160 Epoch
0.0
0.5
1.0
1.5
2.0
2.5
3.0
Us ab
le in
fo rm
at io
n (b
its )
fine coarse
60
70
80
90
Va lid
at io
n ac
cu ra
cy
Batch size: 256
(c)
0 20 40 60 80 100 120 140 160 Epoch
0.0
0.5
1.0
1.5
2.0
2.5
3.0
Us ab
le in
fo rm
at io
n (b
its )
fine coarse
55
60
65
70
75
80
85
90
Va lid
at io
n ac
cu ra
cy
Batch size: 512
(d)
Learning rate: 7.5e-2
(e)
Learning rate: 5e-2
(f)
Learning rate: 1e-2
not ask the network to acquire information about the fine labels, SGD naturally led the network to learn information about the fine label, and then decreased this information later in training.
Together, these results show that SGD tends to result in minimal representations, which may be guided by interesting learning dynamics. To achieve this minimality, the network displays a learning motif where it learns additional information early in training, then discards it later on. We next investigate how these findings depend on hyper-parameter choices, architecture, and task.

heading: 4.3 SENSITIVITY OF USABLE INFORMATION TRAINING DYNAMICS TO HYPER-PARAMETERS, ARCHITECTURE, AND TASK

text:
Using this framework, we evaluated how hyper-parameter choices affected the learning dynamics in deep networks. We focus on the ResNet-18 trained on CIFAR-10 in Figure 3. We varied the batch sizes from 64 to 1024 and found that a small batch size led to dynamics similar to that of Fig 3, while a larger batch size did not lead to minimal representations (Fig 4a-c). Results for a batch size of 1024 are shown in Fig 8. The learning rate also affected the learning dynamics. We found that all networks increased the information about the fine labels during training. However, we found that only for large initial learning rates did the network “forget” the superfluous information. Results for a learning rate of 0.001 are also shown in the appendix in Fig 8. We found that small learning rates (0.001) or large batch sizes (512 or larger) led to lower validation accuracy. Thus, the implicit regularization coming from the use of SGD with a small batch size and large learning rate, which is common in practical settings, is crucial for learning minimal sufficient representations. Here we have provided an underpinning for these choices by exposing their associated learning dynamics.
Additionally, we investigated whether the phenomenon of acquiring “superfluous” task information was common across different architectures and tasks. We used an All-CNN (Springenberg et al., 2015) trained on CIFAR-10 to output the binary coarse label, observing a similar trend with an
(a)
(b)
(c)
increase and decrease in the usable information about the fine label (Fig 5a). In this case, the information about the fine label did not decrease to minimality, but nonetheless, there was a significant reduction in the fine label information, suggesting that SGD naturally compresses additional input information. Finally, we evaluated how a ResNet-18 represented task information using the CIFAR100 dataset. This dataset is accompanied with 100 fine labels and 20 coarse labels, corresponding to groupings of the 100 classes. We used the same hyper-parameters as in Fig 3. We trained the network to output the coarse labels, observing an increase to approximately 3.5 bits of usable information. The network achieved a nearly minimal representation (Fig 5b).
It is important to note that for this setting of hyper-parameters in the CIFAR-100 task (the same as in the CIFAR-10 case), SGD did not show a visible increase followed by a decrease in usable information in the fine labels, a result different than what we observed in CIFAR-10. We conjectured this could be due to at least three potential reasons: (1) the hyper-parameter settings may be suboptimal, which we observed may result in learning dynamics that do not increase then decrease fine information (Fig 4c, f). (2) In CIFAR-100, coarse and fine labels are semantically similar, so there may not be not much more information to be naturally learned in the fine than the coarse labels, and further that it is possible that while the information about the fine labels remains approximately flat, the network is forgetting information about aspects of the fine labels while learning other parts of fine label information in the process of increasing coarse label information and arriving at a nearly minimal representation. (3) CIFAR-100 has relatively few examples, 500 per fine label, impacting the learning of fine label information. Despite these limitations, our results from CIFAR-10 suggest that SGD learning dynamics that increase then decrease information about the fine label should result in more optimal representations and higher validation accuracy. To test this, we performed an experiment where we pretrained the network to output fine label information until epoch 20, after which the network then was trained to output coarse information. This training process resulted in learning dynamics that resembled SGD learning in Fig 3. We observed that these learning dynamics resulted in networks with a 2.1% increase in validation accuracy (compare Fig 5b and c). These results support that learning dynamics that increase, and then decrease, information about inputs, may result in more optimal representations that achieve higher validation accuracy.

heading: 5 DISCUSSION

text:
We introduced a notion of the usable information in the representation, which reflects the amount of information that can be extracted by a learned decoder, for understanding the training dynamics
in deep networks. This definition is appealing, in part, due to its flexibility. For instance, if it is important to understand how accessible the information is to a linear decoder, it suffices to apply our formulation of usable information using a linear decoder trained with cross-entropy loss. In contrast, if the goal is to extract all information present in a representation, regardless of how accessible this information is, one can train a high capacity nonlinear decoder. Since neural networks are powerful function approximators, as the function approximation improves, the decoder will approach the optimal decoder. In this case, the usable information approaches Shannon mutual information, as the lower bound becomes tight (Section A.1). Future theoretical and empirical work should investigate the tightness of this bound and its dependence on training parameters.
In our case, we used a relatively small nonlinear neural network as the decoder, which provided insight into the evolution of optimal representations through training on simple tasks inspired by neuroscience literature and on image classification tasks. These tasks allowed us to show that the implicit regularization of SGD plays an important role in learning minimal sufficient representations. In particular, in standard hyper-parameter settings, we observed learning dynamics where the network learns to encode semantically meaningful but ultimately irrelevant information early in training, before later discarding this information to arrive at a minimal sufficient representation.
Monkeys performing the checkerboard task, like our networks, also had minimal sufficient representations in an output (motor) area (Chandrasekaran et al., 2017; Kleinman et al., 2019). Despite the obvious implementation differences of both information processing systems, we speculate that the general effects coming from a noisy learning process, which led to minimal sufficient representations in our artificial networks, may be an important factor leading to minimal sufficient representations in biological networks.
It is remarkable that in the CIFAR-10 task, SGD naturally exploited the semantically meaningful structure of the fine labels, in order to solve the coarse labels task. In general, it is difficult to identify the features that are being learned during training, and whether they correspond to something semantically meaningful. However by defining a coarse label, our task setup allowed us to study how semantically meaningful information was represented during training. During training, the network increased the information about the semantically important part of the input, even when only asked to output the coarse label. It then decreased the information later in training. We did not notice such a major increase in CIFAR-100, perhaps due to the nature of the dataset or hyperparameter configuration. However, by inducing the network to follow similar learning dynamics to Fig 3 by pretraining the network to output the fine labels, we were able to improve the performance on the coarse labelling task. This suggests that a detailed understanding of the training dynamics and the features learned is important for learning optimal representations and successfully transferring representations between tasks.
Using usable information, we observed an increase and decrease in the information about an irrelevant variable, which has been proposed (Shwartz-Ziv & Tishby, 2017), but has been debated, largely due to controversies over the estimation of Shannon’s mutual information (Saxe et al., 2018). Our observation is in accordance with the ideas of Shwartz-Ziv & Tishby (2017), and importantly we have observed these dynamics on modern architectures and realistic tasks. Our results are also consistent with a complementary view of information in the weights, where it has been observed that the Fisher Information increased and decreased during training (Achille et al., 2019), corresponding to a critical period in neural network training.

heading: ACKNOWLEDGMENTS

text: We thank the anonymous reviewers for their thoughtful feedback and suggestion of new experiments. MK was supported by the National Sciences and Engineering Research Council (NSERC). JCK was supported by an NSF CAREER Award 1943467, a UCLA Computational Medicine Amazon Web Services Award, and a NVIDIA GPU Grant.

heading: A PROOFS

heading: A.1 USABLE INFORMATION LOWER BOUNDS THE MUTUAL INFORMATION

text:
The entropy of a distribution is defined as H(x) = Ex∼p(x) [ log 1
p(x)
] . (2)
The mutual information, I(X;Y ), can be written in terms of an entropy term and a conditional entropy term:
I(Z;Y ) = H(Y )−H(Y |Z). (3)
We want to show that:
I(Z;Y ) ≥ Iu(Z;Y ) := H(Y )− LCE(p(y|z), q(y|z)) (4) It suffices to show that:
H(Y |Z) ≤ LCE (5)
where LCE is the cross-entropy loss on the test set. For our study, H(Y ) represented the known distribution of output classes, which in our case were equiprobable.
H(Y |Z) := E(z,y)∼p(z,y) [ log 1
p(y|z)
] (6)
= E(z,y)∼p(z,y) [ log 1
q(y|z) ] ︸ ︷︷ ︸
cross-entropy loss
−Ez∼p(z) [KL(p(y|z)||q(y|z)]︸ ︷︷ ︸ ≥0 , (7)
≤ E(z,y)∼p(z,y) [ log 1
q(y|z)
] := LCE (8)
To approximate H(Y |Z), we first trained a neural network with cross-entropy loss to predict the output, Y , given the hidden activations, Z, learning a distribution q(y|z). The KL denotes the Kullback-Liebler divergence. We multiplied (and divided) by an arbitrary variational distribution q(y|z) in the logarithm of equation 6, leading to equation 7. The first term in equation 7 is the crossentropy loss commonly used for training neural networks. The second term is a KL divergence and is therefore non-negative. In our approximator, the distribution q(y|x) is parametrized by a neural network. When the distribution q(y|z) = p(y|z), our variational approximation of H(Y |Z), and hence approximation of I(Z;Y ) is exact (Barber & Agakov, 2003; Poole et al., 2019).

heading: B ADDITIONAL RESULTS AND DETAILS IN THE CHECKERBOARD TASK

heading: B.1 SGD WITH NON-RANDOM INITIALIZATION MAY NOT FORM MINIMAL

text:
REPRESENTATIONS IN THE CB TASK
Implicit regularization in SGD is hypothesized to result in a minimal representation through compression of irrelevant input information, also called a “forgetting” phase (Shwartz-Ziv & Tishby, 2017; Achille & Soatto, 2018; Achille et al., 2019). We tested this hypothesis by initializing networks with significant color information, and subsequently performing SGD on the CB task. We then evaluated whether SGD resulted in networks with minimal color representations. We initialized the weights by pretraining the network to output the color decision for 20 epochs, which required the network to represent color information. After 20 epochs, we reverted to training on the CB task, where only the direction decision was reported. Since the learning rate was kept constant, the pretrained weights can be viewed as a different initialization in parameter space for the modified task.
Strikingly, we found that the resulting representations were not minimal for the n = 2 checkerboard case (Fig 6a). This result also held for the CB task with n = 10 and n = 20 (Fig 7b,c). While we observed some compression of usable color information through training, the asymptotic representations had significantly greater than zero color information. In Fig 7b, we observed all layers
(a)
(b)
(c)
had more usable color information than the direction information in the first layer. The network therefore solved the task using an alternative representation that was not minimal. We visualized the activations corresponding to the asymptotic non-minimal representations of Small FC in Fig 6d. In the early epochs the red and green points converge (both crosses and dots) as a result of successful pretraining. However, when we trained the CB task starting at epoch 20, the representations changed. While the dot clusters for red and green checkerboards are overlapping, the cross clusters are not. This representation is not minimal as color information can be decoded above chance.
These results show that the initialization affects the asymptotic representation of neural networks. SGD, under particular initializations, may not lead to minimal representations of the task inputs. This suggests there is a trade-off between learning a minimal representation and simply reusing the existing representations present in the initial weights. Initial structure in the network representations from pretraining, such as the separation of the red and green crosses in the last layer representation, was maintained even when performing SGD to train a different task. Together, these results suggest that while SGD compresses representations towards minimality, it finds a solution that is functionally related to the initial representation. This may correspond to a optima in the neighborhood of the initialization.

heading: B.2 RELATIONSHIP BETWEEN PRETRAINING, MINIMALITY, AND GENERALIZATION IN THE

text:
CB TASK
Our results show that the minimality of network representations, and therefore solutions, depends on initialization. All trained networks (for n larger than 2), however, achieved zero training er-
(a)
(b)
ror. A natural question to ask is how does the pretraining affect the resulting representation and generalization performance?
To answer this, we varied the number of epochs that we pretrained the CB tasks of n = 2, n = 10, and n = 20 classes, and quantified the usable color and direction information, as well as the trained network’s test accuracy to understand how the network generalizes (Fig. 7). We found that networks trained with longer pretraining had less minimal representations and worse generalization performance. This was true regardless of the number of classes, but the effect was more pronounced (in terms of absolute difference in accuracy) when the network did not solve the task perfectly without pretraining. We note that regardless of how long the networks were pretrained for, the networks were subsequently trained for the same number of epochs (80), with the same learning rate throughout training. One interpretation is that when using existing structure to solve the task, the network learned a suboptimal solution to solving the task, increasing the chance of overfitting. Another interpretation is that the pretraining changed the distribution of the weights, affecting the minimality and generalization.

heading: B.3 DETAILS OF NEURAL NETWORK FOR USABLE INFORMATION IN THE CB TASK

text: To estimate usable information, we computed the cross-entropy loss of a decoder q(y|z) that predicts Y from Z. The decoder was a three-layer neural network, with 128, 64, and 32 units per layer, with Leaky-ReLU activations (slope = 0.2), batch-norm and dropout (p = 0.7). At each epoch, 1250 training samples were generated and supplied to the decoder, along with either the corresponding correct direction or color choice. We evaluated the cross-entropy loss on 3750 test samples to minimize overfitting. We trained the network for 100 epochs using a learning rate of 0.5 for ‘Medium FC’ and 0.05 for ‘Small FC.’

heading: B.4 CHECKERBOARD TASK DESCRIPTION

text: Following the conventions of Kleinman et al. (2019), we modeled the CB task (Fig 1a), inputting the checkerboard color and target configuration to a neural network that outputted the direction choice (Fig 1b). We minimized the cross-entropy loss of the network output and the ground truth output. We extended the checkerboard task to the n checkerboard task by increasing the number of checkerboards. Each target was 1 out of the n colors, with the targets forming an ‘n-polygon’. The correct direction corresponds to the direction of the target having the same color of the checkerboard. We specified the color of each target using a one-hot encoding, and the color of the checkerboard as a one-hot encoding. Noise with mean 0 and standard deviation of 0.1 was added to the checkerboard inputs. The target and checkerboard color inputs were concatenated to form an input vector. The correct direction of the target was the output.

heading: B.5 DETAILS OF CB EXPERIMENTS

text: The following are the hyper-parameters used in our experiments. We trained two different network architectures, ‘Small FC’: 5 layers, with 10−7−5−4−3 units in each layer, ‘Medium FC’: 100− 20− 20− 20. We trained networks using SGD with a constant learning rate throughout training.

heading: FC Small, n = 2:

text:
• batch size: 32, learning rate: 0.05, number of data samples: 10000 (90% train, 10% validation)
Medium FC, n = 10:
• batch size: 64, learning rate: 0.5, number of data samples: 25000 (90% train, 10% validation)
Medium FC, n = 20:
• batch size: 128, learning rate: 0.5, number of data samples: 50000 (90% train, 10% validation)
Medium FC, n = 25:
• batch size: 128, learning rate: 0.5, number of data samples: 75000 (90% train, 10% validation)

heading: B.6 DEFINITION OF RELEVANT AND IRRELEVANT INFORMATION IN THE CB TASK

text: In the CB task, the color of the checkerboard and target configuration (inputs) are necessary to determine the correct direction to reach (output). While both a color and direction decision are made, after the direction is determined, the color decision no longer needs to be represented: the network can generate the correct output with only the direction representation. Formally, the output y is conditionally independent of the color representation, Zc, given the direction representation Zd (i.e., y ⊥⊥ (Zc, Zt)|Zd, as illustrated by the graph in Fig 1b). Hence, given a representation of the direction choice, the color choice (and target configuration) no longer needs to be represented. We emphasize that, in general, the output is not independent of the color representation and target configuration representation Zt, i.e., y 6⊥⊥ (Zc, Zt), hence information about the dominant color of the checkerboard is necessary to compute y. When this conditional independence holds, we call the conditionally independent variable “irrelevant.” We therefore refer to the color choice as “irrelevant” and the direction choice as “relevant.” We study how these components evolve together throughout training.

heading: C ADDITIONAL RESULTS AND DETAILS IN THE CIFAR-10 AND CIFAR-100 TASK

heading: C.1 CIFAR-10 AND CIFAR-100 TASK DESCRIPTION

text: We trained a ResNet-18 and an All-CNN architecture to output a superclass corresponding to the twenty coarse-grained classes in CIFAR-100 and, in CIFAR-10, to an arbitrary superclass corresponding to the even and odd classes. Accordingly, a minimal representation should only encode the superclass.

heading: C.2 DETAILS OF NEURAL NETWORK FOR USABLE INFORMATION

text: To estimate usable information, we computed the cross-entropy loss of a decoder q(y|z) that predicts Y from Z. We used a two-layer neural network, with 200 and 100 with Leaky-ReLU activations (slope = 0.2), batch-norm and dropout (p = 0.7). At each epoch, 7500 samples were supplied to the decoder, along with either the corresponding correct direction or color choice. We evaluated the cross-entropy loss on 2500 test samples. We trained the network for 50 epochs using Adam with a learning rate of 0.01 and weight decay of 0.001.

heading: C.3 DETAILS OF NEURAL NETWORK TRAINING

text:
In our experiments, unless otherwise stated, we trained a ResNet-18 (He et al., 2016) with an initial learning rate of 0.1 decaying smoothly with a factor of 0.97 at each epoch, batch size of 128, momentum of 0.9 and weight decay with coefficient 0.0005. For the All-CNN (Springenberg et al., 2015) we used a batch size of 128, initial learning rate of 0.05 decaying smoothly by a factor of 0.97 at each epoch, momentum of 0.9, and weight decay with coefficient 0.001. We used standard data augmentation with random translations up to 4 pixels and random horizontal flipping. These parameter configurations were taken directly from prior work (Achille et al., 2019).
D ADDITIONAL PLOTS