Funded by: Afeyan Foundation
PI: Rafayel Petrosyan
Researchers: Varazdat Stepanyan, Anahit Khachatryan
In this project, one of the approaches we took was the application of the well-established framework in a nonequilibrium statistical mechanics of reconstructing the equilibrium free energy profile of the system by performing repetitive nonequilibrium measurements on it. These methods were developed by Hummer and Szabo using Jarzynski equality or Crooks fluctuation theorem. If we state it in a simple way, in Deep Learning (DL), we are searching for a minimum on a hypersurface using a gradient descent algorithm, for example. This analogy between the free energy landscape reconstruction in statistical mechanics and the search for the minimum on an analogous landscape in DL gives an opportunity to study the geometry of neural network training using the tools of nonequilibrium statistical mechanics.
In order to demonstrate this, we used a simple Convolutional Neural Network (CNN) for classifying handwritten digits from the MNIST dataset, which is a model dataset in deep learning. Our CNN has two convolutional blocks (Conv → ReLU → MaxPool) followed by a fully connected classifier with dropout. Input is a 1×28×28 image, and output is one of the digits 0–9. After three epochs, it reached around 99% test accuracy.
In analogy to a physical system, the cross-entropy loss corresponds to the energy landscape, the weights corresponding to microstates, the learning rate is related to temperature, and stochastic gradient descent is analogous to simulated annealing toward a Boltzmann distribution. Next we would like to know whether we could probe the global structure of the landscape. We built a steered Langevin scheme, namely, we train two endpoint networks, A and B, then drag the weights from A toward B using a moving harmonic spring while measuring the work done. From 24 such trajectories, combined with the Jarzynski estimator, gave us a free-energy reconstruction along the path between the two minima. We then added the standard variance-reduction by pulling in both directions (A→B and B→A), applying the Crooks fluctuation theorem to find where the forward and reverse work distributions cross (which equals ΔF), and using the Bennett Acceptance Ratio (BAR) to combine forward and reverse samples. Although the Jarzynski/BAR approach is computationally more expensive, it reveals the structure of the loss landscape between minima. The figure below shows that the naive linear interpolation between two trained networks results in a sharp barrier of ~1.7, while the Jarzynski free-energy reconstruction along the same path reports a smooth bump of ~0.6. The barrier is largely an artifact of pointwise evaluation.