Visualizing High-Dimensional Trajectories on the Loss-Landscape of ANNs

1 INTRODUCTION . Artificial neural networks ( ANNs ) have been successfully used to solve a number of complex tasks in a diverse array of domains , such as object recognition , machine translation , image generation , 3D protein structure prediction and many more . Despite being highly overparameterized for the tasks they solve , and having the capacity to memorize the entire training data , ANNs tend to generalize to unseen data . This is a spectacular feat since the highly non-convex optimization typically encountered in them should ( theoretically ) be a significant obstacle to using these models ( Blum & Rivest , 1993 ) . Questions such as why ANNs favor generalization over memorization and why they find good minima even with intricate loss functions still remain largely unanswered . One promising research direction for answering them is to look at the loss-landscape of deep learning models . Recent work tried to approach this task by proposing various visualization methods . An emerging challenge here is how to look at such an extremely high dimensional optimization landscape ( linear in the number of parameters of the network ) with respect to minimized loss . In past work , loss functions and their level lines were visualized via random directions starting at a minimum , or by means of linear methods like PCA . In some case , this approach proved effective in uncovering underlying structures in the loss-landscape and link them to network characteristics , such as generalization capabilities or structural features ( Keskar et al. , 2016 ; Li et al. , 2018 ) . However , these methods have two major key drawbacks : ( 1 ) they are linear in that they only choose directions that are linear combinations of parameter axes while the loss landscape itself is highly nonlinear , and ( 2 ) they choose only two among thousands ( if not millions ) of axes to visualize and ignore all others . In this work , we utilize and adapt the PHATE dimensionality reduction method ( Moon et al. , 2019 ) , which relies on diffusion-based manifold learning , to study ANN loss landscapes by visualizing the evolution of network weights during training in low dimensions . In general , visualizations like PHATE Moon et al . ( 2019 ) are specifically designed to squeeze as much variability as possible into two dimensions , and thus provide an advantage over previous approaches . In particular our choice of using PHATE over other popular methods , such as tSNE ( van der Maaten & Hinton , 2008 ) , is due to its ability to capture global and local structures of data , and in particular to keep intact the training trajectories that are traversed through during gradient descent . Indeed , during training , the high-dimensional neural networks weights change significantly while remaining on a connected manifold defined by the support of viable configurations ( e.g. , with sufficiently low training loss ) , which we refer to when discussing the geometry of the loss landscape . We show that PHATE is suitable to track such continuous weight trajectories , as opposed to tSNE or UMAP that tend to shatter them . Moreover , our approach provides general view of relevant geometric patterns that emerge in the high-dimensional parameter space , providing insights regarding the properties of ANN training and reflecting on their impact on the loss landscape . Contributions : We propose a novel loss-landscape visualization based on a variation of PHATE , implemented with cosine distance in Section 4 . Our method is , to our knowledge , different from all other proposed methods for loss visualization in that it is naturally nonlinear and captures data characteristics from all dimensions . In Section 5.1 , we show that our method uncovers key geometric patterns characterizing loss-landscape regions surrounding good and bad generalization optima , as well as memorization optima . Finally , we establish the robustness of our method by applying it to numerous tasks , architectures , and optimizers in Sections 5.3 and 5.2 , suggesting that our method can be used in a consistent manner to validate training and design choices . 2 RELATED WORK . Loss landscape visualization methods have been proposed in numerous contexts . Goodfellow et al . ( 2014 ) proposed the “ linear path experiment ” where the loss of an ANN is evaluated at a series of points θ = ( 1− α ) θi + αθf for different values of α ∈ [ 0 , 1 ] and θi , θf corresponding to the initial parameters of the model and the found optima in parameter space respectively . This one-dimensional linear interpolation method has allowed them to show that popular state of the art ANNs typically do not encounter significant obstacles along a straight path from initialization to convergent solution . They also used the method to visualize the loss along directions connecting two distinct minima and to show that these are linearly separated by a region of higher valued loss . This method was further developed by Im et al . ( 2016 ) , who adapted it to enable the visualization of two-dimensional projections of the loss-landscape using barycentric and bilinear interpolation for groups of three or four points in parameter space . This analysis method has allowed them to establish that despite starting with the same parameter initialization , different optimization algorithms find different minima . Furthermore , they noticed that the loss-landscape around minima have characteristic shapes that are optimizer-specific and that batch-normalizations smooths the loss function . More recently , Li et al . ( 2018 ) have addressed the scale invariance and network symmetries problems discussed in Neyshabur et al . ( 2017 ) ; Dinh et al . ( 2017 ) , which prevented meaningful comparisons between loss-landscape plots from different networks . They proposed 1D and 2D linear interpolation plots , similar to past techniques , but where they used filter-wise normalized directions to remove the scaling effect . This method has allowed them to visualize and compare the regions on the loss-landscape surrounding minima coming from multiple networks in a meaningful way and to correlate the “ flatness ” of the region to the generalization capabilities of the corresponding network . Furthermore , they studied the effects of the network depth , width and the presence of skip connections on the geometry of the loss-landscape and on network generalization . The importance of loss landscape visualization methods like the one presented in this paper increases with the growing scientific community interest in furthering our understanding of ANNs . As our understanding of this landscape gets deeper , we start uncovering more and more high-dimensional and complex geometric and topological characteristics . For instance , Draxler et al . ( 2018 ) have found nonlinear pathways in parameter space connecting distinct minima , along which the training and test errors remain small and losses are consistently low . This suggests that minima are not situated in isolated valleys but rather on connected manifolds representing low loss regions . However , such characteristics are intrinsically high-dimensional , making linear methods inadequate to visualize these structures . Even in standard applications of the linear methods one inevitably asks if the thousands ( or even millions ) of unseen directions do not hide critical features of the landscape , and if the visualized linear path is relevant to what happened during training . In this paper , we suggest the use of modern dimensionality reduction techniques to study the complex loss-landscape of ANNs . Such techniques have been extensively used in recent years to successfully study the internal learned representations of deep networks and highlight their complex geometric structures and intrinsic dimensionality ( Gigante et al. , 2019 ; Horoi et al. , 2020 ; Recanatesi et al. , 2018 ; Farrell et al. , 2019 ; Maheswaranathan et al. , 2019 ) , and here we further advance this line of work to provide new applications and insights in the study of ANN optimization and generalization . 3 PRELIMINARIES : PHATE DIMENSIONALITY REDUCTION & VISUALIZATION . Given a n×m data matrix N , where n is the number of data points and m the number of features , PHATE computes low-dimensional embeddings of the data points ( dimension d to be specified by the user ) which preserve both their global and their local structure . It first computes the pairwise distance matrix D using a specified distance function φ , each element of the matrix being given by Dij = φ ( Ni , : , Nj , : ) . To better capture local affinities between data points , an affinity matrix A is computed using an α-decaying kernel with a locally-adaptive bandwidth k , i corresponding to the k-NN distance of the i-th data point . The elements of A are given by : Ai , j = Kk , α ( i , j ) = 1 2 exp ( − ( Di , j k , i ) α ) + 1 2 exp ( − ( Di , j k , j ) α ) The decaying factor α regulates the decay rate of the kernel ( smaller α⇒ kernel with lighter tails ) , α = 2 corresponding to the Gaussian . The affinity matrix is then row-normalized to obtain the diffusion operator P , a row-stochastic Markov transition matrix with element Pi , j giving the probability of jumping from the i-th to the j-th data point in one time step . One of the reasons PHATE excels at capturing global structures in data , especially high-dimensional trajectories and branches , is that it leverages the diffusion operator ( also used , Coifman & Lafon , 2006 , to construct diffusion maps ) by running the implicit Markov chain forward in time . This is accomplished by raising the matrix P to the power t , effectively taking t random walk steps and revealing intrinsic geometric structure through the affinity of the data at a scale that grows with the number of steps . The optimal value for t is automatically chosen to be the knee point of the von Neumann entropy of P. To enable dimensionality reduction while retaining diffusion geometry information from the operator , PHATE leverages information geometry to define a pairwise potential distance as an M-divergence IDi , j = ‖ logPi , : − logPj , :‖2 between corresponding t-step diffusion probability distributions of the two points , which provides a global context to each data point . The resulting information distance matrix ID is finally embedded into a visualizable low dimensional ( 2D or 3D ) space by metric multidimensional scaling ( MDS ) , thereby squeezing the intrinsic geometric information into two dimensions ( as is also done in tSNE ) , rather than producing a high dimensional alternative coordinate system like PCA or diffusion maps . We reason that since PHATE learns a data manifold via diffusion , and places points in the global context of other points , it can keep trajectories connected from beginning to end without shattering them . Indeed , this property was established quantitatively in Moon et al . ( 2019 ) via a Denoised Manifold Affinity Preservation ( DeMAP ) metric . Here , it leads to the preservation trajectory flow directions and relations between them , as demonstrated in Figure 1 , in comparison with other methods , on trajectory-like data and ANN training trajectories , as well as in our results presented in the remainder of the paper . 4 METHODS . The loss-landscape is a high dimensional function f ( θ ) that assigns a loss value to every possible vector θ in the high-dimensional parameter space of an ANN denoted Θ . The dimensionality of the space renders the task of completely visualizing the loss-landscape virtually impossible . By gathering locally connected loss-landscape patches throughout training and around minima found via gradient descent , we hope to reconstruct manifolds on which these trajectories lie . Despite the thousands , if not millions , of dimensions of the parameter space Θ , we hypothesize that these trajectories and manifolds are intrinsically of low dimensional and thus PHATE , which reduced dimensionality via manifold learning and information geometry , will help facilitate their visualization in two specific settings : 1 . To characterize the region of the loss-landscape surrounding minima ; and 2 . To simultaneously visualize multiple training trajectories corresponding to different parameter initializations and optimizers in order to make training choices . In setting ( 1 ) we use the following “ jump and retrain ” experiment . Given a trained ANN and a minimum found during training with an optimizer : • Let θo represent the vector of network parameters at the minimum ; • For seed ∈ { 0 , 1 , 2 , 3 , 4 } : – For step_size ∈ { 0.25 , 0.5 , 0.75 , 1 } : 1 . Choose the random vector vseed in the parameter space of the network Θ corre- sponding to the random seed seed ; 2 . Filter normalize the vector as proposed in Li et al . ( 2018 ) to obtain vseed ; 3 . Set the network parameters at each jump initialization to be θjump-init = θo + step_size · vseed ; 4 . Retrain for 50 epochs using optimizer with 0.9 momentum on the original training set ; 5 . Record network parameters θ for each seed and step size combination and at the end of each epoch . • Apply PHATE to the data matrix [ θseed , step_size , epoch ] and visualize the embeddings . In setting ( 2 ) , we simultaneously visualize multiple training trajectories corresponding to different parameter initializations and optimizers . Here , we simply follow the optimizer from initialization to convergence , and record the parameter vector at each step . We note that methods that pick random directions or surfaces around the minima projected to random directions are not able to visualize entire trajectories as the spaces and planes change . This is a distinct advantage of using all of the dimensions of the parameter space . In both settings , PHATE has allowed us to bypass the key drawbacks of previously proposed linear interpolation methods by : 1 . Capturing variance in data from all relevant dimensions and embed it in a low-dimensional space ; and 2. preserving high-dimensional trajectories and global structures of data in parameter space . All modern dimensionality-reduction techniques would have achieved ( 1 ) in one way or another and with various degrees of success . The reason our method is based on cosine-distance PHATE , is to better accomplishes ( 2 ) by preserving trajectories and global relationships in high-dimensional data to a greater extent than any other state-of-the-art dimensionality-reduction technique . Figure 1 demonstrates this by showing a comparison of multiple such techniques , namely PHATE , PCA , t-SNE ( van der Maaten & Hinton , 2008 ) and UMAP ( McInnes & Healy , 2018 ) , and how they each embed the results of a jump and retrain experiment ( Figure 1A and B ) and an artificial data set having a tree-like structure ( Figure 1C ) in a 2D space . While some trajectory-like structure is visible in all low-dimensional embeddings , only PHATE properly captures intra-trajectory variance . PHATE is also the only technique that captures the global relationships in between trajectories while t-SNE and UMAP have a tendency to cluster points that are close in parameter space and disregard the global structure of the data . On the artificial data set , what we observe is that the embeddings of the linear method PCA are highly affected by the noise in the data while t-SNE and UMAP have a tendency of shattering trajectories that should be connected .

Summary: This paper uses PHATE to visualize the progression of neural net parameters during learning in neural networks to provide insight into generalizable vs. non-generalizable minima, and the behaviors of different optimization algorithms.  PHATE is an improvement over previous visualization techniques due to its approach to finding a manifold, allowing it to plot in two dimensions multiple trajectories that do not otherwise share a plane.  This is then used to plot trajectories in "jump and retrain" experiments, in which a minimum is found, and then perturbations made to the network parameters, before restarting training.  It is shown that in the networks experimented on, minima with good test set performance reliably funnel the new learning trajectories into the same minimum, while minima with poor test set performance see the perturbed initializations find other minima, thus demonstrating the "flatness vs sharpness" of minima.  Trajectories produced by SGD, SGD with momentum, and Adam are also compared.  Adam is shown to travel further, but along a smoother trajectory.