In this paper, the authors have tried to demystify the role played by the intermediate layers of a neural network. To understand the dynamics of intermediate layers, they propose a new method to understand the design of Neural networks better and help us understand common heuristics (such as adding auxiliary losses). The technique uses linear probes or classifiers, which use the hidden units at specific intermediate layers as discriminating features. They explain how these linear probes help visualize the model's state during the training phase.
Information Theory
The authors begin with the fundamental question of the information contained in neural networks. More specifically, their questions revolve around the following: What would be the effect of more layers? Do they help in training, and how important is each layer to the model's performance? How does information flow from layer 1 to the output layer? Does having multiple auxiliary losses help in training a neural network model? The deep-ness of a model is often associated with its complexity, but this isn't necessarily correlated with the information contained in each layer. For example, one might claim that the higher layers have more information about the ground truth, but that would be incorrect. The authors show why the notion of entropy fails to capture the essence of the above questions using Data Processing Inequality. The Data processing inequality can be expressed in a toy example like the following. We calculate the mutual information between (X, Y ) and (X, Z) where X, Y, and Z are three random variables satisfying the following relationship.
X → Y → Z
We observe that the mutual information between X and Z is less than the mutual information between X and Y,
I(X : Z) < I(X : Y )
where I(X, Y ) is the mutual information between X and Y . Intuitively, this would mean that adding layers isn't contributing anything meaningful while training the model. The notion of entropy also fails to explain the working efficiently: Entropy helps distil the essence of information to a quantity that doesn't depend on the representations. This idea is in direct conflict with the idea of a Neural Network relying on a series of simple deterministic transformations that would affect the final representation of the input data (which is then fed to a linear classifier).
Thus it is essential to have a better notion of information in the context of neural networks, where representation matters and computational complexity is also dependent on it. Hence, the authors propose a new idea of the information that depends on the ability of optimal linear classifiers to classify features at intermediate layers. They chose a linear classifier as it is the most basic form of a classifier that isn't trivial.
Understanding models using Linear Probes
To understand whether the transformations made to the data (in a Neural network) are made suddenly or progressively, the authors introduce the concept of linear probes. The probes are added to every layer of a model. The probes aren't supposed to change the model's training, and the gradient does not propagate back through these connections.
The authors try to show that information in the context of neural networks can be seen as an immediate suitability for a linear classifier. The parameters of the linear probes(classifiers) are updated after freezing the gradients of the original model (neural network). This happens at each training step.
The probes are trained till convergence with pixel model parameters, and the prediction error is calculated on the training set.
The authors perform three sets of experiments:
- Probes on an untrained model: The authors take a 32-layer multi-layer perceptron with 128 hidden units. All the layers are fully connected, and LeakyRELU is used as an activation function. They initialize the MLP using glorot normal initialization and do not perform any training on the model. The probes are optimized with RMSProp and a low learning rate(<0.001). The same experiment is conducted 100 times with a different toy dataset. The experimental results demonstrate that each successive layer's input signal is mangled and becomes less and less valuable as it reaches the final layer. The results complement the theory of vanishing/exploding gradients while training a neural network. As the usefulness of the layers decays as one goes deeper, the values contained in the last layer also influence the upcoming gradients. As the derivative is based on garbage values, the model performance deteriorates with each pass.
- Probes on a trained model (MNIST convnet) The authors used a model comprised of Conv2D, Maxpool and dense layers, with ReLU as an activation between intermediate layers. The experimental results show that the randomly initialized layers provide functional transformations at the start of training. Random projections of the data help us to understand it better. The prediction error drops from 80% to 20%(calculated using linear probes) by just using the random features. This also helps us understand that a significant chunk of the learning occurs in the first few layers of the model.
- Probes with auxiliary loss add to the model. The authors construct a 128-layer model, adding auxiliary loss at every 16 layers. A probe is placed at every layer, and the performance is estimated in three steps. In this experiment, the probes contribute to loss minimization. This is analogous to auxiliary heads used in Szegedy et al. The addition of auxiliary losses helps the model to converge successfully.
Conclusion
The authors present the usage of linear classifiers or probes as a tool to understand the dynamics of a neural net. They show how each probe can be used to identify unexpected behaviour in models (non-convergence, garbage predictions etc.). They also plan on using multiple probes in each intermediate layer as a part of their future experiments.