Original: Towards Theoretically Inspired Neural Initialization Optimization
Revised: Towards Theoretically Inspired Neural Initialization Optimization

Original: Abstract
Revised: Abstract

Original: Automated machine learning has been widely explored to reduce human efforts in designing neural architectures and looking for proper hyperparameters.
Revised: Automated machine learning has been widely explored to reduce human efforts in designing neural architectures and looking for proper hyperparameters.

Original: In the domain of neural initialization, however, similar automated techniques have rarely been studied.
Revised: In the domain of neural initialization, however, similar automated techniques have rarely been studied.

Original: Most existing initialization methods are handcrafted and highly dependent on speciﬁc architectures.
Revised: Most existing initialization methods are handcrafted and highly dependent on speciﬁc architectures.

Original: In this paper, we propose a differentiable quantity, named GradCosine, with theoretical insights to evaluate the initial state of a neural network.
Revised: In this paper, we propose a differentiable quantity, named GradCosine, with theoretical insights to evaluate the initial state of a neural network.

Original: Speciﬁcally, GradCosine is the cosine similarity of sample-wise gradients with respect to the initialized parameters.
Revised: Speciﬁcally, GradCosine is the cosine similarity of sample-wise gradients with respect to the initialized parameters.

Original: By analyzing the samplewise optimization landscape, we show that both the training and test performance of a network can be improved by maximizing GradCosine under gradient norm constraint.
Revised: By analyzing the samplewise optimization landscape, we show that both the training and test performance of a network can be improved by maximizing GradCosine under gradient norm constraint.

Original: Based on this observation, we further propose the neural initialization optimization (NIO) algorithm.
Revised: Based on this observation, we further propose the neural initialization optimization (NIO) algorithm.

Original: Generalized from the sample-wise analysis into the real batch setting, NIO is able to automatically look for a better initialization with negligible cost compared with the training time.
Revised: Generalized from the sample-wise analysis into the real batch setting, NIO is able to automatically look for a better initialization with negligible cost compared with the training time.

Original: With NIO, we improve the classiﬁcation performance of a variety of neural architectures on CIFAR10,
Revised: With NIO, we improve the classiﬁcation performance of a variety of neural architectures on CIFAR10,

Original: CIFAR-100, and ImageNet.
Revised: CIFAR-100, and ImageNet.

Original: Moreover, we ﬁnd that our method can even help to train large vision Transformer architecture without warmup.
Revised: Moreover, we ﬁnd that our method can even help to train large vision Transformer architecture without warmup.

Original:  Introduction
Revised:  Introduction

Original: For a deep neural network, architecture [19, 20, 21] and parameter initialization [18, 16] are two initial elements that largely account for the ﬁnal model performance.
Revised: For a deep neural network, architecture [19, 20, 21] and parameter initialization [18, 16] are two initial elements that largely account for the ﬁnal model performance.

Original: Lots of human efforts have been devoted to ﬁnding better answers with respect to the two aspects.
Revised: Lots of human efforts have been devoted to ﬁnding better answers with respect to the two aspects.

Original: To automatically produce better architectures, neural architecture search [62, 38, 5] has been a research focus.
Revised: To automatically produce better architectures, neural architecture search [62, 38, 5] has been a research focus.

Original: However, on the other hand, using automated techniques for parameter initialization has rarely been studied.
Revised: However, on the other hand, using automated techniques for parameter initialization has rarely been studied.

Original: Previous initialization methods are mostly handcrafted.
Revised: Previous initialization methods are mostly handcrafted.

Original: They focus on ﬁnding proper variance patterns of randomly initialized weights [18, 16, 40, 36] or rely on empirical evidence derived from certain architectures [58, 23, 15, 7].
Revised: They focus on ﬁnding proper variance patterns of randomly initialized weights [18, 16, 40, 36] or rely on empirical evidence derived from certain architectures [58, 23, 15, 7].

Original: Recently, [60, 8] propose learning-based initialization that learns to tune the norms of the initial weights so as to minimize a quantity that is intimately related to favorable training dynamics.
Revised: Recently, [60, 8] propose learning-based initialization that learns to tune the norms of the initial weights so as to minimize a quantity that is intimately related to favorable training dynamics.

Original: Despite being architecture-agnostic and gradient-based, these methods do not consider sample-wise landscape.
Revised: Despite being architecture-agnostic and gradient-based, these methods do not consider sample-wise landscape.

Original: Their derived quantities lack theoretical supports for being related to model performance.
Revised: Their derived quantities lack theoretical supports for being related to model performance.

Original: It is unclear whether optimizing these quantities can indeed lead to better training or generalization performance.
Revised: It is unclear whether optimizing these quantities can indeed lead to better training or generalization performance.

Original: In order to ﬁnd a better initialization, a theoretically sound quantity that intimately evaluates both the training and test performance should be designed.
Revised: In order to ﬁnd a better initialization, a theoretically sound quantity that intimately evaluates both the training and test performance should be designed.

Original: Finding such a quantity is a non-trivial problem and shows many challenges.
Revised: Finding such a quantity is a non-trivial problem and shows many challenges.

Original: First, the test performance is mostly decided by the converged parameters after training, while the training dynamic is more related to the parameters at initialization or during training.
Revised: First, the test performance is mostly decided by the converged parameters after training, while the training dynamic is more related to the parameters at initialization or during training.

Original: Second, to efﬁciently ﬁnd a better starting point, the quantity is expected to be differentiable to enable the optimization in the continuous parameter space.
Revised: Second, to efﬁciently ﬁnd a better starting point, the quantity is expected to be differentiable to enable the optimization in the continuous parameter space.

Original: To this end, we leverage the recent advances in optimization landscape analysis [2] to propose a novel differentiable quantity and develop a corresponding algorithm for its optimization at initialization.
Revised: To this end, we leverage the recent advances in optimization landscape analysis [2] to propose a novel differentiable quantity and develop a corresponding algorithm for its optimization at initialization.

Original: Speciﬁcally, our quantity is inspired by analyzing the optimization landscapes of individual training samples [59].
Revised: Speciﬁcally, our quantity is inspired by analyzing the optimization landscapes of individual training samples [59].

Original: Through generalizing prior theoretical results on batch-wise optimization [2] to samplewise optimization, we prove that both the network’s training and generalization error are upper bounded by a theoretical quantity that correlates with the cosine similarity of the sample-wise local optima.
Revised: Through generalizing prior theoretical results on batch-wise optimization [2] to samplewise optimization, we prove that both the network’s training and generalization error are upper bounded by a theoretical quantity that correlates with the cosine similarity of the sample-wise local optima.

Original: Moreover, this quantity also relates to the training dynamic since it reﬂects the optimization path consistency [35, 37] from the starting point.
Revised: Moreover, this quantity also relates to the training dynamic since it reﬂects the optimization path consistency [35, 37] from the starting point.

Original: Unfortunately, the sample-wise local optima are intractable.
Revised: Unfortunately, the sample-wise local optima are intractable.

Original: With the hypothesis that the sample-wise local optima can be reached by the ﬁrst-order approximation from the initial parameters, we can approximate the quantity via the sample-wise gradients at initialization.
Revised: With the hypothesis that the sample-wise local optima can be reached by the ﬁrst-order approximation from the initial parameters, we can approximate the quantity via the sample-wise gradients at initialization.

Original: Our ﬁnal result shows that, under a limited gradient norm, both the training and test performance of a network can be improved by maximizing the cosine similarity of sample-wise gradients, named GradCosine, which is differentiable and easy to implement.
Revised: Our ﬁnal result shows that, under a limited gradient norm, both the training and test performance of a network can be improved by maximizing the cosine similarity of sample-wise gradients, named GradCosine, which is differentiable and easy to implement.

Original: We then propose the Neural Initialization Optimization (NIO) algorithm based on GradCosine to ﬁnd a better initialization agnostic of architecture.
Revised: We then propose the Neural Initialization Optimization (NIO) algorithm based on GradCosine to ﬁnd a better initialization agnostic of architecture.

Original: We generalize the algorithm from the sample-wise analysis into the batch-wise setting by dividing a batch into sub-batches for friendly implementation.
Revised: We generalize the algorithm from the sample-wise analysis into the batch-wise setting by dividing a batch into sub-batches for friendly implementation.

Original: We follow [8, 60] using gradient descent to learn a set of scalar coefﬁcients of the initialized parameters.
Revised: We follow [8, 60] using gradient descent to learn a set of scalar coefﬁcients of the initialized parameters.

Original: These coefﬁcients are optimized to maximize GradCosine for better training dynamic and expected performance while constraining the gradient norm from explosion.
Revised: These coefﬁcients are optimized to maximize GradCosine for better training dynamic and expected performance while constraining the gradient norm from explosion.

Original: Experiments show that for a variety of deep architectures including ResNet [19], DenseNet [21], and
Revised: Experiments show that for a variety of deep architectures including ResNet [19], DenseNet [21], and

Original: WideResNet [56], our method achieves better classiﬁcation results on CIFAR-10/100 [27] than prior heuristic [18] and learning-based [8, 60] initialization methods.
Revised: WideResNet [56], our method achieves better classiﬁcation results on CIFAR-10/100 [27] than prior heuristic [18] and learning-based [8, 60] initialization methods.

Original: We can also initialize ResNet-50 [19] on ImageNet [9] for better performance.
Revised: We can also initialize ResNet-50 [19] on ImageNet [9] for better performance.

Original: Moreover, our method is able to help the recently proposed
Revised: Moreover, our method is able to help the recently proposed

Original: Swin-Transformer [32] achieve stable training and competitive results on ImageNet even without warmup [17], which is crucial for the successful training of Transformer architectures [31, 52].
Revised: Swin-Transformer [32] achieve stable training and competitive results on ImageNet even without warmup [17], which is crucial for the successful training of Transformer architectures [31, 52].

Original:  Related Work 2.1 Network Initialization
Revised:  Related Work 2.1 Network Initialization

Original: Existing initialization methods are designed to control the norms of network parameters via Gaussian initialization [16, 18] or orthonormal matrix initialization [40, 36] with different variance patterns.
Revised: Existing initialization methods are designed to control the norms of network parameters via Gaussian initialization [16, 18] or orthonormal matrix initialization [40, 36] with different variance patterns.

Original: These analyses are most effective for simple feed-forward networks without skip connections or normalization layers.
Revised: These analyses are most effective for simple feed-forward networks without skip connections or normalization layers.

Original: Recently, initialization techniques speciﬁed for some complex architectures are proposed.
Revised: Recently, initialization techniques speciﬁed for some complex architectures are proposed.

Original: For example, [58] studied how to initialize networks with skip connections and [23] generalized the results into Transformer architecture [49].
Revised: For example, [58] studied how to initialize networks with skip connections and [23] generalized the results into Transformer architecture [49].

Original: However, these heuristic methods are restricted to speciﬁc architectures.
Revised: However, these heuristic methods are restricted to speciﬁc architectures.

Original: Automated machine learning has achieved success in looking for hyperparameters [3, 13] and architectures [62, 38, 5, 54, 55, 22], while similar techniques for neural network initialization deserve more exploration.
Revised: Automated machine learning has achieved success in looking for hyperparameters [3, 13] and architectures [62, 38, 5, 54, 55, 22], while similar techniques for neural network initialization deserve more exploration.

Original: Current learning based initialization methods [8, 60] optimize the curvature [8] or the loss reduction of the ﬁrst stochastic step [60] using gradient descent to tune the norms of the initial parameters.
Revised: Current learning based initialization methods [8, 60] optimize the curvature [8] or the loss reduction of the ﬁrst stochastic step [60] using gradient descent to tune the norms of the initial parameters.

Original: However, these methods lack theoretical foundations to be related to the model performance.
Revised: However, these methods lack theoretical foundations to be related to the model performance.

Original: Different from these methods, our proposed GradCosine is derived from a theoretical quantity that is the upper bound of both training and generalization error.
Revised: Different from these methods, our proposed GradCosine is derived from a theoretical quantity that is the upper bound of both training and generalization error.

Original: 2.2 Evaluating Model Performance at Initialization
Revised: 2.2 Evaluating Model Performance at Initialization

Original: Evaluating the performance of a network at initialization has been an important challenge and widely applied in zero-shot neural architecture search [1, 34, 6, 42, 59, 41, 29] and pruning [50, 28, 44].
Revised: Evaluating the performance of a network at initialization has been an important challenge and widely applied in zero-shot neural architecture search [1, 34, 6, 42, 59, 41, 29] and pruning [50, 28, 44].

Original: The evaluation quantities in these studies are mainly based on the initial gradient norm [44, 41], the eigenvalues of neural tangent kernel [6, 41], and the Fisher information matrix [47, 48, 45].
Revised: The evaluation quantities in these studies are mainly based on the initial gradient norm [44, 41], the eigenvalues of neural tangent kernel [6, 41], and the Fisher information matrix [47, 48, 45].

Original: However, these quantities cannot reﬂect optimization landscape, which is crucial for training dynamic and generalization [4, 12, 14, 43, 30, 2].
Revised: However, these quantities cannot reﬂect optimization landscape, which is crucial for training dynamic and generalization [4, 12, 14, 43, 30, 2].

Original: [2] provided theoretical evidence that for a sufﬁciently large neighborhood of a random initialization, the optimization landscape is nearly convex and semi-
Revised: [2] provided theoretical evidence that for a sufﬁciently large neighborhood of a random initialization, the optimization landscape is nearly convex and semi-

Original: 𝜽 𝟏∗ 𝒍(𝒇 𝜽 𝒙 𝟏 ,𝒚 𝟏 ) 𝒍(𝒇 𝜽 𝒙 𝟐 ,𝒚 𝟐 ) 𝜽 𝟐∗ 𝜽 𝟐∗
Revised: 𝜽 𝟏∗ 𝒍(𝒇 𝜽 𝒙 𝟏 ,𝒚 𝟏 ) 𝒍(𝒇 𝜽 𝒙 𝟐 ,𝒚 𝟐 ) 𝜽 𝟐∗ 𝜽 𝟐∗

Original: 𝜽
Revised: 𝜽

Original: 𝟏∗ 𝒍 𝒍 𝒍(𝒇 𝜽 𝒙 𝟐 ,𝒚 𝟐 )
Revised: 𝟏∗ 𝒍 𝒍 𝒍(𝒇 𝜽 𝒙 𝟐 ,𝒚 𝟐 )

Original: 𝒍(𝒇 𝜽 𝒙 𝟏 ,𝒚 𝟏 ) 𝜽 𝜽 𝜽 𝒍
Revised: 𝒍(𝒇 𝜽 𝒙 𝟏 ,𝒚 𝟏 ) 𝜽 𝜽 𝜽 𝒍

Original: 𝒍(𝒇 𝜽 𝒙 𝟏 ,𝒚 𝟏 ) 𝒍(𝒇 𝜽 𝒙 𝟐 ,𝒚 𝟐 ) 𝜽 𝟏∗ 𝜽 𝟐∗ 𝜽 ∗ |𝜽 𝟏∗ −𝜽 𝟐∗
Revised: 𝒍(𝒇 𝜽 𝒙 𝟏 ,𝒚 𝟏 ) 𝒍(𝒇 𝜽 𝒙 𝟐 ,𝒚 𝟐 ) 𝜽 𝟏∗ 𝜽 𝟐∗ 𝜽 ∗ |𝜽 𝟏∗ −𝜽 𝟐∗

Original: | 𝜽 𝒍 𝒍(𝒇 𝜽 𝒙 𝟏 ,𝒚 𝟏 ) 𝒍(𝒇 𝜽 𝒙 𝟐 ,𝒚 𝟐 ) 𝜽 𝟏∗ 𝜽 𝟐∗ 𝜽 ∗ |𝜽 𝟏∗ −𝜽 𝟐∗ | (𝒂) (𝒃) (𝒄) (𝒅) 𝜽𝜽𝐜𝐨𝐬(𝜽 𝟏∗ −𝜽 𝟎 ,𝜽 𝟐∗ −𝜽 𝟎 ) 𝐜𝐨𝐬(𝜽
Revised: | 𝜽 𝒍 𝒍(𝒇 𝜽 𝒙 𝟏 ,𝒚 𝟏 ) 𝒍(𝒇 𝜽 𝒙 𝟐 ,𝒚 𝟐 ) 𝜽 𝟏∗ 𝜽 𝟐∗ 𝜽 ∗ |𝜽 𝟏∗ −𝜽 𝟐∗ | (𝒂) (𝒃) (𝒄) (𝒅) 𝜽𝜽𝐜𝐨𝐬(𝜽 𝟏∗ −𝜽 𝟎 ,𝜽 𝟐∗ −𝜽 𝟎 ) 𝐜𝐨𝐬(𝜽

Original: 𝟏∗ −𝜽 𝟎 ,𝜽 𝟐∗ −𝜽 𝟎 ) smooth.
Revised: 𝟏∗ −𝜽 𝟎 ,𝜽 𝟐∗ −𝜽 𝟎 ) smooth.

Original: Based on the result, [59] proposed to use the density of sample-wise local optima to evaluate and rank neural architectures.
Revised: Based on the result, [59] proposed to use the density of sample-wise local optima to evaluate and rank neural architectures.

Original: Our study also performs sample-wise landscape analysis, but differs from [59] in that our proposed quantity is differentiable and reﬂects optimization path consistency, while the quantity in [59] is non-differentiable so cannot serve for initialization optimization.
Revised: Our study also performs sample-wise landscape analysis, but differs from [59] in that our proposed quantity is differentiable and reﬂects optimization path consistency, while the quantity in [59] is non-differentiable so cannot serve for initialization optimization.

Original: We make a comparison in more details in Appendix C.
Revised: We make a comparison in more details in Appendix C.

Original:  Theoretical Foundations 3.1 Sample-wise Optimization Landscape Analysis
Revised:  Theoretical Foundations 3.1 Sample-wise Optimization Landscape Analysis

Original: Conventional optimization landscape analyses [30, 2] mainly focus on the objective across a minibatch of samples and miss potential evidence hidden in the optimization landscapes of individual samples.
Revised: Conventional optimization landscape analyses [30, 2] mainly focus on the objective across a minibatch of samples and miss potential evidence hidden in the optimization landscapes of individual samples.

Original: As recently pointed out in [59], by decomposing a mini-batch objective into the summation of sample-wise objectives across individual samples in the mini-batch, they ﬁnd that the network with denser sample-wise local optima tends to reach a better local optima for the mini-batch objective, as shown in the Figure (1) (a)-(b).
Revised: As recently pointed out in [59], by decomposing a mini-batch objective into the summation of sample-wise objectives across individual samples in the mini-batch, they ﬁnd that the network with denser sample-wise local optima tends to reach a better local optima for the mini-batch objective, as shown in the Figure (1) (a)-(b).

Original: Based on this insight, they propose to use sample-wise local optima density to judge model performance at initialization.
Revised: Based on this insight, they propose to use sample-wise local optima density to judge model performance at initialization.

Original: Density of sample-wise local optima.
Revised: Density of sample-wise local optima.

Original: For a batch of training samples S
Revised: For a batch of training samples S

Original: = { ( x i , y i ) } i ∈ [ n ] , a loss function l (ˆ y i , y i ) , and a network f θ ( · ) parameterized by θ ∈ R m , the sample-wise local optima density Ψ S,l ( f θ 0 ( · )) is measured by the averaged Manhattan distance between the pair-wise local optima { θ ∗ i } i ∈ [ n ] of all n samples near the random initialization θ 0 [59], i.e.,
Revised: = { ( x i , y i ) } i ∈ [ n ] , a loss function l (ˆ y i , y i ) , and a network f θ ( · ) parameterized by θ ∈ R m , the sample-wise local optima density Ψ S,l ( f θ 0 ( · )) is measured by the averaged Manhattan distance between the pair-wise local optima { θ ∗ i } i ∈ [ n ] of all n samples near the random initialization θ 0 [59], i.e.,

Original: [Equation]
Revised: [Equation]

Original: where H is the smoothness upper bound: ∀ k ∈ [ m ] , i ∈ [ n ] , [ ∇ 2 l ( f θ ( x i ) , y i )]
Revised: where H is the smoothness upper bound: ∀ k ∈ [ m ] , i ∈ [ n ] , [ ∇ 2 l ( f θ ( x i ) , y i )]

Original: k,k ≤ H and H always exists with the smoothness assumption that for a neighborhood Γ θ 0 of a random initialization θ 0 , the sample-wise optimization landscapes are nearly convex and semi-smooth [2].
Revised: k,k ≤ H and H always exists with the smoothness assumption that for a neighborhood Γ θ 0 of a random initialization θ 0 , the sample-wise optimization landscapes are nearly convex and semi-smooth [2].

Original: Based on this assumption, the training error of the network can be upper bounded by n 3 2 Ψ 2 S,l ( f θ 0 ( · )) .
Revised: Based on this assumption, the training error of the network can be upper bounded by n 3 2 Ψ 2 S,l ( f θ 0 ( · )) .

Original: Moreover, with probability 1 − δ , the generalization error measured by population loss E ( x u ,y u ) ∼D
Revised: Moreover, with probability 1 − δ , the generalization error measured by population loss E ( x u ,y u ) ∼D

Original: [ l ( f ∗ θ ( x u ) , y u )] is upper bounded by n 3 2 Ψ 2 S,l ( f θ 0 ( · ))
Revised: [ l ( f ∗ θ ( x u ) , y u )] is upper bounded by n 3 2 Ψ 2 S,l ( f θ 0 ( · ))

Original: + σ √ nδ , where
Revised: + σ √ nδ , where

Original: D = { ( x u , y u ) } u ∈ [ U ] is the underlying data distribution of test samples and σ 2 is the upper bound of V ar ( x u ,y u )
Revised: D = { ( x u , y u ) } u ∈ [ U ] is the underlying data distribution of test samples and σ 2 is the upper bound of V ar ( x u ,y u )

Original: ∼D [ || θ ∗ − θ ∗ u || 21 ] [59].
Revised: ∼D [ || θ ∗ − θ ∗ u || 21 ] [59].

Original: Although the sample-wise local optima density is theoretically related to both the network training and generalization error, we point out that it may not be a suitable quantity to evaluate network initialization due to the following reasons.
Revised: Although the sample-wise local optima density is theoretically related to both the network training and generalization error, we point out that it may not be a suitable quantity to evaluate network initialization due to the following reasons.

Original: First, it ignores the optimization path consistency from initialization θ 0 to each sample-wise local optimum θ ∗ i .
Revised: First, it ignores the optimization path consistency from initialization θ 0 to each sample-wise local optimum θ ∗ i .

Original: As shown in Figure 1 (c)-(d), while both (c) and (d) have the same local optima density, in Figure 1 (d) the optimization path from the initialization to the local optima of the two samples are more consistent.
Revised: As shown in Figure 1 (c)-(d), while both (c) and (d) have the same local optima density, in Figure 1 (d) the optimization path from the initialization to the local optima of the two samples are more consistent.

Original: Training networks on samples with more consistent optimization paths using batch gradient descent naturally corresponds to a better training dynamic that enjoys faster and more stable convergence [39, 37].
Revised: Training networks on samples with more consistent optimization paths using batch gradient descent naturally corresponds to a better training dynamic that enjoys faster and more stable convergence [39, 37].

Original: Second, the sample-wise local optima in Eq.
Revised: Second, the sample-wise local optima in Eq.

Original: (1) are intractable.
Revised: (1) are intractable.

Original: It can be approximated by measuring the consistency of sample-wise gradient signs [59].
Revised: It can be approximated by measuring the consistency of sample-wise gradient signs [59].

Original: But it is non-differentiable and cannot serve for initialization optimization.
Revised: But it is non-differentiable and cannot serve for initialization optimization.

Original: Based on this observation, we aim to derive a new quantity that directly reﬂects optimization path consistency and is a differentiable function of the initialization θ 0 .
Revised: Based on this observation, we aim to derive a new quantity that directly reﬂects optimization path consistency and is a differentiable function of the initialization θ 0 .

Original: Our proposed quantity is based on the cosine similarity of the paths from initialization to sample-wise local optima.
Revised: Our proposed quantity is based on the cosine similarity of the paths from initialization to sample-wise local optima.

Original: Cosine similarity of sample-wise local optima.
Revised: Cosine similarity of sample-wise local optima.

Original: Concretely, our quantity can be formulated as:
Revised: Concretely, our quantity can be formulated as:

Original: [Equation]
Revised: [Equation]

Original: where α and β are the maximal and minimal (cid:96) 2 -norms of the sample-wise optimization paths, i.e., α = max
Revised: where α and β are the maximal and minimal (cid:96) 2 -norms of the sample-wise optimization paths, i.e., α = max

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: ||
Revised: ||

Original: [Equation]
Revised: [Equation]

Original: || 2 ) , ∀ i ∈ [1 , n ] , and cos ∠ ( θ ∗ i − θ 0 , θ ∗ j − θ 0 ) refers to the cosine similarity of the paths from initialization θ 0 to sample-wise local optima, θ ∗ i and θ ∗ j .
Revised: || 2 ) , ∀ i ∈ [1 , n ] , and cos ∠ ( θ ∗ i − θ 0 , θ ∗ j − θ 0 ) refers to the cosine similarity of the paths from initialization θ 0 to sample-wise local optima, θ ∗ i and θ ∗ j .

Original: The cosine term in Eq.
Revised: The cosine term in Eq.

Original: (2) reﬂects the optimization path consistency.
Revised: (2) reﬂects the optimization path consistency.

Original: Together with the distance term αβ , it is also able to measure the density of sample-wise local optima.
Revised: Together with the distance term αβ , it is also able to measure the density of sample-wise local optima.

Original: In the ideal case, when all the local optima are located at the same point, Θ S,l ( f θ 0 ( · )) =
Revised: In the ideal case, when all the local optima are located at the same point, Θ S,l ( f θ 0 ( · )) =

Original: Ψ S,l ( f θ 0 ( · ))
Revised: Ψ S,l ( f θ 0 ( · ))

Original: = 0 .
Revised: = 0 .

Original: Hence, compared with Ψ , our Θ is more suitable for evaluating the initialization quality.
Revised: Hence, compared with Ψ , our Θ is more suitable for evaluating the initialization quality.

Original: 3.2 Main Results
Revised: 3.2 Main Results

Original: In this subsection, we theoretically illustrate how minimizing the quantity Θ in Eq.
Revised: In this subsection, we theoretically illustrate how minimizing the quantity Θ in Eq.

Original: (2) corresponds to better training and generalization performance.
Revised: (2) corresponds to better training and generalization performance.

Original: Similar to [59], our derivations are also based on the evidence that there exists a neighborhood for a random initialization such that the sample-wise optimization landscapes are nearly convex and semi-smooth [2].
Revised: Similar to [59], our derivations are also based on the evidence that there exists a neighborhood for a random initialization such that the sample-wise optimization landscapes are nearly convex and semi-smooth [2].

Original: Lemma 1.
Revised: Lemma 1.

Original: There exists no saddle point in a sample-wise optimization landscape and every local optimum is a global optimum [59].
Revised: There exists no saddle point in a sample-wise optimization landscape and every local optimum is a global optimum [59].

Original: Based on Lemma 1, we can draw a relation between the training error and Θ S,l ( f θ 0 ( · )) .
Revised: Based on Lemma 1, we can draw a relation between the training error and Θ S,l ( f θ 0 ( · )) .

Original: Moreover, we show that the proposed quantity is also related to generalization performance as the upper bound of population error.
Revised: Moreover, we show that the proposed quantity is also related to generalization performance as the upper bound of population error.

Original: We present the following two theoretical results.
Revised: We present the following two theoretical results.

Original: Theorem 2.
Revised: Theorem 2.

Original: The training loss L = 1 n (cid:80) i
Revised: The training loss L = 1 n (cid:80) i

Original: l ( f θ ∗ ( x i ) , y i ) of a trained network f θ ∗ on a dataset
Revised: l ( f θ ∗ ( x i ) , y i ) of a trained network f θ ∗ on a dataset

Original: S = { ( x i , y i ) } i ∈ [ n ] is upper bounded by Θ S,l ( f θ 0 ( · )) , and the bound is tight when Θ S,l ( f θ 0 ( · ))
Revised: S = { ( x i , y i ) } i ∈ [ n ] is upper bounded by Θ S,l ( f θ 0 ( · )) , and the bound is tight when Θ S,l ( f θ 0 ( · ))

Original: = 0 .
Revised: = 0 .

Original: Theorem 3.
Revised: Theorem 3.

Original: Suppose that σ 2 is the upper bound of V ar ( x u ,y u )
Revised: Suppose that σ 2 is the upper bound of V ar ( x u ,y u )

Original: ∼D [ || θ ∗ − θ ∗ u || 22 ] , where θ ∗ u is the local optimum in the convex neighborhood of θ 0 for test sample ( x u , y u ) .
Revised: ∼D [ || θ ∗ − θ ∗ u || 22 ] , where θ ∗ u is the local optimum in the convex neighborhood of θ 0 for test sample ( x u , y u ) .

Original: With probability 1 − δ , the population loss E ( x u ,y u ) ∼D
Revised: With probability 1 − δ , the population loss E ( x u ,y u ) ∼D

Original: [ l ( f θ ∗ ( x u ) , y u )] is upper bounded by Θ S,l ( f θ 0 ( · ))
Revised: [ l ( f θ ∗ ( x u ) , y u )] is upper bounded by Θ S,l ( f θ 0 ( · ))

Original: + σ √ nδ .
Revised: + σ √ nδ .

Original: We provide proofs of these two theorems in Appendix A. Combining both theorems, we can conclude that Θ S,l ( f θ 0 ( · )) is the upper bound of both training and generalization errors of the network f θ ∗ .
Revised: We provide proofs of these two theorems in Appendix A. Combining both theorems, we can conclude that Θ S,l ( f θ 0 ( · )) is the upper bound of both training and generalization errors of the network f θ ∗ .

Original: Therefore, minimizing Θ theoretically helps to improve the model performance.
Revised: Therefore, minimizing Θ theoretically helps to improve the model performance.

Original: Albeit theoretically sound, Θ S,l ( f θ 0 ( · )) requires sample-wise optima θ ∗ i which are intractable at initialization.
Revised: Albeit theoretically sound, Θ S,l ( f θ 0 ( · )) requires sample-wise optima θ ∗ i which are intractable at initialization.

Original: To this end, we will show how to develop a differentiable and tractable objective based on Eq.
Revised: To this end, we will show how to develop a differentiable and tractable objective based on Eq.

Original: (2) in Section 4, and introduce the initialization optimization algorithm in Section 5.
Revised: (2) in Section 4, and introduce the initialization optimization algorithm in Section 5.

Original:  GradCosine 4.1 First-Order Approximation of Sample-Wise Optimization
Revised:  GradCosine 4.1 First-Order Approximation of Sample-Wise Optimization

Original: Since we are dealing with sample-wise optimization, it is reasonable to calculate each sample-wise optimum by the ﬁrst-order approximation.
Revised: Since we are dealing with sample-wise optimization, it is reasonable to calculate each sample-wise optimum by the ﬁrst-order approximation.

Original: We hypothesize that each sample-wise optimum can be reached via only one-step gradient descent from the initialized parameters.
Revised: We hypothesize that each sample-wise optimum can be reached via only one-step gradient descent from the initialized parameters.

Original: Its rationale lies in that it is very easy for a deep neural network to learn the optimum for only one training sample with gradient descent.
Revised: Its rationale lies in that it is very easy for a deep neural network to learn the optimum for only one training sample with gradient descent.

Original: Based on this hypothesis, we can approximate each local optimum as:
Revised: Based on this hypothesis, we can approximate each local optimum as:

Original: [Equation]
Revised: [Equation]

Original: where g i = ∇ θ l ( f θ ( x i ) , y i ) | θ 0 is the sample-wise gradient at initialization, and η is the learning rate.
Revised: where g i = ∇ θ l ( f θ ( x i ) , y i ) | θ 0 is the sample-wise gradient at initialization, and η is the learning rate.

Original: With Eq.
Revised: With Eq.

Original: (3), the upper bound quantity in Eq.
Revised: (3), the upper bound quantity in Eq.

Original: (2) can be simpliﬁed as:
Revised: (2) can be simpliﬁed as:

Original: [Equation]
Revised: [Equation]

Original: where g max and g min are the maximal and minimal sample-wise gradient norms at initialization, i.e., g max = max( || g
Revised: where g max and g min are the maximal and minimal sample-wise gradient norms at initialization, i.e., g max = max( || g

Original: i || 2 ) , g min = min( || g
Revised: i || 2 ) , g min = min( || g

Original: i || 2 ) , ∀ i ∈ [1 , n ] .
Revised: i || 2 ) , ∀ i ∈ [1 , n ] .

Original: 4.2 Guidance for Initialization
Revised: 4.2 Guidance for Initialization

Original: Suppose that the sample-wise gradient at initialization is upper bounded by a constant γ , i.e., g max ≤ γ , and then we have:
Revised: Suppose that the sample-wise gradient at initialization is upper bounded by a constant γ , i.e., g max ≤ γ , and then we have:

Original: [Equation]
Revised: [Equation]

Original: From Eq.
Revised: From Eq.

Original: (2) to Eq.
Revised: (2) to Eq.

Original: (5), we have almost converted the upper bound into a quantity that can be easily calculated by the sample-wise gradient at initialization, except for the term g min and the constraint g max ≤ γ .
Revised: (5), we have almost converted the upper bound into a quantity that can be easily calculated by the sample-wise gradient at initialization, except for the term g min and the constraint g max ≤ γ .

Original: We will explicitly add corresponding optimization objective and constraint in our neural initialization optimization algorithm in Section 5 to ensure a small γ/g min .
Revised: We will explicitly add corresponding optimization objective and constraint in our neural initialization optimization algorithm in Section 5 to ensure a small γ/g min .

Original: From Eq.
Revised: From Eq.

Original: (5), we can draw some useful guidance for a better initialization:
Revised: (5), we can draw some useful guidance for a better initialization:

Original: (1) The sample-wise gradient norms of the initialized parameters should be large and close to the maximal value bounded by γ to induce a small γ/g min ; (2) The cosine similarity of the sample-wise gradients should be as large as possible.
Revised: (1) The sample-wise gradient norms of the initialized parameters should be large and close to the maximal value bounded by γ to induce a small γ/g min ; (2) The cosine similarity of the sample-wise gradients should be as large as possible.

Original: Relations to Favorable Training Dynamic.
Revised: Relations to Favorable Training Dynamic.

Original: There are signiﬁcant evidences in existing studies that support the two rules for better training dynamic.
Revised: There are signiﬁcant evidences in existing studies that support the two rules for better training dynamic.

Original: The ﬁrst rule is intimately related to the neural tangent kernel (NTK) [25] that has recently been shown to determine the early learning behavior [51, 41, 6, 42].
Revised: The ﬁrst rule is intimately related to the neural tangent kernel (NTK) [25] that has recently been shown to determine the early learning behavior [51, 41, 6, 42].

Original: Speciﬁcally, [41] ﬁnds that the training dynamic of a neural network can be characterized by the trace norm of NTK at initialization and further approximates it via the initial gradient norm as a gauge to search for neural architectures.
Revised: Speciﬁcally, [41] ﬁnds that the training dynamic of a neural network can be characterized by the trace norm of NTK at initialization and further approximates it via the initial gradient norm as a gauge to search for neural architectures.

Original: [60] initializes a network by maximizing the loss reduction of the ﬁrst gradient descent step, which also corresponds to gradient norm in the ﬁrst-order approximation of the loss function.
Revised: [60] initializes a network by maximizing the loss reduction of the ﬁrst gradient descent step, which also corresponds to gradient norm in the ﬁrst-order approximation of the loss function.

Original: The two rules together prefer a model whose sample-wise gradients have close (cid:96) 2 and cosine distances.
Revised: The two rules together prefer a model whose sample-wise gradients have close (cid:96) 2 and cosine distances.

Original: It is in line with prior observations that the initial gradient variance should be small to enable a large learning rate [31, 61].
Revised: It is in line with prior observations that the initial gradient variance should be small to enable a large learning rate [31, 61].

Original: However, gradient variance is not a proper optimization objective due to its high sensitivity to gradient norm.
Revised: However, gradient variance is not a proper optimization objective due to its high sensitivity to gradient norm.

Original: While the ﬁrst rule of improving the initial gradient norm has been suggested for better training dynamic in prior studies [60, 41], the second rule is completely new.
Revised: While the ﬁrst rule of improving the initial gradient norm has been suggested for better training dynamic in prior studies [60, 41], the second rule is completely new.

Original: We refer this novel quantity, i.e., cos ∠ ( g i , g j ) , for evaluating the network initialization as GradCosine.
Revised: We refer this novel quantity, i.e., cos ∠ ( g i , g j ) , for evaluating the network initialization as GradCosine.

Original: Intuitively, in the ideal case where the initial gradients of all samples are identical, we have the smallest upper bound of training and generalization error.
Revised: Intuitively, in the ideal case where the initial gradients of all samples are identical, we have the smallest upper bound of training and generalization error.

Original: Moreover, both gradient norm and GradCosine are differentiable and thus enable optimization to ﬁnd a better initialization.
Revised: Moreover, both gradient norm and GradCosine are differentiable and thus enable optimization to ﬁnd a better initialization.

Original: We provide the calculation of gradient norm (GN) and GradCosine (GC) for an initialized network in Algorithm 1.
Revised: We provide the calculation of gradient norm (GN) and GradCosine (GC) for an initialized network in Algorithm 1.

Original: Require: network parameters θ 0 , a batch of samples S
Revised: Require: network parameters θ 0 , a batch of samples S

Original: = { ( x i , y i ) } i ∈
Revised: = { ( x i , y i ) } i ∈

Original: [ B ] , the number of batches D , and overlap ratio r .
Revised: [ B ] , the number of batches D , and overlap ratio r .

Original: [Figure]
Revised: [Figure]

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: experiment is conducted
Revised: experiment is conducted

Original:  Neural Initialization Optimization
Revised:  Neural Initialization Optimization

Original: Based on the GradCosine, we propose our Neural Initialization Optimization (NIO) algorithm.
Revised: Based on the GradCosine, we propose our Neural Initialization Optimization (NIO) algorithm.

Original: Our algorithm follows [60, 8] to rectify the variance of the initialized parameters in a network via gradient descent.
Revised: Our algorithm follows [60, 8] to rectify the variance of the initialized parameters in a network via gradient descent.

Original: Concretely, we introduce a set of learnable scalar coefﬁcients denoted as
Revised: Concretely, we introduce a set of learnable scalar coefﬁcients denoted as

Original: M = { ω 1 , · · · , ω m } .
Revised: M = { ω 1 , · · · , ω m } .

Original: The initial parameters θ 0 =
Revised: The initial parameters θ 0 =

Original: { W 1 , · · · , W m } rectiﬁed by these coefﬁcients can be formulated as θ M = { ω 1 W 1 , · · · , ω m W m } .
Revised: { W 1 , · · · , W m } rectiﬁed by these coefﬁcients can be formulated as θ M = { ω 1 W 1 , · · · , ω m W m } .

Original: The NIO algorithm mainly solves the following constrained optimization problem:
Revised: The NIO algorithm mainly solves the following constrained optimization problem:

Original: [Equation]
Revised: [Equation]

Original: where GC and GN refer to GradCosine and gradient norm, respectively, as calculated in Algorithm 1, and B is the batchsize of the batch data S .
Revised: where GC and GN refer to GradCosine and gradient norm, respectively, as calculated in Algorithm 1, and B is the batchsize of the batch data S .

Original: With the maximal sample-wise gradient norm bounded by γ , we maximize both GradCosine and gradient norm.
Revised: With the maximal sample-wise gradient norm bounded by γ , we maximize both GradCosine and gradient norm.

Original: Although the problem in Eq.
Revised: Although the problem in Eq.

Original: (6) is differentiable and tractable, inferring GradCosine and gradient norm in the sample-wise manner can be very time-consuming especially for networks trained on large datasets such as ImageNet.
Revised: (6) is differentiable and tractable, inferring GradCosine and gradient norm in the sample-wise manner can be very time-consuming especially for networks trained on large datasets such as ImageNet.

Original: Moreover, using gradient-based methods to solve Eq.
Revised: Moreover, using gradient-based methods to solve Eq.

Original: (6) requires deriving the second-order gradient with respect to the network parameters, which leads to unbearable memory consumption in the sample-wise manner.
Revised: (6) requires deriving the second-order gradient with respect to the network parameters, which leads to unbearable memory consumption in the sample-wise manner.

Original: To this end, we introduce the batch-wise relaxation.
Revised: To this end, we introduce the batch-wise relaxation.

Original: 5.1 Generalizing to Batch Settings
Revised: 5.1 Generalizing to Batch Settings

Original: To reduce the time and space complexity of gradient calculation, we further generalize GC and
Revised: To reduce the time and space complexity of gradient calculation, we further generalize GC and

Original: GN from the sample-wise manner into a more efﬁcient batch setting.
Revised: GN from the sample-wise manner into a more efﬁcient batch setting.

Original: Speciﬁcally, for a mini-batch of training samples S
Revised: Speciﬁcally, for a mini-batch of training samples S

Original: = { ( x i , y i ) } i ∈
Revised: = { ( x i , y i ) } i ∈

Original: [ B ] , we split it into D sub-batches S 1 , ..., S d , ..., S D with an overlap ratio of r that is used to relieve the violation of the sample-wise setting.
Revised: [ B ] , we split it into D sub-batches S 1 , ..., S d , ..., S D with an overlap ratio of r that is used to relieve the violation of the sample-wise setting.

Original: The number of samples in each sub-batch is N = (cid:100) BD − r (cid:101) , and each sub-batch can be denoted as:
Revised: The number of samples in each sub-batch is N = (cid:100) BD − r (cid:101) , and each sub-batch can be denoted as:

Original: [Equation]
Revised: [Equation]

Original: ) where d = 1 , ..., D .
Revised: ) where d = 1 , ..., D .

Original: The sample-wise input can be regarded as the special case of the batch split when setting D = B, r = 0 .
Revised: The sample-wise input can be regarded as the special case of the batch split when setting D = B, r = 0 .

Original: We further illustrate some split examples in Figure 2.
Revised: We further illustrate some split examples in Figure 2.

Original: We empirically ﬁnd that making the sub-batches overlapped with each other stabilize the optimization and results in better initialization.
Revised: We empirically ﬁnd that making the sub-batches overlapped with each other stabilize the optimization and results in better initialization.

Original: As shown in Figure 2, the test accuracies of the models trained using our initialization with different batch settings further conﬁrm this practice.
Revised: As shown in Figure 2, the test accuracies of the models trained using our initialization with different batch settings further conﬁrm this practice.

Original: The GradCosine and gradient norm under the batch setting are illustrated in Algorithm 2.
Revised: The GradCosine and gradient norm under the batch setting are illustrated in Algorithm 2.

Original: The main difference is that, in the batch setting, the gradient is computed as the average of gradients in a sub-batch.
Revised: The main difference is that, in the batch setting, the gradient is computed as the average of gradients in a sub-batch.

Original: Our optimization problem in the batch setting can be formulated as:
Revised: Our optimization problem in the batch setting can be formulated as:

Original: [Equation]
Revised: [Equation]

Original: where B-GC and B-GN refer to the batch-wise GradCosine and gradient norm as calculated in
Revised: where B-GC and B-GN refer to the batch-wise GradCosine and gradient norm as calculated in

Original: Algorithm 2.
Revised: Algorithm 2.

Original: 5.2 Main Algorithm
Revised: 5.2 Main Algorithm

Original: The ﬁnal neural initialization optimization is illustrated in Algorithm 3.
Revised: The ﬁnal neural initialization optimization is illustrated in Algorithm 3.

Original: It rectiﬁes the initialized parameters by gradient descent to solve the constrained optimization problem in Eq.
Revised: It rectiﬁes the initialized parameters by gradient descent to solve the constrained optimization problem in Eq.

Original: Concretely, we iterate for T iterations.
Revised: Concretely, we iterate for T iterations.

Original: At each iteration, a random batch data S t is sampled from the training set, and divided into sub-batches according to D and r .
Revised: At each iteration, a random batch data S t is sampled from the training set, and divided into sub-batches according to D and r .

Original: We calculate B-GN and B-GC by
Revised: We calculate B-GN and B-GC by

Original: Algorithm 2.
Revised: Algorithm 2.

Original: If the maximal gradient norm is larger than the predeﬁned upper bound γ , we minimize the averaged batch gradient norm to avoid gradient explosion at initialization.
Revised: If the maximal gradient norm is larger than the predeﬁned upper bound γ , we minimize the averaged batch gradient norm to avoid gradient explosion at initialization.

Original: When the constraint is satisﬁed, we maximize the GradCosine (B-GC) and gradient norm (B-GN) objectives simultaneously, in order to minimize the upper bound of the quantity Θ in Eq.
Revised: When the constraint is satisﬁed, we maximize the GradCosine (B-GC) and gradient norm (B-GN) objectives simultaneously, in order to minimize the upper bound of the quantity Θ in Eq.

Original: (5), which intimately corresponds to not only a lower training and generalization error, but also a better training dynamic.
Revised: (5), which intimately corresponds to not only a lower training and generalization error, but also a better training dynamic.

Original:  Experiments 6.1 Initialization for CNN
Revised:  Experiments 6.1 Initialization for CNN

Original: Dataset and Architectures.
Revised: Dataset and Architectures.

Original: We validate our method on three widely used datasets including
Revised: We validate our method on three widely used datasets including

Original: CIFAR10/100 [27] and ImageNet [9].
Revised: CIFAR10/100 [27] and ImageNet [9].

Original: We select three kinds of convolution neural networks including ResNet [19], DenseNet [21], and WideResNet [56].
Revised: We select three kinds of convolution neural networks including ResNet [19], DenseNet [21], and WideResNet [56].

Original: On CIFAR10/100, we adopt ResNet110,
Revised: On CIFAR10/100, we adopt ResNet110,

Original: DenseNet100, and the 28-layer Wide ResNet with Widen Factor 10 (WRN-28-10) as three main architectures for evaluation.
Revised: DenseNet100, and the 28-layer Wide ResNet with Widen Factor 10 (WRN-28-10) as three main architectures for evaluation.

Original: To further show that our initialization method helps training dynamic for better training stability, we also conduct experiments on the same architectures but remove the batch normalization (BN) [24] layers.
Revised: To further show that our initialization method helps training dynamic for better training stability, we also conduct experiments on the same architectures but remove the batch normalization (BN) [24] layers.

Original: Moreover, to illustrate that our method can be extensible to large-scale benchmarks, we test our proposed NIO on ImageNet using ResNet-50.
Revised: Moreover, to illustrate that our method can be extensible to large-scale benchmarks, we test our proposed NIO on ImageNet using ResNet-50.

Original: The detailed settings for different architectures and datasets are described in Appendix B.
Revised: The detailed settings for different architectures and datasets are described in Appendix B.

Original: Experiment Results on CIFAR-10/100.
Revised: Experiment Results on CIFAR-10/100.

Original: We select four different initialization methods for comparison: (1) Kaiming Initialization [18]; (2) First train the network for one epoch with a linear warmup learning rate, denoted as Warmup; (3) MetaInit [8]; and (4) GradInit [60].
Revised: We select four different initialization methods for comparison: (1) Kaiming Initialization [18]; (2) First train the network for one epoch with a linear warmup learning rate, denoted as Warmup; (3) MetaInit [8]; and (4) GradInit [60].

Original: (3) and (4) are learning-based initialization as ours.
Revised: (3) and (4) are learning-based initialization as ours.

Original: We re-implement their methods using the code provided in [60].
Revised: We re-implement their methods using the code provided in [60].

Original: For MetaInit, GradInit, and our proposed NIO, the initialization is rectiﬁed based on the Kaiming initialized parameters.
Revised: For MetaInit, GradInit, and our proposed NIO, the initialization is rectiﬁed based on the Kaiming initialized parameters.

Original: After initialization, we train these models for 500 epochs with the same
Revised: After initialization, we train these models for 500 epochs with the same

Original: [Table] [Table]
Revised: [Table] [Table]

Original: [Figure]
Revised: [Figure]

Original: training setting.
Revised: training setting.

Original: Each model is trained four times with different seeds.
Revised: Each model is trained four times with different seeds.

Original: We report the average and standard deviation numbers of the accuracies on the test set.
Revised: We report the average and standard deviation numbers of the accuracies on the test set.

Original: As shown in Table 1 and 2, compared with the Kaiming initialization, we observe that MetaInit and
Revised: As shown in Table 1 and 2, compared with the Kaiming initialization, we observe that MetaInit and

Original: GradInit do not achieve improvement for some cases.
Revised: GradInit do not achieve improvement for some cases.

Original: As a comparison, our proposed NIO consistently improves upon the Kaiming initialization on both CIFAR-10 and CIFAR-100 with and without BN.
Revised: As a comparison, our proposed NIO consistently improves upon the Kaiming initialization on both CIFAR-10 and CIFAR-100 with and without BN.

Original: Especially for ResNet-110-BN and DenseNet-100-noBN, the performance improvements upon the
Revised: Especially for ResNet-110-BN and DenseNet-100-noBN, the performance improvements upon the

Original: Kaiming initialization are more than 0.5% on CIFAR-10 and 1.0% on CIFAR-100.
Revised: Kaiming initialization are more than 0.5% on CIFAR-10 and 1.0% on CIFAR-100.

Original: Compared with
Revised: Compared with

Original: GradInit, we also achieve better results by more than 1.0% for ResNet-110-noBN and DenseNet-110noBN on CIFAR-100.
Revised: GradInit, we also achieve better results by more than 1.0% for ResNet-110-noBN and DenseNet-110noBN on CIFAR-100.

Original: We also make some empirical comparisons between the Kaiming initialization and our NIO with
Revised: We also make some empirical comparisons between the Kaiming initialization and our NIO with

Original: ResNet-110 on CIFAR-10.
Revised: ResNet-110 on CIFAR-10.

Original: As shown in Figure 3, compared with the Kaiming initialization, NIO enjoys signiﬁcantly smaller ratios of the maximal gradient norm to the minimal one, and larger cosine similarities of sample-wise gradients.
Revised: As shown in Figure 3, compared with the Kaiming initialization, NIO enjoys signiﬁcantly smaller ratios of the maximal gradient norm to the minimal one, and larger cosine similarities of sample-wise gradients.

Original: Therefore, after NIO, the gradient norm ratio is reduced towards 1.0, while the cosine similarity is improved towards 1.0, which indicates that the quantity in
Revised: Therefore, after NIO, the gradient norm ratio is reduced towards 1.0, while the cosine similarity is improved towards 1.0, which indicates that the quantity in

Original: Eq.
Revised: Eq.

Original: (2) is minimized, and the sample-wise gradients are close with respect to length and direction.
Revised: (2) is minimized, and the sample-wise gradients are close with respect to length and direction.

Original: Ablation Studies on Batch Setting.
Revised: Ablation Studies on Batch Setting.

Original: We further test the sensitivity of our method to the batch setting hyper-parameters in Eq.
Revised: We further test the sensitivity of our method to the batch setting hyper-parameters in Eq.

Original: (7), the number of sub-batches D and overlap ratio r .
Revised: (7), the number of sub-batches D and overlap ratio r .

Original: We show the model performances and the time and space complexities of NIO with ResNet-110 on CIFAR-100 using different D and r in Tabel 3.
Revised: We show the model performances and the time and space complexities of NIO with ResNet-110 on CIFAR-100 using different D and r in Tabel 3.

Original: The performances do not vary signiﬁcantly with the change of D and r .
Revised: The performances do not vary signiﬁcantly with the change of D and r .

Original: The worst accuracy (when D = 2 and r = 0 ) 75.56% in Table 3 is still higher than the Kaiming initialization result 74.54% by 1%.
Revised: The worst accuracy (when D = 2 and r = 0 ) 75.56% in Table 3 is still higher than the Kaiming initialization result 74.54% by 1%.

Original: As D increases, the time and memory consumption of NIO also grows, because more batch gradients need to be calculated by back propagation.
Revised: As D increases, the time and memory consumption of NIO also grows, because more batch gradients need to be calculated by back propagation.

Original: When D is small, a larger r should be adopted to relieve the violation of the sample-wise hypothesis.
Revised: When D is small, a larger r should be adopted to relieve the violation of the sample-wise hypothesis.

Original: Experiment Results on ImageNet.
Revised: Experiment Results on ImageNet.

Original: To verify that our method can be extensible to large-scale task, we further evaluate our approach on the ImageNet dataset.
Revised: To verify that our method can be extensible to large-scale task, we further evaluate our approach on the ImageNet dataset.

Original: Using different initialization methods,
Revised: Using different initialization methods,

Original: A smaller number of sub-batches D leads to faster speed and less memory consumption, and almost does not harm the ﬁnal performance.
Revised: A smaller number of sub-batches D leads to faster speed and less memory consumption, and almost does not harm the ﬁnal performance.

Original: “Time” refers to the time consumed for each iteration in NIO.
Revised: “Time” refers to the time consumed for each iteration in NIO.

Original: [Table]
Revised: [Table]

Original: The results of FixUp and MetaInit are taken from their papers [58, 8]
Revised: The results of FixUp and MetaInit are taken from their papers [58, 8]

Original: [Table]
Revised: [Table]

Original: and the other methods are reimplemented by us under the same training setting.
Revised: and the other methods are reimplemented by us under the same training setting.

Original: FixUp removes the BN layers as their default and the other models are trained with BN layers.
Revised: FixUp removes the BN layers as their default and the other models are trained with BN layers.

Original: [Table]
Revised: [Table]

Original: [Table]
Revised: [Table]

Original: we train ResNet-50 for 100 epochs with a batchsize of 256.
Revised: we train ResNet-50 for 100 epochs with a batchsize of 256.

Original: Alll models are trained under the same training setting.
Revised: Alll models are trained under the same training setting.

Original: As shown in Table 4, we improve upon the Kaiming initialization by 0.3%.
Revised: As shown in Table 4, we improve upon the Kaiming initialization by 0.3%.

Original: We also re-implement GradInit using the default conﬁgurations in their paper [60], and have a performance close to the Kaiming initialization baseline.
Revised: We also re-implement GradInit using the default conﬁgurations in their paper [60], and have a performance close to the Kaiming initialization baseline.

Original: Moreover, NIO takes the shortest initialization time compared with the other two learning-based initialization methods, MetaInit and GradInit.
Revised: Moreover, NIO takes the shortest initialization time compared with the other two learning-based initialization methods, MetaInit and GradInit.

Original: It is because only 100 iterations are needed for NIO in our implementation, while MetaInit and GradInit require 1000 and 2000 iterations, respectively, as their default settings [8, 60].
Revised: It is because only 100 iterations are needed for NIO in our implementation, while MetaInit and GradInit require 1000 and 2000 iterations, respectively, as their default settings [8, 60].

Original: Compared with the training time, the initialization time of NIO is negligible.
Revised: Compared with the training time, the initialization time of NIO is negligible.

Original: 6.2 Initialization for Transformer
Revised: 6.2 Initialization for Transformer

Original: Vision Transformer architectures have been popular and proven to be effective for many vision tasks [11, 46, 32, 53].
Revised: Vision Transformer architectures have been popular and proven to be effective for many vision tasks [11, 46, 32, 53].

Original: But Transformer architectures suffer from unstable training due to the large gradient variance [31, 61] in the Adam optimizer [26], so have to resort to the warmup trick, which adopts a small learning rate at initialization to stabilize the training process.
Revised: But Transformer architectures suffer from unstable training due to the large gradient variance [31, 61] in the Adam optimizer [26], so have to resort to the warmup trick, which adopts a small learning rate at initialization to stabilize the training process.

Original: In order to test the effectiveness of
Revised: In order to test the effectiveness of

Original: NIO to improve the training stability, we perform NIO for Swin-Transformer [32] with and without warmup on ImageNet.
Revised: NIO to improve the training stability, we perform NIO for Swin-Transformer [32] with and without warmup on ImageNet.

Original: Detailed training and initialization settings are described in Appendix B. As shown in Table 5, when warmup is not enabled for training, both Kaiming and Truncated normal initialized models cannot converge.
Revised: Detailed training and initialization settings are described in Appendix B. As shown in Table 5, when warmup is not enabled for training, both Kaiming and Truncated normal initialized models cannot converge.

Original: As a comparison, the model with NIO achieves a top-1 accuracy of 80.9%, which is very close to the standard training result with warmup.
Revised: As a comparison, the model with NIO achieves a top-1 accuracy of 80.9%, which is very close to the standard training result with warmup.

Original: These results indicate that
Revised: These results indicate that

Original: NIO is able to produce a better initialization agnostic of architecture and dataset.
Revised: NIO is able to produce a better initialization agnostic of architecture and dataset.

Original: 6.3 Discussions
Revised: 6.3 Discussions

Original: Why not Gradient Variance?
Revised: Why not Gradient Variance?

Original: As indicated by Eq.
Revised: As indicated by Eq.

Original: (2) and Figure 3, the aim of NIO is to reduce the variability of sample-wise gradients.
Revised: (2) and Figure 3, the aim of NIO is to reduce the variability of sample-wise gradients.

Original: But why do we not directly minimize the gradient variance or just the pair-wise Euclidean distances?
Revised: But why do we not directly minimize the gradient variance or just the pair-wise Euclidean distances?

Original: We point out that gradient variance is highly sensitive to gradient norm.
Revised: We point out that gradient variance is highly sensitive to gradient norm.

Original: Directly minimizing gradient variance would lead to a trivial gradient norm near zero, which disables the training.
Revised: Directly minimizing gradient variance would lead to a trivial gradient norm near zero, which disables the training.

Original: So the ability of reducing gradient variance while keeping gradient norm at a proper range is one crucial property of our proposed NIO.
Revised: So the ability of reducing gradient variance while keeping gradient norm at a proper range is one crucial property of our proposed NIO.

Original: Limitations and Societal Impacts.
Revised: Limitations and Societal Impacts.

Original: Similar to [60], the choice of γ has not been determined from a theoretical point yet.
Revised: Similar to [60], the choice of γ has not been determined from a theoretical point yet.

Original: Calculating GradCosine accurately requires much memory consumption for large models, so it needs the batch setting relaxation.
Revised: Calculating GradCosine accurately requires much memory consumption for large models, so it needs the batch setting relaxation.

Original: Besides, performing NIO without a dataset will be also a promising breakthrough that deserves future exploration.
Revised: Besides, performing NIO without a dataset will be also a promising breakthrough that deserves future exploration.

Original: Our study is general and not concerned with malicious practices with respect to privacy, public health, fairness, etc .
Revised: Our study is general and not concerned with malicious practices with respect to privacy, public health, fairness, etc .

Original: The proposed algorithm is lightweight and does not bring much environment burden.
Revised: The proposed algorithm is lightweight and does not bring much environment burden.

Original:  Conclusion
Revised:  Conclusion

Original: In this paper, we propose a differentiable quantity with theoretical insight to evaluate the initial state of a neural network.
Revised: In this paper, we propose a differentiable quantity with theoretical insight to evaluate the initial state of a neural network.

Original: Based on our sample-wise landscape analysis, we show that maximizing the quantity with the maximal gradient norm upper bounded is able to improve both training dynamic and model performance.
Revised: Based on our sample-wise landscape analysis, we show that maximizing the quantity with the maximal gradient norm upper bounded is able to improve both training dynamic and model performance.

Original: Accordingly, we develop an initialization optimization algorithm to rectify the initial parameters.
Revised: Accordingly, we develop an initialization optimization algorithm to rectify the initial parameters.

Original: Experimental results show that our method is able to automatically produce a better initialization for variant architectures to improve the performance on multiple datasets with negligible cost.
Revised: Experimental results show that our method is able to automatically produce a better initialization for variant architectures to improve the performance on multiple datasets with negligible cost.

Original: It also helps a Transformer architecture train stably even without warmup.
Revised: It also helps a Transformer architecture train stably even without warmup.

Original: Acknowledgments and Disclosure of Funding
Revised: Acknowledgments and Disclosure of Funding

Original: Z. Lin was supported by the NSF China (No.s 62276004 and 61731018), NSFC Tianyuan Fund for
Revised: Z. Lin was supported by the NSF China (No.s 62276004 and 61731018), NSFC Tianyuan Fund for

Original: Mathematics (No. 12026606), and Project 2020BD006 supported by PKU-Baidu Fund.
Revised: Mathematics (No. 12026606), and Project 2020BD006 supported by PKU-Baidu Fund.

Original: Checklist
Revised: Checklist

Original: For all authors... (a) Do the main claims made in the abstract and introduction accurately reﬂect the paper’s contributions and scope?
Revised: For all authors... (a) Do the main claims made in the abstract and introduction accurately reﬂect the paper’s contributions and scope?

Original: [Yes] See the last three paragraphs in Section 1.
Revised: [Yes] See the last three paragraphs in Section 1.

Original: (b) Did you describe the limitations of your work?
Revised: (b) Did you describe the limitations of your work?

Original: [Yes] See Section 6.3.
Revised: [Yes] See Section 6.3.

Original: (c) Did you discuss any potential negative societal impacts of your work?
Revised: (c) Did you discuss any potential negative societal impacts of your work?

Original: [Yes] See Section 6.3.
Revised: [Yes] See Section 6.3.

Original: (d) Have you read the ethics review guidelines and ensured that your paper conforms to them?
Revised: (d) Have you read the ethics review guidelines and ensured that your paper conforms to them?

Original: [Yes] We are ensured.
Revised: [Yes] We are ensured.

Original: If you are including theoretical results... (a) Did you state the full set of assumptions of all theoretical results?
Revised: If you are including theoretical results... (a) Did you state the full set of assumptions of all theoretical results?

Original: [Yes] See Section 3.
Revised: [Yes] See Section 3.

Original: (b) Did you include complete proofs of all theoretical results?
Revised: (b) Did you include complete proofs of all theoretical results?

Original: [Yes] See Appendix A. 3.
Revised: [Yes] See Appendix A. 3.

Original: If you ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
Revised: If you ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?

Original: [Yes] See supplementary material.
Revised: [Yes] See supplementary material.

Original: (b) Did you specify all the training details (e
Revised: (b) Did you specify all the training details (e

Original: g., data splits, hyperparameters, how they were chosen)?
Revised: g., data splits, hyperparameters, how they were chosen)?

Original: [Yes] See Appendix B. (c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
Revised: [Yes] See Appendix B. (c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?

Original: [Yes] See Table 1 and 2.
Revised: [Yes] See Table 1 and 2.

Original: (d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)?
Revised: (d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)?

Original: [Yes] See Table 4.
Revised: [Yes] See Table 4.

Original: If you are using existing assets (e.g., code, data, models) or curating/releasing new assets... (a) If your work uses existing assets, did you cite the creators?
Revised: If you are using existing assets (e.g., code, data, models) or curating/releasing new assets... (a) If your work uses existing assets, did you cite the creators?

Original: [Yes] We cite the corresponding papers.
Revised: [Yes] We cite the corresponding papers.

Original: (b) Did you mention the license of the assets?
Revised: (b) Did you mention the license of the assets?

Original: [N/A] (c) Did you include any new assets either in the supplemental material or as a URL?
Revised: [N/A] (c) Did you include any new assets either in the supplemental material or as a URL?

Original: [N/A] (d) Did you discuss whether
Revised: [N/A] (d) Did you discuss whether

Original: and how consent was obtained from people whose data you’re using/curating?
Revised: and how consent was obtained from people whose data you’re using/curating?

Original: [N/A]
Revised: [N/A]

Original: (e) Did you discuss whether the data you are using/curating contains personally identiﬁable information or offensive content?
Revised: (e) Did you discuss whether the data you are using/curating contains personally identiﬁable information or offensive content?

Original: [N/A] 5.
Revised: [N/A] 5.

Original: If you used crowdsourcing or conducted research with human subjects... (a) Did you include the full text of instructions given to participants and screenshots, if applicable?
Revised: If you used crowdsourcing or conducted research with human subjects... (a) Did you include the full text of instructions given to participants and screenshots, if applicable?

Original: [N/A] (b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?
Revised: [N/A] (b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?

Original: [N/A] (c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?[N/A] (Appendix A Proofs of Theoretical Results A.1 Lemma 1 Proof
Revised: [N/A] (c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?

Original: Please refer to [59] for the proof.

Original: A.2 Theorem

Original: Proof

Original: Let θ ∗ denote an overall local optimum of the training loss L = 1 n (cid:80) ni l ( f θ ( x i ) , y i ) via gradient descent from the initialization θ 0 .

Original: Since we have assumed that ∀ k ∈ [ m ] , i ∈ [ n ] , [ ∇ 2 l ( f θ ( x i ) , y i )]

Original: k,k ≤

Original: H , through the expansion of l ( f θ ∗ ( x i ) , y i ) at the sample optimum θ ∗ i , we have:

Original: [Equation]
Revised: [N/A]

Original: At the sample optimum θ ∗ i , we have l ( f θ ∗ i ( x i ) , y i )

Original: [-= 
-]
Revised: [-= 
-]

Original: and ∇ θ l ( f θ ∗ i ( x i ) , y i )

Original: = 0 .

Original: Then the above inequality can be rewrote as:

Original: [Equation]

Original: and thus we have L ≤ H 2 n (cid:80) ni

Original: || θ ∗ − θ ∗

Original: i

Original: || 22 .

Original: Due to the assumption that the sample-wise optimization landscapes are nearly convex and semi-smooth at the neighborhood of the initialization θ 0 , the overall optimization landscape is also convex and semi-smooth, and the overall local optimum θ ∗ lies in the convex hull of the sample-wise optima θ ∗ i , ∀ i ∈ [1 , n ] .

Original: Then we have:

Original: [Equation]

Original: 11) and

Original: [Equation]

Original: We use α and β to denote the maximal and minimal (cid:96) 2 -norms of the sample-wise optimization paths, respectively, i.e., α = max( || θ ∗ i − θ 0 || 2 ) , β = min( || θ ∗ i − θ 0 || 2 ) , ∀ i ∈ [1 , n ] .

Original: We further denote the normalized sample-wise optimization paths as p i

Original: [-= θ ∗ i − θ 
-]
Revised: [-= θ ∗ i − θ 
-]

Original: α , 0 < || p

Original: i || 2 ≤ 1 .

Original: Then we have:

Original: [Equation]

Original: [Equation]

Original: where a ≤ holds because for any pair of ( i, j ) , i, j ∈ [1 , n ] , we have |

Original: [Equation]

Original: p

Original: [Equation]

Original: ≤ αβ

Original: [Equation]

Original: A.3 Theorem 3 Proof

Original: We use the population error E ( x u ,y u ) ∼D [ l ( f θ ∗ ( x u ) , y u )] to measure the generalization performance, where D is the data distribution for both training and test data, and ( x u , y u ) refers to the test sample.

Original: With probability 1 − δ , we have:

Original: [Equation]

Original: (16) where n is the number of training samples and σ is a positive constant such that V ar ( x

Original: [Equation]

Original: ,y u )

Original: ∼D [ || θ ∗ − θ ∗ u || 22 ] ≤ σ 2 .

Original: Eq.

Original: (14) is based on Eq.

Original: (10), Eq.

Original: (15) holds following Chebyshev’s inequality, and Eq.

Original: (16) is derived from the result of Theorem 2.

Original: With Theorem 2 and Theorem 3, we can conclude that Θ S,l ( f θ 0 ( · )) is the upper bound for both training and generalization error, and thus can be used as an indicator to evaluate the quality of the initialization θ 0 .

Original: B Experimental Details

Original: In this section, we provide more information about the implementation details, the architectures, the conﬁgurations of NIO, and the training details after NIO.

Original: B.1 Implementation Details

Original: Selection of Hyper-parameters.

Original: The hypeparameters used in the NIO Algorithm 3 are mainly selected according to the following principles:

Original: (1) For the number of sub-batches D and overlap ratio r in the batch-wise relaxation, while a smaller D corresponds to a faster inference speed, it also means more violation of the sample-wise analysis.

Original: Therefore, the smaller D is, the larger r should be adopted to relieve the violation.

Original: Our ablation study also validates this practice.

Original: For all experiments, we set D as 2 to obtain the fastest initialization speed and r is correspondingly selected from 0 .

Original: 6 ∼ 0 . 8 , proportional to the batch size B .

Original: For the iteration steps T , we set it to a number that traverses all the training samples on CIFAR-10/100.

Original: For the large-scale dataset

Original: ImageNet, traversing all the training samples is time-consuming.

Original: So we set T as 100 for experiments on ImageNet.

Original: For the upper bound of the gradient norm γ , we follow the choices suggested in [60].

Original: The learning rate τ of scale coefﬁcients is set within the range from 10 − 3 to 1 .

Original: Networks with more parameters are initialized with a smaller learning rate τ .

Original: B.2 On CIFAR-10/CIFAR-

Original: Architecture Details.

Original: The architectures adopted include a ResNet-110 with an initial width ofand two convolution layers in each residual block; a 28-layer Wide ResNet with a widen factor of(WRN-28-10); and a DenseNet-100.

Original: Following [60], for experiments without BN layers, we replace the BN layers with learnable bias parameters in the corresponding places to isolate the normalization effect from the afﬁne function in BN.

Original: When testing NIO on these architectures, we ﬁrst adopt the non-zero Kaiming initialization to all convolution and FC layers, and then perform NIO upon the initialized parameters.

Original: These architecture settings remain unchanged for CIFAR-100 except that the channel size of the last linear layer is adjusted according to the number of categories.

Original: NIO Conﬁgurations.

Original: For architectures with BN layers enabled, we select the learning rate τ from relatively larger values { 3 × 10 − 1 , 2 . 5 × 10 − 1 , 2 × 10 − 1 , 1 . 5 × 10 − 1 , 1 × 10 − 1 } .

Original: We ﬁnd that the best τ for ResNet-110 (w/ BN), DenseNet-100 (w/ BN), and WRN-28-10 (w/ BN) are 1 × 10 − 1 , 1 . 5 × 10 − 1 , and 2 . 5 × 10 − 1 , respectively.

Original: For architectures without BN layers, we try τ from smaller values { 3 × 10 − 2 , 2 . 5 × 10 − 2 , 2 × 10 − 2 , 1 . 5 × 10 − 2 , 1 × 10 − 2 } .

Original: The best τ for ResNet-110 (w/o BN),

Original: DenseNet-100 (w/o BN), and WRN-28-10 (w/o BN) are 1 .

Original: 5 × 10 − 2 , 1 . 5 × 10 − 2 , and 3 × 10 − 1 , respectively.

Original: To speed up NIO for a fast initialization, we use the batch setting for all architectures with and without BN layers.

Original: We set the number of sub-batches D as 2 under a batch size of B = 128 .

Original: To relieve the violation of the sample-wise analysis, we adopt an overlap between the two sub-batches.

Original: Speciﬁcally, the best overlap ratios r for ResNet-110 (w/ BN), DenseNet-100 (w/ BN), and WRN28-10 (w/ BN) are 0 .

Original: 6 , 0 . 6 , and 0 .

Original: 7 , respectively.

Original: The best overlap ratios r for ResNet-110 (w/o

Original: BN), DenseNet-100 (w/o BN), and WRN-28-10 (w/o BN) are 0 .

Original: 75 , and 0 .

Original: 8 respectively.

Original: Finally, for the upper bound of the gradient norm γ , we follow the choices suggested in [60].

Original: The best γ for ResNet-110 (w/ BN), ResNet-110 (w/o BN), DenseNet-100 (w/ BN), DenseNet-100 (w/o BN),

Original: WRN-28-10 (w/ BN), and WRN-28-10 (w/o BN) are 5 , 4 , 3 , 3 . 5 , 2 , and 2 , respectively for CIFAR-10.

Original: CIFAR-100 has 10 2 classes, so the initial cross entropy loss will be about 2 times as large as the one on CIFAR-10 for the same architecture.

Original: Following the principle in [60], we accordingly set γ 2 times as large as that for CIFAR-10.

Original: NIO is iterated until all the training samples are covered.

Original: Training details.

Original: After NIO, we train these networks on CIFAR-10 and CIFAR-100 for 500 epochs with a batch size of 128 on a single GPU.

Original: The initial learning rate is 0.1 and decays after each iteration following the cosine annealing learning rate schedule without restart [33].

Original: We set weight decay to 10 − 4 in all experiments.

Original: For networks without BN layers, we apply gradient clipping (with a maximum gradient norm of 1) to prevent from gradient explosion.

Original: Following the standard data augmentation schemes, we adopt random cropping, random ﬂipping, and cutout [10] for all networks.

Original: We do not use dropout in any of our experiments.

Original: B.3 On ImageNet

Original: Architecture Details.

Original: To show the extensibility of NIO to large datasets and novel architectures, we also conduct experiments on ImageNet with ResNet-50 and Swin-Transformer.

Original: Speciﬁcally, we use the standard ResNet-50 architecture that is stacked by 3-layer bottleneck blocks.

Original: BN layers are enabled as default.

Original: For experiments of Swin-Transformer, we adopt the Swin-Tiny conﬁguration, where the window size M is 7, the query dimension of each head d is 32, and the expansion ratio of each MLP block is 4.

Original: The initial channel number is 96, and the numbers of blocks for four stages are { 2 , 2 , 6 , 2 } .

Original: We follow the ofﬁcial implementation in [32].

Original: NIO Conﬁgurations.

Original: Since the networks on ImageNet have more parameters than the architectures on CIFAR10/100, we adopt a relatively smaller learning rate τ .

Original: Concretely, we try τ from { 5 × 10 − 2 , 1 × 10 − 2 , 5 × 10 − 3 , 3 × 10 − 3 , 1 × 10 − 3 } and choose τ = 3 × 10 − 3 for both ResNet-50 and

Original: Swin-Transformer.

Original: We set the number of sub-batches D = 4 under the batch size B = 64 and use an overlap ratio of r = 0 .

Original: 2 for both architectures.

Original: Following the previous principle, we change the upper bound of the gradient norm γ three times as large as that for CIFAR-10 due to the increased number of categories.

Original: Finally, we run NIO for 100 iterations on ImageNet to keep the initialization process efﬁcient.

Original: Training Details.

Original: Following the standard implementations, we train ResNet-50 for 100 epochs and Swin-Transformer for 300 epochs.

Original: The batchsize is 256 for ResNet-50 and 1024 for SwinTransformer among 8 NVIDIA A-100 GPUs.

Original: ResNet-50 is trained by the SGD optimizer with a weight decay of 10 − 4 and a momentum of 0.9.

Original: Swin-Transformer is trained by the AdamW optimizer with a weight decay of 0.05 and a momentum of (0.9, 0.999).

Original: The initial learning for

Original: ResNet-50 is 0.1, and is divided by a factor of 10 at epoch 30, 60, and 90.

Original: The base learning rate for Swin-Transformer is 1 × 10 − 3 , and follows the cosine annealing learning rate schedule without restart.

Original: For Swin-Transformer with warmup enabled, a linear learning rate warmup is used for the ﬁrst 20 epochs of training.

Original: Swin-Transformer adopts gradient clipping with a maximum gradient norm of 5, and the standard augmentation schemes, such as Mixup [57].

Original: C More Explanations and Discussions

Original: As suggested by the reviewers of this paper, we offer more explanations and discussions here for some contents in the main paper.

Original: Explanations for Figure 1.

Original: In Figure 1 (c) and (d), it seems that an initialization with a smaller cosine similarity (larger angle) is closer to the global optimum of the two samples.

Original: But we should note that: (1) The optimum of a given network is not unique.

Original: It is dependent on its initialization.

Original: The model with different initialization points will converge to different optimal parameters, even though their performances may be close.

Original: So, what we do here is to look for an initialization whose sample-wise optimization paths are more consistent (smaller angle), instead of the inverse way—looking for a point closer to the global optimum of some given sample-wise optima and landscapes, which is impossible for a real network.

Original: The larger optimization path consistency induces a smaller gradient variance, whose beneﬁts have been supported by prior studies [31, 61].

Original: (2) The illustration in Figure 1 is a toy example with only two samples and simple landscapes.

Original: In a real network, the landscapes are highly non-convex in a high dimension, and there are a larger number of training samples.

Original: In this case, the global optimum does not necessarily lie in the convex hull.

Original: The point with a larger angle to sample-wise optima is not ensured to be closer to the global optimum.

Original: In conclusion, an initialization with more consistent sample-wise optimization paths is more favorable.

Original: Our aim is to look for such initialization, instead of the global optimum of given sample-wise landscapes.

Original: Figure 1 (c) and (d) are toy examples that are only used to show that the metric Eq.

Original: (1) is agnostic of initialization, while our Eq.

Original: (2) reﬂects the optimization path consistency and is a function of initialization.

Original: Explanations for some concepts.

Original: Sample-wise optimization refers to training the model only on each single sample.

Original: So, the optimization path for sample i can be characterized by θ ∗ i − θ 0 , where θ 0 is the initialized point, and θ ∗ i is the converged optimum of only training on sample i .

Original: Because all sample-wise optimization paths have the same starting point θ 0 , we use the averaged Cosine similarity to measure their consistency.

Original: The consistency of sample-wise optimization path can be formulated as 1 n 2 (cid:80) i,j cos ∠ (cid:0) θ ∗ i − θ 0 , θ ∗ j − θ 0 (cid:1) .

Original: So, if the angle between the paths from the initialization θ 0 to each local optimum θ ∗ i is small, we will have more consistent sample-wise optimization paths.

Original: It is approximated by our GradCosine in

Original: Line 6 in Algorithm 1.

Original: It is shown that our proposed Eq.

Original: (2) reﬂects the optimization path consistency and is a function of initialization θ 0 .

Original: The aim of our NIO is to look for an initialization θ 0 that maximizes GradCosine.

Original: Discussions about the rationality of the approximation in Eq.

Original: Note that θ ∗ i in Eq.

Original: (3) is not the global optimum.

Original: It is the local optimum for sample i .

Original: If we optimize a model on only one training sample, it is very easy to ﬁnish the training.

Original: Only several iterations are needed to attain a zero loss.

Original: So, we make the ﬁrst-order approximation for the sample-wise optimization, i.e. , the sample-wise optimum can be reached via only one step of gradient descent.

Original: Discussions about the comparison with GradSign [59].

Original: In GradSign [59], the metric of sample-wise optima density (Eq.

Original: (1) in our paper) is proposed and proved to be an upper bound of training and generalization error.

Original: Because sample-wise optima are not tractable, they approximate it by counting the gradient sign, which is non-differentiable.

Original: The metric is used to rank neural architectures without training for a neural architecture search purpose.

Original: However, the goal of our work is to develop a method for initialization optimization, instead of architecture search.

Original: So, our theoretical metric needs to be differentiable and be a function of initialization.

Original: The differences between [59] and our work are summarized as follows:

Original: (1) The metric Eq.

Original: (1) proposed in [59] is agnostic of initialization (as shown in Figure 1 (c) and (d

Original: Its approximated quantity, GradSign, is non-differentiable.

Original: So, it cannot be used for initialization optimization.

Original: In contrast, our metric Eq.

Original: (2) reﬂects the optimization path consistency as a function of initialization.

Original: It is approximated by the differentiable GradCosine, and thus can serve for our initialization optimization purpose.

Original: (2) Our method considers optimization path consistency, which is intimately related to a favorable training dynamic.

Original: The larger optimization path consistency induces a smaller gradient variance, whose beneﬁts have been supported by prior studies [31, 61].

Original: See more details in the “Relations to

Original: Favorable Training Dynamic” paragraph.

Original: As a comparison, [59] does not enjoy this advantage.

Original: Although both metrics Eq.

Original: (1) and Eq.

Original: (2) are related to model performance, because our work and [59] aim for different tasks, the other parts of our work including the GradCosine quantity, the objective of our neural initialization optimization, and the algorithm to solve it, are original and have no similarity with [59].