Original: Visualizing high-dimensional trajectories on the loss-landscape of ANNs
Revised: Visualizing high-dimensional trajectories on the loss-landscape of ANNs

Original: Abstract
Revised: Abstract

Original: Training artiﬁcial neural networks requires the optimization of highly non-convex loss functions.
Revised: Training artiﬁcial neural networks requires the optimization of highly non-convex loss functions.

Original: Throughout the years, the scientiﬁc community has developed an extensive set of tools and architectures that render this optimization task tractable and a general intuition has been developed for choosing hyper parameters that help the models reach minima that generalize well to unseen data.
Revised: Throughout the years, the scientiﬁc community has developed an extensive set of tools and architectures that render this optimization task tractable and a general intuition has been developed for choosing hyper parameters that help the models reach minima that generalize well to unseen data.

Original: However, for the most part, the difference in trainability in between architectures, tasks and even the gap in network generalization abilities still remain unexplained.
Revised: However, for the most part, the difference in trainability in between architectures, tasks and even the gap in network generalization abilities still remain unexplained.

Original: Visualization tools have played a key role in uncovering key geometric characteristics of the losslandscape of ANNs and how they impact trainability and generalization capabilities.
Revised: Visualization tools have played a key role in uncovering key geometric characteristics of the losslandscape of ANNs and how they impact trainability and generalization capabilities.

Original: However, most visualizations methods proposed so far have been relatively limited in their capabilities since they are of linear nature and only capture features in a limited number of dimensions.
Revised: However, most visualizations methods proposed so far have been relatively limited in their capabilities since they are of linear nature and only capture features in a limited number of dimensions.

Original: We propose the use of the modern dimensionality reduction method PHATE which represents the SOTA in terms of capturing both global and local structures of high-dimensional data.
Revised: We propose the use of the modern dimensionality reduction method PHATE which represents the SOTA in terms of capturing both global and local structures of high-dimensional data.

Original: We apply this method to visualize the loss landscape during and after training.
Revised: We apply this method to visualize the loss landscape during and after training.

Original: Our visualizations reveal differences in training trajectories and generalization capabilities when used to make comparisons between optimization methods, initializations, architectures, and datasets.
Revised: Our visualizations reveal differences in training trajectories and generalization capabilities when used to make comparisons between optimization methods, initializations, architectures, and datasets.

Original: Given this success we anticipate this method to be used in making informed choices about these aspects of neural networks.
Revised: Given this success we anticipate this method to be used in making informed choices about these aspects of neural networks.

Original:  Introduction
Revised:  Introduction

Original: Artiﬁcial neural networks (ANNs) have been successfully used to solve a number of complex tasks in a diverse array of domains.
Revised: Artiﬁcial neural networks (ANNs) have been successfully used to solve a number of complex tasks in a diverse array of domains.

Original: However, questions such as why ANNs favor generalization over memorization and why they ﬁnd good minima even with intricate loss functions still remain largely unanswered.
Revised: However, questions such as why ANNs favor generalization over memorization and why they ﬁnd good minima even with intricate loss functions still remain largely unanswered.

Original: One promising research direction for answering them is to look at the loss-landscape of deep learning models.
Revised: One promising research direction for answering them is to look at the loss-landscape of deep learning models.

Original: Recent work has approached this task by proposing various visualization methods.
Revised: Recent work has approached this task by proposing various visualization methods.

Original: In most papers, loss functions and their level lines have been visualized via linear methods like PCA.
Revised: In most papers, loss functions and their level lines have been visualized via linear methods like PCA.

Original: In some case, this approach proved effective in uncovering underlying structures in the loss-landscape and linking them to ANN generalization or architecture [1, 2].
Revised: In some case, this approach proved effective in uncovering underlying structures in the loss-landscape and linking them to ANN generalization or architecture [1, 2].

Original: However, these methods have two major key drawbacks: (1) 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 millions of axes to visualize and ignore all others.
Revised: However, these methods have two major key drawbacks: (1) 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 millions of axes to visualize and ignore all others.

Original: In this work, we utilize and adapt the PHATE dimensionality reduction method [3], which relies on diffusion-based manifold learning, to study ANN loss landscapes by visualizing the evolution of network weights during training in low dimensions.
Revised: In this work, we utilize and adapt the PHATE dimensionality reduction method [3], which relies on diffusion-based manifold learning, to study ANN loss landscapes by visualizing the evolution of network weights during training in low dimensions.

Original: In general, unlike PCA, visualizations like
Revised: In general, unlike PCA, visualizations like

Original: PHATE [3] are speciﬁcally designed to squeeze as much variability as possible into two dimensions, and thus provide an advantage over previous approaches.
Revised: PHATE [3] are speciﬁcally designed to squeeze as much variability as possible into two dimensions, and thus provide an advantage over previous approaches.

Original: In particular our choice of using PHATE over other popular methods, such as tSNE [4], is due to its ability to retain both global and local structures of data, and in particular to keep intact the training trajectories that are traversed through during gradient descent.
Revised: In particular our choice of using PHATE over other popular methods, such as tSNE [4], is due to its ability to retain both global and local structures of data, and in particular to keep intact the training trajectories that are traversed through during gradient descent.

Original: Indeed, During training, the high-dimensional neural networks weights change signiﬁcantly while remaining on a connected manifold deﬁned by the support of viable conﬁgurations which we refer to when discussing the geometry of the loss landscape.
Revised: Indeed, During training, the high-dimensional neural networks weights change signiﬁcantly while remaining on a connected manifold deﬁned by the support of viable conﬁgurations which we refer to when discussing the geometry of the loss landscape.

Original: We show that
Revised: We show that

Original: PHATE is suitable to track such continuous weight trajectories, as opposed to tSNE or UMAP that tend to shatter them.
Revised: PHATE is suitable to track such continuous weight trajectories, as opposed to tSNE or UMAP that tend to shatter them.

Original: Moreover, our approach provides general view of relevant geometric patterns that emerge in the high-dimensional parameter space, providing insights regarding the properties of
Revised: Moreover, our approach provides general view of relevant geometric patterns that emerge in the high-dimensional parameter space, providing insights regarding the properties of

Original: ANN training and reﬂecting on their impact on the loss landscape.
Revised: ANN training and reﬂecting on their impact on the loss landscape.

Original:  Related work
Revised:  Related work

Original: Loss landscape visualization methods in literature are mostly based on linear interpolations, either 1D interpolations between the initial and ﬁnal network parameters [5], or scanning a 2D grid around the optimum [2].
Revised: Loss landscape visualization methods in literature are mostly based on linear interpolations, either 1D interpolations between the initial and ﬁnal network parameters [5], or scanning a 2D grid around the optimum [2].

Original: There is growing interest in uncovering the high-dimensional and complex geometric and topological characteristics of loss landscapes beyond linear paths.
Revised: There is growing interest in uncovering the high-dimensional and complex geometric and topological characteristics of loss landscapes beyond linear paths.

Original: Nonlinear pathways with consistently low losses connecting distinct minima were found in [6], suggesting the existence of connected low loss manifolds.
Revised: Nonlinear pathways with consistently low losses connecting distinct minima were found in [6], suggesting the existence of connected low loss manifolds.

Original: Such high-dimensional manifolds are beyond the capability of linear methods.
Revised: Such high-dimensional manifolds are beyond the capability of linear methods.

Original: Even in standard applications of the linear methods one inevitably asks if the 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.
Revised: Even in standard applications of the linear methods one inevitably asks if the 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.

Original: We suggest the use of dimensionality reduction techniques to study the complex loss-landscape of ANNs.
Revised: We suggest the use of dimensionality reduction techniques to study the complex loss-landscape of ANNs.

Original: Such techniques have been extensively used recently to successfully study the internal learned representations of ANNs and highlight their complex geometric structures and intrinsic dimensionality [7, 8, 9, 10, 11].
Revised: Such techniques have been extensively used recently to successfully study the internal learned representations of ANNs and highlight their complex geometric structures and intrinsic dimensionality [7, 8, 9, 10, 11].

Original: Here we further advance this line of work to provide new applications and insights in the study of ANN optimization and generalization.
Revised: Here we further advance this line of work to provide new applications and insights in the study of ANN optimization and generalization.

Original:  Background: PHATE dimensionality reduction & visualization
Revised:  Background: PHATE dimensionality reduction & visualization

Original: Given a data matrix N , PHATE ﬁrst computes the pairwise similarity matrix A (e.g. via a Gaussian kernel), then row-normalize A to obtain the diffusion operator P , a row-stochastic Markov transition matrix where P i,j denotes the probability of moving from the i -th to the j -th data point in one time step.
Revised: Given a data matrix N , PHATE ﬁrst computes the pairwise similarity matrix A (e.g. via a Gaussian kernel), then row-normalize A to obtain the diffusion operator P , a row-stochastic Markov transition matrix where P i,j denotes the probability of moving from the i -th to the j -th data point in one time step.

Original: Powers of the matrix P t represents moving along the Markov chain for t steps, where t is selected automatically as the knee point of the Von Neumann Entropy of the diffusion operator.
Revised: Powers of the matrix P t represents moving along the Markov chain for t steps, where t is selected automatically as the knee point of the Von Neumann Entropy of the diffusion operator.

Original: To enable dimensionality reduction while retaining diffusion geometry information from the operator,
Revised: To enable dimensionality reduction while retaining diffusion geometry information from the operator,

Original: PHATE leverages information geometry to deﬁne a pairwise potential distance as an M-divergence
Revised: PHATE leverages information geometry to deﬁne a pairwise potential distance as an M-divergence

Original: ID i,j = (cid:107) log P i, : − log P j, : (cid:107) 2 between corresponding t -step diffusion probability distributions of the two points, which provides a global context to each data point.
Revised: ID i,j = (cid:107) log P i, : − log P j, : (cid:107) 2 between corresponding t -step diffusion probability distributions of the two points, which provides a global context to each data point.

Original: The resulting information distance matrix ID is ﬁnally embedded into a visualizable low dimensional (2D or 3D) space by metric multidimensional scaling (MDS), thereby squeezing the intrinsic geometric information to calculate the ﬁnal 2D or 3D embeddings of the data.
Revised: The resulting information distance matrix ID is ﬁnally embedded into a visualizable low dimensional (2D or 3D) space by metric multidimensional scaling (MDS), thereby squeezing the intrinsic geometric information to calculate the ﬁnal 2D or 3D embeddings of the data.

Original: For further details, see Moon et al.
Revised: For further details, see Moon et al.

Original:  Methods
Revised:  Methods

Original: The dimensionality of the space renders the task of completely visualizing the loss-landscape of
Revised: The dimensionality of the space renders the task of completely visualizing the loss-landscape of

Original: ANNs virtually impossible.
Revised: ANNs virtually impossible.

Original: 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.
Revised: 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.

Original: Despite the high dimensions of the parameter space Θ , we hypothesize that these trajectories and manifolds are intrinsically low dimensional and thus PHATE, which reduced dimensionality via diffusion-based manifold learning and information geometry, will help facilitate their visualization in two speciﬁc settings: 1.
Revised: Despite the high dimensions of the parameter space Θ , we hypothesize that these trajectories and manifolds are intrinsically low dimensional and thus PHATE, which reduced dimensionality via diffusion-based manifold learning and information geometry, will help facilitate their visualization in two speciﬁc settings: 1.

Original: To characterize the region of the loss-landscape surrounding minima; 2.
Revised: To characterize the region of the loss-landscape surrounding minima; 2.

Original: To simultaneously visualize multiple training trajectories corresponding to different parameter initializations and optimizers in order to make training choices.
Revised: To simultaneously visualize multiple training trajectories corresponding to different parameter initializations and optimizers in order to make training choices.

Original: In setting (1) we use the following “jump and retrain” experiment.
Revised: In setting (1) we use the following “jump and retrain” experiment.

Original: Given a trained ANN and the ﬁnal parameters θ o ,
Revised: Given a trained ANN and the ﬁnal parameters θ o ,

Original: • For combinations of seed ∈ { 0 , 1 , 2 , 3 , 4 } and step _ size ∈ { 0 .25 , 0 . 5 , 0 .
Revised: • For combinations of seed ∈ { 0 , 1 , 2 , 3 , 4 } and step _ size ∈ { 0 .

Original: Choose the random v seed in Θ and ﬁlter normalize it as proposed in [2] to obtain v seed ; 2.
Revised: Choose the random v seed in Θ and ﬁlter normalize it as proposed in [2] to obtain v seed ; 2.

Original: Set the network parameters at each jump initialization to be θ jump-init =
Revised: Set the network parameters at each jump initialization to be θ jump-init =

Original: θ o + step _ size · v seed ; 3.
Revised: θ o + step _ size · v seed ; 3.

Original: Retrain for 50 epochs on the original training set and record θ at the end of each epoch.
Revised: Retrain for 50 epochs on the original training set and record θ at the end of each epoch.

Original: • Apply PHATE to the data matrix [ θ seed,step _ size,epoch ] and visualize the embeddings.
Revised: • Apply PHATE to the data matrix [ θ seed,step _ size,epoch ] and visualize the embeddings.

Original: In setting (2) , we simultaneously visualize multiple training trajectories corresponding to different parameter initializations and optimizers.
Revised: In setting (2) , we simultaneously visualize multiple training trajectories corresponding to different parameter initializations and optimizers.

Original: In both settings, PHATE has allowed us to bypass the key drawbacks of previously proposed linear interpolation methods by: 1.
Revised: In both settings, PHATE has allowed us to bypass the key drawbacks of previously proposed linear interpolation methods by: 1.

Original: Capturing variance in data from all relevant dimensions and embed it in a low-dimensional space; and 2.
Revised: Capturing variance in data from all relevant dimensions and embed it in a low-dimensional space; and 2.

Original: Preserving high-dimensional trajectories and global relationships to a greater extent than any other state-of-the-art dimensionality reduction technique.
Revised: Preserving high-dimensional trajectories and global relationships to a greater extent than any other state-of-the-art dimensionality reduction technique.

Original: Figure 1 demonstrates this by showing a comparison of multiple such techniques, namely PHATE, PCA, t-SNE [4] and UMAP [12], and how they each embed the results of a jump and retrain experiment.
Revised: Figure 1 demonstrates this by showing a comparison of multiple such techniques, namely PHATE, PCA, t-SNE [4] and UMAP [12], and how they each embed the results of a jump and retrain experiment.

Original: While some trajectory-like structure is visible in all low-dimensional embeddings, only PHATE properly captures intra-trajectory variance and 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.
Revised: While some trajectory-like structure is visible in all low-dimensional embeddings, only PHATE properly captures intra-trajectory variance and 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.

Original:  Generalization vs. Memorization
Revised:  Generalization vs. Memorization

Revised: PHATE embeddings of the jump and retrain experiment conducted on one good ( a ) and one bad ( b ) generalization minima as well as one memorization minimum ( c ).

Revised: The initial points of each retraining are marked by diamonds of different colors based on step _ size .

Revised: All trajectories are colored by their respective direction of the initial jump and descending hue, meaning the color gets whiter as retraining progresses.

Revised: See Appendix B for a higher resolution version.

Original: Arguably one of the most important and general questions in machine learning at the moment is: by looking only at a given trained model and its training data, can we predict how well the model will perform on unseen data?
Revised: Arguably one of the most important and general questions in machine learning at the moment is: by looking only at a given trained model and its training data, can we predict how well the model will perform on unseen data?

Original: Previous work offers a partial and vague answer to this question by linking the “ﬂatness” of neural network optima to the extent to which the network is able to generalize.
Revised: Previous work offers a partial and vague answer to this question by linking the “ﬂatness” of neural network optima to the extent to which the network is able to generalize.

Original: This idea was initially presented in [13] and has recently been the focus of numerous papers [e.g., 14, 1, 15, 16, 2].
Revised: This idea was initially presented in [13] and has recently been the focus of numerous papers [e.g., 14, 1, 15, 16, 2].

Original: Flat minima, around which the value of the loss function is fairly constant, are believed to correspond to conﬁgurations where the network generalizes well to the unseen data while the opposite is believed for sharp minima.
Revised: Flat minima, around which the value of the loss function is fairly constant, are believed to correspond to conﬁgurations where the network generalizes well to the unseen data while the opposite is believed for sharp minima.

Original: Following the underlying idea of these propositions, we believe that studying the geometry of the loss landscape around minima has the potential of uncovering key characteristics about the network’s ability to generalize.
Revised: Following the underlying idea of these propositions, we believe that studying the geometry of the loss landscape around minima has the potential of uncovering key characteristics about the network’s ability to generalize.

Original: Here, we show that jump-and-retrain experiments combined with our visualization method reveal stark differences between neural network conﬁgurations that generalize well versus ones that memorize (or overﬁt).
Revised: Here, we show that jump-and-retrain experiments combined with our visualization method reveal stark differences between neural network conﬁgurations that generalize well versus ones that memorize (or overﬁt).

Original: We trained multiple WResNet28-2 networks on CIFAR10 with and without scrambled labels to obtain generalization and memorization minima.
Revised: We trained multiple WResNet28-2 networks on CIFAR10 with and without scrambled labels to obtain generalization and memorization minima.

Original: Varying the training procedure for the generalization task allowed us to ﬁnd "good" and "bad" minima each reaching ~95% and ~85% accuracy on the test set respectively, for more details see
Revised: Varying the training procedure for the generalization task allowed us to ﬁnd "good" and "bad" minima each reaching ~95% and ~85% accuracy on the test set respectively, for more details see

Original: Appendix E.
Revised: Appendix E.

Original: While the starting points are all set on  distinct directions from the minimum, the plotted retraining trajectories are not restricted to any subspace of the parameter space.
Revised: While the starting points are all set on ﬁve distinct directions from the minimum, the plotted retraining trajectories are not restricted to any subspace of the parameter space.

Original: Despite having only a ~7% difference in their accuracy on the test data set, the PHATE embedding of the trajectories surrounding each generalization minimum is signiﬁcantly different.
Revised: Despite having only a ~7% difference in their accuracy on the test data set, the PHATE embedding of the trajectories surrounding each generalization minimum is signiﬁcantly different.

Original: The minimum that shows good generalization has a distinctive ﬂower-petal convergence pattern indicating that even when points are thrown far from the minimum, the landscape is wide enough that they return to it in roughly the same direction they were pushed out.
Revised: The minimum that shows good generalization has a distinctive ﬂower-petal convergence pattern indicating that even when points are thrown far from the minimum, the landscape is wide enough that they return to it in roughly the same direction they were pushed out.

Original: This structure is evocative of a high-dimensional “wide valley”.
Revised: This structure is evocative of a high-dimensional “wide valley”.

Original: In the bad minimum, the trajectories start off near the middle of the plot, but during the retraining the trajectories diverge toward the edges before coming back towards the middle of the plot.
Revised: In the bad minimum, the trajectories start off near the middle of the plot, but during the retraining the trajectories diverge toward the edges before coming back towards the middle of the plot.

Original: This outward movement is stark contrast with the consistent retraining trajectories surrounding the good gen. minimum.
Revised: This outward movement is stark contrast with the consistent retraining trajectories surrounding the good gen. minimum.

Original: This indicates that the region around the bad gen. minimum contains a non-negligible amount of non-convexities, preventing training trajectories to systematically converge towards the minimum.
Revised: This indicates that the region around the bad gen. minimum contains a non-negligible amount of nonconvexities, preventing training trajectories to systematically converge towards the minimum.

Original: The “memorization” minimum plot Fig.
Revised: The “memorization” minimum plot Fig.

Original: 2c, f looks similar to the bad generalization plot, with trajectories that go outward at small step sizes of the "jump".
Revised: 2c, f looks similar to the bad generalization plot, with trajectories that go outward at small step sizes of the "jump".

Original: However, curiously at larger step sizes of the "jump," the trajectories seem to return without going outward ﬁrst, but they do not return immediately, i.e., they show some lateral movement, potentially indicating bumpiness in the landscape that they are avoiding.
Revised: However, curiously at larger step sizes of the "jump," the trajectories seem to return without going outward ﬁrst, but they do not return immediately, i.e., they show some lateral movement, potentially indicating bumpiness in the landscape that they are avoiding.

Original: See supplemental video for an animation of these trajectories, and Appendix B for trajectories colored by loss.
Revised: See supplemental video for an animation of these trajectories, and Appendix B for trajectories colored by loss.

Original: Our method is robust across initializations and tasks since same patterns arise for experiments ran on CIFAR100 and for multiple initializations of the network on CIFAR10, see Appendix F for ﬁgures.
Revised: Our method is robust across initializations and tasks since same patterns arise for experiments ran on CIFAR100 and for multiple initializations of the network on CIFAR10, see Appendix F for ﬁgures.

Original:  Optimization
Revised:  Optimization

Original: We analyzed three stochastic gradient-descent optimization methods commonly used for training
Revised: We analyzed three stochastic gradient-descent optimization methods commonly used for training

Original: ANNs: SGD (vanilla), SGD with momentum (SGD_M), and Adam.
Revised: ANNs: SGD (vanilla), SGD with momentum (SGD_M), and Adam.

Original: The major difference is that
Revised: The major difference is that

Original: Adam uses adaptive step sizes based off previous iterations while SGD and SGD_M use ﬁxed step sizes.
Revised: Adam uses adaptive step sizes based off previous iterations while SGD and SGD_M use ﬁxed step sizes.

Original: First we examine individual trajectories of network weights from the same seed using different optimizers (Fig. 3A).
Revised: First we examine individual trajectories of network weights from the same seed using different optimizers (Fig. 3A).

Original: Compared to the smooth trajectory from using Adam, SGD with or without momentum seems to exhibit scattered patterns.
Revised: Compared to the smooth trajectory from using Adam, SGD with or without momentum seems to exhibit scattered patterns.

Original: It should be noted that the magnitudes of network weights learned by Adam are much larger than by SGD or SGD_M (Appendix C).
Revised: It should be noted that the magnitudes of network weights learned by Adam are much larger than by SGD or SGD_M (Appendix C).

Original: The evolution and resulting network weights from the Adam optimizer appears to differ more signiﬁcantly from SGD or
Revised: The evolution and resulting network weights from the Adam optimizer appears to differ more signiﬁcantly from SGD or

Original: SGD_M, where the resulting parameters appear to be further away from the initialization.
Revised: SGD_M, where the resulting parameters appear to be further away from the initialization.

Original: A (1)
Revised: A (1)

Original: A (2)
Revised: A (2)

Original: A (3)
Revised: A (3)

Original: B (1)
Revised: B (1)

Original: B (2)
Revised: B (2)

Original: B (3)
Revised: B (3)

Original: Adam trajectories, plotted by themselves are continuous, however distances between ﬁnal weights reached by Adam are much larger than other optimizers.

Original: See Appendix D for higher resolution.

Original: In Fig.
Revised: In Fig.

Original: 3B the distinctiveness between optima reached from different initializations and optimizers are clearly shown.
Revised: 3B the distinctiveness between optima reached from different initializations and optimizers are clearly shown.

Original: We see in the Fig.
Revised: We see in the Fig.

Original: 3B(1) that when starting from the same seed, SGD and SGD_M travel in similar directions.
Revised: 3B(1) that when starting from the same seed, SGD and SGD_M travel in similar directions.

Original: However, we see a shorter and straighter training path with momentum.
Revised: However, we see a shorter and straighter training path with momentum.

Original: Without momentum, a similar minima is reached.
Revised: Without momentum, a similar minima is reached.

Original: However, the path towards the minima is longer with more moves sideways, indicating that the space is searched less effectively.
Revised: However, the path towards the minima is longer with more moves sideways, indicating that the space is searched less effectively.

Original: In Fig.
Revised: In Fig.

Original: 3B(2)
Revised: 3B(2)

Original: the fact that Adam travels so far in the weight space, causes the initial points look very close together.
Revised: the fact that Adam travels so far in the weight space, causes the initial points look very close together.

Original: From the initial points, the SGD_M travels in shorter paths towards different minima.
Revised: From the initial points, the SGD_M travels in shorter paths towards different minima.

Original: However,
Revised: However,

Original: Adam seems to exhibit a disconnected jump, potentially ﬂying off into an orthogonal space to what is visualized, but interestingly returning to the same region of minima as SGD_M. The seemingly discontinuity of the trajectories in the middle plots indicates that Adam searches space very differently than SGD.
Revised: Adam seems to exhibit a disconnected jump, potentially ﬂying off into an orthogonal space to what is visualized, but interestingly returning to the same region of minima as SGD_M. The seemingly discontinuity of the trajectories in the middle plots indicates that Adam searches space very differently than SGD.

Original: Fig.
Revised: Fig.

Original: 3B(3) shows that the Adam trajectories are indeed connected in a continuous line, so the trajectory itself is not discontinuous, but rather in a different subspace than the SGD_M trajectory.
Revised: 3B(3) shows that the Adam trajectories are indeed connected in a continuous line, so the trajectory itself is not discontinuous, but rather in a different subspace than the SGD_M trajectory.

Original:  Discussion and conclusion
Revised:  Discussion and conclusion

Original: We proposed a novel approach to visualize ANN loss landscapes based on the state-of-the-art PHATE dimensionality reduction, which is able to capture branch-like structures in high-dimensional data in two dimensional representations by leveraging diffusion-based manifold learning and information geometry.
Revised: We proposed a novel approach to visualize ANN loss landscapes based on the state-of-the-art PHATE dimensionality reduction, which is able to capture branch-like structures in high-dimensional data in two dimensional representations by leveraging diffusion-based manifold learning and information geometry.

Original: Our approach enables geometric exploration of retrained trajectories surrounding generalization and memorization optima, found via ANN training, to provide insight into generalization capabilities of the network.
Revised: Our approach enables geometric exploration of retrained trajectories surrounding generalization and memorization optima, found via ANN training, to provide insight into generalization capabilities of the network.

Original: Further, it enables meaningful comparison of training trajectories across optimizers.
Revised: Further, it enables meaningful comparison of training trajectories across optimizers.

Original: Finally, our proposed paradigm is demonstrably consistent and robust across initializations and supervised learning tasks.
Revised: Finally, our proposed paradigm is demonstrably consistent and robust across initializations and supervised learning tasks.

Original: We expect in future work this visualization approach to enable more methodical paradigms for the development of deep models that generalize better, train faster, and provide fundamental understanding of their capabilities.  
Revised: We expect in future work this visualization approach to enable more methodical paradigms for the development of deep models that generalize better, train faster, and provide fundamental understanding of their capabilities.  Acknowledgments

Original: {+This work was partially funded by NSERC’s CGS M scholarship [S.H.]; IVADO Professor startup & operational funds [ G.W. ]; CZI grants 182702 & CZF2019-
+}
Revised: {+This work was partially funded by NSERC’s CGS M scholarship [S.H.]; IVADO Professor startup & operational funds [ G.W. ]; CZI grants 182702 & CZF2019-
+}

Revised: [ S.K. ]; and NIH grants

Revised: R01GM135929 & R01GM130847 [ G.W., S.K. ].

Revised: The content provided here is solely the responsibility of the authors and does not necessarily represent the ofﬁcial views of the funding agencies.

Original: E Experimental details
Revised: E Experimental details

Original: For our empirical results of Section 5 we trained multiple WResNet28-2 networks (wide ResNet of depth 28 and 2 times as wide as a standard ResNet [17]) presented in [18] on the CIFAR10 dataset [19] using batches of size 128 for 200 epochs with the SGD optimizer with 0.9 momentum and a learning rate of 0.1 which decays by a factor of 10 at epoch 150.
Revised: For our empirical results of Section 5 we trained multiple WResNet28-2 networks (wide ResNet of depth 28 and 2 times as wide as a standard ResNet [17]) presented in [18] on the CIFAR10 dataset [19] using batches of size 128 for 200 epochs with the SGD optimizer with 0.9 momentum and a learning rate of 0.1 which decays by a factor of 10 at epoch 150.

Original: In some cases we did not use data augmentation and weight decay and in others we used 1e-4 weight decay and random crops and horizontal ﬂips for the training data.
Revised: In some cases we did not use data augmentation and weight decay and in others we used 1e-4 weight decay and random crops and horizontal ﬂips for the training data.

Original: Adding the data augmentation and the weight decay has allowed us to ﬁnd optima that generalize better, increasing the accuracy on the test data from around 85% to around 95%.
Revised: Adding the data augmentation and the weight decay has allowed us to ﬁnd optima that generalize better, increasing the accuracy on the test data from around 85% to around 95%.

Original: We will be referring the optima with ~95% test accuracy as “good optima” since they generalize better, and that with lower test accuracy as “bad optima”.
Revised: We will be referring the optima with ~95% test accuracy as “good optima” since they generalize better, and that with lower test accuracy as “bad optima”.

Original: Using the same training procedure as the one for the bad optima, we also trained a WResNet28-2 network for memorization by completely randomizing the labels of the training data.
Revised: Using the same training procedure as the one for the bad optima, we also trained a WResNet28-2 network for memorization by completely randomizing the labels of the training data.

Original: This modiﬁcation was used in [20, 21, 22] to show that neural networks have the ability to completely memorize the training data sets.
Revised: This modiﬁcation was used in [20, 21, 22] to show that neural networks have the ability to completely memorize the training data sets.

Original: We then ran the jump-and-retrain experiment for these three types of minima and visualized the results using
Revised: We then ran the jump-and-retrain experiment for these three types of minima and visualized the results using

Original: PHATE in Figure 2.
Revised: PHATE in Figure 2.

Original: This allows us to analyze and compare the regions in parameter space around the good, bad and memorization minima, which all achieve ~0 loss on their respective training sets.
Revised: This allows us to analyze and compare the regions in parameter space around the good, bad and memorization minima, which all achieve ~0 loss on their respective training sets.

Original: For the results of Section 6 we trained multiple Wide-ResNet networks of depth 28 and twice as wide, similar to the ones in the generali.
Revised: For the results of Section 6 we trained multiple Wide-ResNet networks of depth 28 and twice as wide, similar to the ones in the generali.

Original: The same initial learning rate of 0.1 and learning rate decay schedule is employed with each optimizer, data augmentation is used to assist training while no weight decay was used.
Revised: The same initial learning rate of 0.1 and learning rate decay schedule is employed with each optimizer, data augmentation is used to assist training while no weight decay was used.

Original: All experiments achieve close to 100% training accuracy and > 90% test accuracy.
Revised: All experiments achieve close to 100% training accuracy and > 90% test accuracy.

Original: A ﬁxed set of ﬁve seeds is used across optimizers.
Revised: A ﬁxed set of ﬁve seeds is used across optimizers.

Original: F Robustness
Revised: F Robustness

Original: One of the main problems encountered in previous attempts to visualize ANN loss landscapes resides in the difﬁculty of comparing them across tasks and initializations.
Revised: One of the main problems encountered in previous attempts to visualize ANN loss landscapes resides in the difﬁculty of comparing them across tasks and initializations.

Original: Here, we show that our visualization approach, applied to the results of jump-and-retrain experiments, alleviates this concern.
Revised: Here, we show that our visualization approach, applied to the results of jump-and-retrain experiments, alleviates this concern.

Original: To this end, we trained multiple WResNet28-2 networks by following the procedure described in
Revised: To this end, we trained multiple WResNet28-2 networks by following the procedure described in

Original: Sec. 2, and starting from multiple parameter initializations.
Revised: Sec. 2, and starting from multiple parameter initializations.

Original: For each initialization, we found one good (~94% test acc.) and one bad optima (~86% test acc.), giving the jump-and-retrain results shown in Fig.
Revised: For each initialization, we found one good (~94% test acc.) and one bad optima (~86% test acc.), giving the jump-and-retrain results shown in Fig.

Original: Further, we also trained a similar network on the CIFAR100 dataset, following the same training procedure, except here the data augmentation was comprised not only of random horizontal ﬂips and crops, but also of random rotations of up to 45 ◦ in each direction and random modiﬁcations to the brightness, contrast, saturation and hue of the images.
Revised: Further, we also trained a similar network on the CIFAR100 dataset, following the same training procedure, except here the data augmentation was comprised not only of random horizontal ﬂips and crops, but also of random rotations of up to 45 ◦ in each direction and random modiﬁcations to the brightness, contrast, saturation and hue of the images.

Original: Since CIFAR100 poses a signiﬁcantly harder classiﬁcation task than CIFAR10, our networks only reaches ~60% test accuracy for the bad optimum trained without data augmentation and weight decay, and ~70% test accuracy for the good optimum.
Revised: Since CIFAR100 poses a signiﬁcantly harder classiﬁcation task than CIFAR10, our networks only reaches ~60% test accuracy for the bad optimum trained without data augmentation and weight decay, and ~70% test accuracy for the good optimum.

Original: The fact that the networks only reached 70% accuracy is just a reﬂection of the complexity of the task, with better accuracy requiring a better ﬁne-tuning of the architecture and the training procedure.
Revised: The fact that the networks only reached 70% accuracy is just a reﬂection of the complexity of the task, with better accuracy requiring a better ﬁne-tuning of the architecture and the training procedure.

Original: To this end we trained multiple networks by following the procedure described in
Revised: To this end we trained multiple networks by following the procedure described in

Original: Sec.5 on CIFAR10 and CIFAR100 to ﬁnd good and bad generalisation minima.
Revised: Sec.5 on CIFAR10 and CIFAR100 to ﬁnd good and bad generalisation minima.

Original: The results of these experiments are shown in Fig.
Revised: The results of these experiments are shown in Fig.

Original: For details see Appendix E.
Revised: For details see Appendix E.

Original: It demonstrates that our approach is robust across tasks.
Revised: It demonstrates that our approach is robust across tasks.

Original: See below for a higher resolution version.
Revised: See below for a higher resolution version.

Original: The plots for different initializations of the WResNet28-2 are consistently showing the same starshaped patterns, characterizing the retraining trajectories surrounding minima that generalize well.
Revised: The plots for different initializations of the WResNet28-2 are consistently showing the same starshaped patterns, characterizing the retraining trajectories surrounding minima that generalize well.

Original: The plots for the bad minima also share a similar pattern, not only across initialization but also across tasks.
Revised: The plots for the bad minima also share a similar pattern, not only across initialization but also across tasks.

Original: Strikingly, the star-shaped pattern is even better deﬁned in Fig.
Revised: Strikingly, the star-shaped pattern is even better deﬁned in Fig.

Original: 9f, corresponding to the good optimum of WResNet28-2 trained on CIFAR100.
Revised: 9f, corresponding to the good optimum of WResNet28-2 trained on CIFAR100.

Original: This offers a key insight into the effect of regularization on the geometric characteristics of the found minima.
Revised: This offers a key insight into the effect of regularization on the geometric characteristics of the found minima.

Original: It seems that regularization techniques enhance generalization by forcing the optimizer into wider valleys.
Revised: It seems that regularization techniques enhance generalization by forcing the optimizer into wider valleys.