Title: TrAct: Making First-layer Pre-Activations Trainable
Venue: NeurIPS 2024
Accepted: yes
Keywords: computer vision

convolution

second-order

optimization
Abstract: We consider the training of the first layer of vision models and notice the clear relationship between pixel values and gradient update magnitudes: the gradients arriving at the weights of a first layer are by definition directly proportional to (normalized) input pixel values. Thus, an image with low contrast has a smaller impact on learning than an image with higher contrast, and a very bright or very dark image has a stronger impact on the weights than an image with moderate brightness. In this work, we propose performing gradient descent on the embeddings produced by the first layer of the model. However, switching to discrete inputs with an embedding layer is not a reasonable option for vision models. Thus, we propose the conceptual procedure of (i) a gradient descent step on first layer activations to construct an activation proposal, and (ii) finding the optimal weights of the first layer, i.e., those weights which minimize the squared distance to the activation proposal. We provide a  closed form solution of the procedure and adjust it for robust stochastic training while computing everything efficiently. Empirically, we find that TrAct (Training Activations) speeds up training by factors between 1.25x and 4x while requiring only a small computational overhead. We demonstrate the utility of TrAct with different optimizers for a range of different vision models including convolutional and transformer architectures.

TrAct: Making First-layer Pre-Activations Trainable

Abstract

We consider the training of the first layer of vision models and notice the clear relationship between pixel values and gradient update magnitudes: the gradients arriving at the weights of a first layer are by definition directly proportional to (normalized) input pixel values. Thus, an image with low contrast has a smaller impact on learning than an image with higher contrast, and a very bright or very dark image has a stronger impact on the weights than an image with moderate brightness. In this work, we propose performing gradient descent on the embeddings produced by the first layer of the model. However, switching to discrete inputs with an embedding layer is not a reasonable option for vision models. Thus, we propose the conceptual procedure of (i) a gradient descent step on first layer activations to construct an activation proposal, and (ii) finding the optimal weights of the first layer, i.e., those weights which minimize the squared distance to the activation proposal. We provide a closed form solution of the procedure and adjust it for robust stochastic training while computing everything efficiently. Empirically, we find that TrAct (Training Activations) speeds up training by factors between 1.25 × and 4 × while requiring only a small computational overhead. We demonstrate the utility of TrAct with different optimizers for a range of different vision models including convolutional and transformer architectures.

1 Introduction

We consider the learning of first-layer embeddings / pre-activations in vision models, and in particular learning the weights with which the input images are transformed in order to obtain these embeddings. In gradient descent, the updates to first-layer weights are directly proportional to the (normalized) pixel values of the input images. As a consequence (assuming that input images are standardized), high contrast, very dark, or very bright images have a greater impact on the trained first-layer weights, while low contrast images with medium brightness have only smaller impact on training.

While, in the past, mainly transformations of the input images, especially various forms of normalization have been considered, either as a preprocessing step or as part of the neural network architecture, our approach targets the training process directly without modifying the model architecture or any preprocessing. The goal of our approach is to achieve a training behavior that is equivalent to training the pre-activations or embedding values themselves. For example, in language models [1], the first layer is an 'Embedding' layer that maps a token id to an embedding vector (via a lookup). When training language models, this embedding vector is trained directly, i.e., the update to the embedding directly corresponds to the gradient of the pre-activation of the first layer. As discussed above, this is not the case in vision models as, here, the updates to the first-layer weight matrix correspond to the outer product between the input pixel values and the gradient of the pre-activation of the first layer. Bridging this gap between the 'Embedding' layer in language models, and 'Conv2D' / 'Linear' / 'Dense' layers in vision models, we propose a novel technique for training the pre-activations of the latter, effectively mimicking training behavior of the 'Embedding' layer in language models. As vision models rely on pixel values rather than tokens, and any discretization of image patches,

̸

Language

Vision

TrAct

=

=

≈

△

△

△

0.50

-1.00

1.10

0.50-1.001.10

0.50

001-

1.10

2.89

+90660-

0.50

-1.00

1.10

0.52 -0.971.06

2.00

0.10-0.50

2.000.10-0.50

2.00

0.10

-0.50

4.12

-0.990.40

2.00

0.10-0.50

1.97

0.08-0.47

M

M

M

hello

0.90.6

0.90.6

world

0.60.4

0.20.4

0.60.4

0.20.4

-0.5

1.0

-0.1-0.5

0.5

1.0

0.1-0.5

<EOS>

0.51.0

0.51.0

△W

△W

△

0.000.000.00

0.90.6

1.65

-0.840.69

0.90.6

0.90.6

0.90.6

Vz

△W=

0.50

2.00

0.10-0.50

-1.00

1.10

0.50

0.000.000.00

-1.00

1.10

= MV

0.20.4

-0.1-0.5

0.50

2.000.10-0.50

-1.00

1.10

0.90-0.16 0.02

△W=

0.20.4

-0.1-0.5

0.20.4

-0.1-0.5

+入·I

0.20.4

-0.1-0.5

0.50

2.000.10-0.50

001-

1.10

2.000.10-0.50

0.51.0

2.25

-0.400.05

0.51.0

0.51.0

0.51.0

△W

△W

△W

0.000.000.00

0.90.6

S91

-0.840.69

0.90.6

0.50-1.50

1.89

0.50-1.00

1.10

0.50-1.001.10

0.20.4

0.90-0.160.02

2.89-0.990.64

0.20.4

0.600.21-0.37

0.52 -0.971.06

000

0.000.00

2.000.10-0.50

-0.1-0.5

4.12

0.990.40

-0.1-0.5

1.09-0.72

1.07

1.970.08-0.47

0.10

0.5

1.0

0.400.05

0.53

-0.92

e.g., via clustering is not a reasonable option, we approach the problem via a modification of the gradient (and therefore a modification of the training behavior) without modifying the original model architectures. We illustratively compare the updates in language and vision models and demonstrate the modification that TrAct introduces in Figure 1.

The proposed method is general and applicable to a variety of vision model architecture types, from convolutional to vision transformer models. In a wide range of experiments, we demonstrate the utility of the proposed approach, effectively speeding up training by factors ranging from 1 . 25 × to 4 × , or, within a given training budget, improving model performance consistently. The approach requires only one hyperparameter λ , which is easy to select, and our default value works consistently well across all 50 considered model architecture + data set + optimizer settings.

The remainder of this paper is organized as follows: in Section 2, we introduce related work, in Section 3, we introduce and derive TrAct from a theoretical perspective, and in Section 3.1 we discuss implementation considerations of TrAct. In Section 4, we empirically evaluate our method in a variety of experiments, spanning a range of models, data sets, and training strategies, including an analysis of the mild behavior of the hyperparameter, an ablation study, and a runtime analysis. We conclude the paper with a discussion in Section 5. The code is publicly available at github.com/Felix-Petersen/tract.

2 Related Work

It is not surprising that the performance of image classification and object recognition models depends heavily on the quality of the input images, especially on their brightness range and contrast. For example, image augmentation techniques generate modified versions of the original images as additional training examples. Some of these techniques work by geometric transformations (rotation, mirroring, cropping), others by adding noise, changing contrast or modifying the image in

the color space [2]. In the area of vision transformers [3], [4] so-called 3-augmentation (Gaussian blur, reduction to grayscale, and solarization) has been shown to be essential to performance [5]. Augmentation approaches are similar to image enhancement as a preprocessing step, because they generate possibly enhanced versions of the images as additional training examples, even though they leave the original images unchanged, which are also still used as training examples.

Another direction related to the problem we deal with in this paper are various normalizations and standardizations, starting with the most common one of standardizing the data to mean 0 and standard deviation 1 (over the training set), and continuing through batch normalization [6], weight normalization [7], layer normalization [8], which are usually applied not just for the first layer, but throughout the network, and in particular patch-wise normalization of the input images [9], which we will draw on for comparisons. We note that, e.g., Dual PatchNorm [9], in contrast to our approach, modifies the actual model architecture, but not the gradient backpropagation procedure.

However, none of these approaches directly addresses the actual concern that weight changes in the first layer are proportional to the inputs, but instead only modify the inputs and architectures to make training easier or faster. In contrast to these approaches, we address the training problem itself and propose a different way of optimizing first-layer weights for unchanged inputs. Of course, this does not mean that input enhancement techniques are superfluous with our method, but only that additional performance gains can be obtained by including TrAct during training.

In the context of deviating from standard gradient descent-based optimization [10], there are different lines of work in the space of second-order optimization [11], e.g., K-FAC [12], ViViT [13], ISAAC [14], Backpack [15], and Newton Losses [16], which have inspired our methodology for modifying the gradient computation. In particular, the proposed approach integrates second-order ideas for solving a (later introduced) sub-optimization-problem in closed-form [17], and has similarities to a special case of ISAAC [14].

3 Method

First, let us consider a regular gradient descent of a vision model. Let z = f ( x ; W ) be the first layer embeddings excluding an activation function and W be the weights of this first layer, i.e., for a fully-connected layer f ( x ; W ) = W · x . Here, we have x ∈ R n × b , z ∈ R m × b , and W ∈ R m × n for a batch size of b . We remark that our input x may be unfolded, supporting convolutional and vision transformer networks. Further, let ˆ y = g (ˆ z ; θ \ W ) = g ( f ( x ; W ); θ \ W ) be the prediction of the entire model. Moreover, let L (ˆ y, y ) be the loss function for a label y and wlog. let us assume it is an averaging loss (i.e., reduction over batch dimension via mean). During backpropagation, the gradient of the loss wrt. z , i.e., ∇ z L ( g ( z ; θ \ W ) , y ) or ∇ z L ( z ) for short, will be computed. Conventionally, the gradient wrt. W , i.e., ∇ W L ( g ( f ( x ; W ); θ \ W ) , y ) or ∇ W L ( W ) for short, is computed during backpropagation as

leading to the gradient descent update step of

Equation 1 clearly shows the direct proportionality between the gradient wrt. the first layer weights and the input (magnitudes), showing that larger input magnitudes produce proportionally larger changes in first layer network weights. We remark that a corresponding relationships also holds in later layers of the neural network, but emphasize that, in later layers, the relationship shows a proportionality to activation magnitude, which is desirable.

To resolve this dependency on the inputs and make training more efficient, we propose to conceptually optimize in the space of first layer embeddings z . In particular, we could perform a gradient descent step on z , i.e.,

Here, b is a multiplier because L ( z ) is (per convention) the empirical expectation over the batch dim.

However, now, z ⋆ depends on the inputs and is not part of the actual model parameters. We can resolve this problem by determining how to update W such that f ( x ; W ) is as close to z ⋆ as possible. Conceptually, we compute the optimal update ∆ W ⋆ by solving the optimization problem

where we (1) want to minimize the distance between z ⋆ the embeddings implied by the change of W by ∆ W , and (2) want to keep the change ∆ W small.

We enforce that weight matrix changes ∆ W are small ( ∥ ∆ W ∥ 2 ≤ ϵ ) by taking the Lagrangian of the problem, i.e.,

with a heuristically selected Lagrangian multiplier λ · b (parameterized with b because the first part is also proportional to b ). We simplify Equation 5 to

and ease the presentation by considering it from a row-wise perspective, i.e., for ∆ W i ∈ R 1 × n :

The problem is separable into a row-wise perspective because the norm ( ∥ · ∥ 2 2 ) is the squared Frobenius norm and the rows have independent solutions.

In the following, we provide a closed-form solution for optimization problem (7), which is related to [17], [18].

Lemma 1. The solution ∆ W ⋆ i of Equation 7 is

Extending the solution to ∆ W , we have

and can accordingly use it for an update step for W , i.e.,

The update in Equation 10 directly inserts the solution of the problem formulated in Equation 4. This computation is efficient as it only requires inversion of an n × n matrix, where n in the case of convolutions correspond to 3 (RGB) times the squared first layer's kernel size, and for vision transformers corresponds to the number of pixels per patch. The values of n typically range from n = 27 (CIFAR ResNet) to n = 768 (ImageNet large-scale vision transformer).

Lemma 2. Using TrAct does not change the set of possible convergence points compared to vanilla (full batch) gradient descent. Herein, we use the standard definition of convergence points as those points where no update is performed because the gradient is zero.

Proof sketch: First, we remark that only the training of the first layer is affected by TrAct. To show the statement, we show that (i) a zero gradient for GD implies that TrAct also performs no update and that (ii) TrAct performing no update implies zero gradients for GD. Proof deferred to SM A.

The statement formalizes that TrAct does not change the set of attainable models, but instead only affects the behavior of the optimization itself.

For an illustration of how TrAct affects updates to W and z , with a comparison to language models and conventional vision models, see Figure 1.

3.1 Implementation Considerations

To implement the proposed update in Equation 10 in modern automatic differentiation frameworks [19], [20], we can make use of a custom backward or backpropagation for the first layer.

Standard gradient for the weights:

implemented via a backward function:

For TrAct, we perform an in-place replacement by:

i.e., we replace the backward of the first layer by:

Details are shown in Figure 2. This applies the TrAct update from Equation 10 when using the SGD optimizer. Moreover, extensions of the update corresponding to optimizers like ADAM [21] (including, e.g., momentum, learning rate scheduler, regularizations, etc.) can be attained by using a respective optimizer and pretending (towards the optimizer) that the TrAct update corresponds to the gradient. As it only requires a modification of the gradient computation of the first layer, the proposed method allows for easy adoption in existing code. All other layers / weights ( θ \ W ) are trained conventionally without modification. Convolutions can be easily expressed as a matrix multiplication via an unfolding of the input; accordingly, we unfold the inputs respectively in the case of convolution.

Moreover, we would like to show a second method of applying TrAct that is exactly equivalent. Typically, we have some batch of data x and a first embedding layer embed , as well as a remaining network net , a loss , and targets gt . In the following, we show how the forward and backward is usually written (left) and how it can be modified to incorporate TrAct (right):

This modifies the input of embed for the gradient computation, but replaces the actual values propagated through the remaining network z.data by the original values, therefore not affecting downstream layers. which modifies the input of embed for the gradient computation, but replaces the actual values propagated through the remaining network z.data by the original values, therefore not affecting downstream layers. This illustrates interesting relationships: TrAct is minimally invasive, can be removed or included at any time without breaking the network, and does not have learnable parameters. TrAct can be seen as in some sense related to normalizing / whitening / inverting the input for the purpose of gradient computation, but then switching the embeddings back to the original first layer embeddings / activations for propagation through the remainder of the network.

We provide an easy-to-use wrapper module that can be applied to the first layer, and automatically provides the TrAct gradient computation replacement procedure. For example, for PyTorch [19], the TrAct module can be applied to nn.Linear and nn.Conv2d layers by wrapping them as

and for existing implementations, we can apply TrAct, e.g., for vision transformers via:

4 Experimental Evaluation

4.1 CIFAR-10

Setup For the evaluation on the CIFAR-10 data set [22], we consider the ResNet18 [23] as well as a small ViT model. We consider training from scratch as the method is particularly designed for this case. We perform training for 100 , 200 , 400 , and 800 epochs. For the ResNet models, we use the Adam and SGD with momentum ( 0 . 9 ) optimizers, both with cosine learning rate schedules; learning rates, due to their significance, will be discussed alongside respective experiments. Further, we use the standard softmax cross-entropy loss. For the ViT, we use Adam with a cosine learning

0

100

200

300

400

500

600

700

800

Epochs

89

90

91

92

93

Test accuracy

=0.05)

=0.1)

=0.2)

SGD   | vanilla

0

100

200

300

400

500

600

700

800

Epochs

89

90

91

92

93

Test accuracy

=0.05)

=0.1)

=0.2)

Adam | vanilla

rate scheduler as well as a softmax cross-entropy loss with label smoothing ( 0 . 1 ). The selected ViT 1 is particularly designed for effective training on CIFAR scales and has 7 layers, 12 heads, and hidden sizes of 384 . Each model is trained with a batch size of 128 on an Nvidia RTX 4090 GPU with PyTorch [19].

As mentioned above, the learning rate is a significant factor in the evaluation. Therefore, throughout this paper, to remove any bias towards the proposed method (and even give an advantage to the baseline), we utilize the optimal learning rate of the baseline also for the proposed method. For the Adam optimizer, we consider a learning rate grid of { 10 -2 , 10 -2 . 5 , 10 -3 , 10 -3 . 5 } ; for SGD with momentum, a learning rate grid of { 0 . 1 , 0 . 09 , 0 . 08 , 0 . 07 } . The optimal learning rate is determined for each number of epochs using regular training; in particular, for Adam, we have { 100 → 10 -2 , 200 → 10 -2 , 400 → 10 -3 , 800 → 10 -3 } , and, for SGD with momentum, we find that a learning rate of 0 . 08 is optimal in each case. For the ViT, we considered a learning rate grid of { 10 -3 , 10 -3 . 1 , 10 -3 . 2 , 10 -3 . 3 , 10 -3 . 4 , 10 -3 . 5 , 10 -3 . 6 , 10 -3 . 7 , 10 -3 . 8 , 10 -3 . 9 , 10 -4 } . Here, the optimal learning rates (based on the baseline) are { 100 → 10 -3 , 200 → 10 -3 . 2 , 400 → 10 -3 . 5 , 800 → 10 -3 . 5 } .

Results In Figure 3 we show the results for ResNet-18 trained on CIFAR-10. We can observe that TrAct improves the test accuracy in every setting, in particular, for both optimizers, for all four numbers of epochs, and for all three choices of the hyperparameter λ ∈ { 0 . 05 , 0 . 1 , 0 . 2 } . Moreover, we can observe that, for SGD, the accuracy after 100 epochs is already better than for the baseline after 800 epochs. For Adam, we can see that TrAct after 100 epochs performs similar to the baseline after 400 epochs, and TrAct after 200 epochs performs similar to the baseline after 800 epochs. Comparing the different choices of λ , λ = 0 . 05 performs best in most cases.

The results for the ViT model are displayed in Figure 4. Again, we can observe that TrAct consistently outperforms the baselines for all λ . Further, we can observe that TrAct with 200 epochs performs comparable to the baseline with 400 epochs. We emphasize that, again, the optimal learning rate has been selected based on the baseline. Overall, here, λ = 0 . 1 performed best.

4.2 CIFAR-100

Setup For CIFAR-100, we consider two experimental settings. First, we consider the training of 36 different convolutional model architectures based on a strong and popular repository 2 for CIFAR-100. We use the same hyperparame-

0

100

200

300

400

500

600

700

800

Epochs

87

88

89

90

91

92

93

Test accuracy

=0.05)

=0.1)

=0.2)

Adam | vanilla

ters as the reference, i.e., SGD with momentum ( 0 . 9 ), weight decay ( 0 . 0005 ), and learning rate schedule with 60 epochs at 0 . 1 , 60 epochs at 0 . 02 , 40 epochs at 0 . 004 , 40 epochs at 0 . 0008 , and a warmup schedule during the first epoch, for a total of 200 epochs. We reproduced each baseline on a set of 5 separate seeds, and discarded the models that produced NaNs on any of the 5 seeds of the baseline. To make the evaluation feasible, we limit the hyperparameter for TrAct to λ = 0 . 1 . Second, we also reproduce the ResNet-18 CIFAR-10 experiment but with CIFAR-100. The results for this are displayed in Figure 10 in the Supplementary Material and demonstrate similar relations as the corresponding Figure 3. Again, all models are trained with a batch size of 128 on a single NVIDIA RTX 4090 GPU.

Results We display the results for the 36 CIFAR-100 models in Table 1. We can observe that TrAct outperforms the baseline wrt. top-1 and top-5 accuracy for 33 and 34 out of 36 models, respectively. Further, except for those 5 models, for which TrAct and the baseline perform comparably (each better on one metric), TrAct is better than vanilla training. Specifically, for 31 models, TrAct outperforms the baseline on both metrics, and the overall best result is also achieved by TrAct. Further, TrAct improves the accuracy on average by 0 . 49% on top-1 accuracy and by 0 . 23% on top-5 accuracy, a statistically very significant improvement over the baseline. The average standard deviations

are 0 . 25% and 0 . 15% for top-1 and top-5 accuracy, respectively.

In addition, we also considered training the models with TrAct for only 133 epochs, i.e., 2 / 3 of the training time. Here, we found that, on average, regular training for 200 epochs is comparable with TrAct for 133 epochs with a small advantage for TrAct. In particular, the average accuracy of TrAct with 133 epochs is 75 . 94% (top-1) and 93 . 34% (top-5), which is a small improvement over regular training for 200 epochs. The individual results are reported in Table 7 in the Supplementary Material.

4.3 ImageNet

Finally, we consider training on the ImageNet data set [40]. We train ResNet-{ 18 , 34 , 50 } , ViT-S and ViT-B models.

ResNet Setup For the ResNet-{ 18 , 34 , 50 } models, we train for { 30 , 60 , 90 } epochs and consider base learning rates in the grid { 0 . 2 , 0 . 141 , 0 . 1 , 0 . 071 , 0 . 05 } and determine the choice for each model / training length combination with standard baseline training. We find that for each model, when training for 30 epochs, 0 . 141 performs best, and, when training for { 60 , 90 } epochs, 0 . 1 performs best as the base learning rate. We use SGD with momentum ( 0 . 9 ), weight decay ( 0 . 0001 ), and the typical learning rate schedule, which decays the learning rate after 1 / 3 and 2 / 3 of training by 0 . 1

each. For TrAct, we (again) use the same learning rate as optimal for the baseline, and consider λ ∈ { 0 . 05 , 0 . 1 , 0 . 2 } . Each ResNet model is trained with a batch size of 256 on a single NVIDIA RTX 4090 GPU.

ResNet Results We start by discussing the ResNet results and then proceed with the vision transformers. We present training plots forResNet-50 in Figure 5. Here, we can observe an effective speedup of a factor of 1 . 5 during training, which we also demonstrate in Table 2. In particular, the difference in accuracy for TrAct ( λ = 0 . 1 ) with 60 compared to the baseline with full 90 epoch training is -0 . 02% and +0 . 04% for top-1 and top-5.

0

10

20

30

40

50

60

70

80

90

Epochs

66

68

70

72

74

76

78

Top-1 accuracy

=0.05)

=0.1)

=0.2)

Adam | vanilla

0

10

20

30

40

50

60

70

80

90

Epochs

82

84

86

88

90

92

94

Top-5 accuracy

=0.05)

=0.1)

=0.2)

Adam | vanilla

ViT Setup For training the ViTs, we reproduce the 'DeiT III' [5], which provides the strongest baseline that is reproducible on a single 8-GPU node. We train each model with the same hyperparameters as in the official source code 3 . We note that the ViT-S and ViT-B are both trained at a batch size of 2 048 and are pre-trained on resolutions of 224 and 192 , respectively, and both models are finetuned on a resolution of 224 . We consider pre-training for 400 and 800 epochs. Finetuning for each model is performed for 50 epochs. For the 400 epoch pre-training with TrAct, we use the stronger λ = 0 . 1 , while for the longer 800 epoch pre-training we use the weaker λ = 0 . 2 . We train the ViT-S models on 4 NVIDIA A40 GPUs and the ViT-B models on 8 NVIDIA V100 (32GB) GPUs.

ViT Results In Table 3 we present the results for training vision transformers. First, we observe that our reproductions following the official code and hyperparameters improved over the originally reported baselines, potentially due to contemporary improvements in the underlying libraries (our hardware only supported more recent versions). Notably, TrAct consistently improves upon our improved baselines. We note that we did not change any hyperparameters for training with TrAct. For ViT-S, using TrAct leads to 36% of the improvement that can be achieved by training the baseline twice as long. These improvesmall, ranging from 0 . 08% to 0 . 25%

ments can be considered quite substantial considering that these are very large models and we modified only the training of the first layer. Notably, here, the runtime overheads were particularly . Finally, we consider the quality of the pre-trained model outside

of ImageNet. We fine-tune the ViT-S ( 800 epoch pre-training) model on the data sets CIFAR-10 and CIFAR-100 [22] ( 200 epochs), Flowers-102 [41] ( 5000 epochs), and Stanford Cars [42] ( 1000 epochs). For the baseline, both pre-training and finetuning were performed with the vanilla

method, and, for TrAct, both pre-training and fine-tuning were performed with TrAct. In Table 4, we can observe consistent improvements for training with TrAct.

4.4 Effect of λ

λ is the only hyperparameter introduced by TrAct. Often, with an additional hyperparameter, the hyperparameter space becomes more difficult to manage. However, for TrAct, the selection of λ is simple and compatible with existing hyperparameters. Therefore, throughout all experiments in this paper, we kept all other hyperparameters equal to the optimal choice for the respective baselines, and only considered λ ∈ { 0 . 05 , 0 . 1 , 0 . 2 } . A general trend is that

.00625.0125

.025

.05

.1

.2

.4

.8

1.6

90.0

90.5

91.0

91.5

Test accuracy

Vanilla

TrAct

with smaller λ s, TrAct becomes more aggressive, which tends to be more favorable in shorter training, and for larger λ s, TrAct is more moderate, which is ideal for longer trainings. However, in many cases, the particular choice of λ ∈ { 0 . 05 , 0 . 1 , 0 . 2 } has only a subtle impact on accuracy as can be seen throughout the figures in this work. Further, going beyond this range of λ s, in Fig. 6, we can observe that TrAct is robust against changes in this parameter. In all experiments, the data was as-per-convention standardized to mean 0 and standard deviation 1 ; deviating from this convention could change the space of λ s. For significantly different tasks and drastically different kernel sizes or number of input channels, we expect that the space of λ s could change. Overall, we recommend λ = 0 . 1 as a starting point and, for long training, we recommend λ = 0 . 2 .

4.5 Ablation Study

As an ablation study, we first compare TrAct to patch-wise layer normalization for ViTs. For this, we normalize the pixel values of each input patch to mean 0 and standard deviation 1 . This is an alternate solution to the conceptual problem of low contrast image regions having a lesser effect on the first layer optimization compared to higher contrast image regions. However, here, we also note that, in contrast to TrAct, the actual neural network inputs are changed through the normalization. Further, we consider DualPatchNorm [9] as a comparison, which additionally includes a second patch normalization layer after the first linear layer, and introduces additional trainable weight parameters for affine transformations into both patch normalization layers.

We use the same setup as for the CIFAR-10 ViT and run each setting for 5 seeds. The results are displayed in Figure 7. Here, we observe that patch normalization improves training for up to 400 epochs compared to the baseline; however, not as much as TrAct does. Further, we find that DualPatchNorm performs equivalently compared to input patch normalization and worse than TrAct, except for the case of 200 epochs where it performs insignificantly better than TrAct. For training for 800 epochs, patch normalization and DualPatchNorm do not improve the baseline and perform insignificantly worse, whereas TrAct still shows accuracy improvements. This effect may be explained by

0

100

200

300

400

500

600

700

800

Epochs

88.0

88.5

89.0

89.5

90.0

90.5

91.0

91.5

92.0

92.5

Test accuracy

=0.05)

=0.1)

=0.2)

Adam | PatchNorm

Adam | DualPatchNorm

the fact that patch normalization is a scalar form of whitening, and whitening can hurt generalization capabilities due to a loss of information [43]. In particular, what may be problematic is that patch normalization also affects the model behavior during inference, which contrasts TrAct.

As a second ablation study, we examine what happens if we (against convention) do not perform standardization of the data set. We train the same ViTs as above on CIFAR-10 for 200 epochs, averaged over 5 seeds. We consider two cases: first, an input value range of [0 , 1] and a quite extreme input value range of [0 , 255] .

We display the results in Figure 8. Here, we observe that TrAct is more robust against a lack of standardization. Interestingly, we observe that TrAct performs better for the range of [0 , 255] than [0 , 1] . The reason for this is that TrAct suffers from obtaining only positive inputs, which affects the xx ⊤ matrix in Equation 10; however, we note that regular training suffers even more from the lack of standardization. When considering the range of [0 , 255] , we observe that TrAct is virtually agnostic to λ , which is caused by the xx ⊤ matrix becoming very large. The reason why TrAct performs so well here (compared to

.00625

.0125

.025

.05

.1

.2

.4

.8

75

80

85

90

Test accuracy

Vanilla  [0, 1]

TrAct     [0, 1]

Vanilla  [0, 255]

TrAct     [0, 255]

the baseline) is that, due to the large xx ⊤ , the updates ∆ W become very small. This is more desirable compared to the standard gradient, which explodes due to its proportionality to the input values, and therefore drastically degrades training.

In each experiment, we used only a single optimizer for the entire model; however, our theory assumes that TrAct is used with SGD. This motivates the question of whether it is advantageous to train the first layer with SGD, while training the remainder of the model, e.g., with Adam.

Thus, as a final ablation study, we extend the experiment from Figure 3 (right) by training the first layer with SGD while training the remaining model with Adam. We display the results in Figure 9 where we can observe small improvements when using SGD for the TrAct layer.

4.6 Runtime Analysis

In this section, we provide a training runtime analysis. Overall, the trend is that, for large models, TrAct adds on a tiny runtime overhead, while it can become more expensive for smaller models. In particular, for the CIFAR-10 ViT, the average training time per 100 epochs increased by 9 . 7% from 1091 s to 1197 s. Much of this can be attributed to the required additional CUDA calls and non-fused operations, which can be expensive for cheaper tasks. However, when considering larger models, this overhead almost entirely amortizes. In particular the ViT-S ( 800 epochs) pre-training cost increased by only 0 . 08% from 133:52 hours to 133:58 hours. The pre-training cost of the ViT-B ( 400 epochs) increased by 0 . 25% from 98:28 hours to 98:43 hours. We can see that, in each case, the training cost overhead is clearly more than worth the reduced requirement of epochs already. Further, fused kernels could drastically reduce the computational overhead; in particular, our current implementation replaces an existing fused operation by multiple calls from the Python space. As TrAct only affects training, and the modification isn't present during forwarding, TrAct has no effect on inference time.

5 Discussion & Conclusion

In this work, we introduced TrAct, a novel training strategy that modifies the optimization behavior of the first layer, leading to significant performance improvements across a range of 50 experimental setups. The approach is efficient and effectively speeds up training by factors between 1 . 25 × and 4 × depending on the model size. We hope that the simplicity of integration into existing training schemes as well as the robust performance improvements motivate the community to adopt TrAct.

0

100

200

300

400

500

600

700

800

Epochs

89

90

91

92

93

Test accuracy

=0.05)

=0.1)

=0.2)

Adam | vanilla

Acknowledgments and Disclosure of Funding

This work was in part supported by the Federal Agency for Disruptive Innovation SPRIN-D, the Land Salzburg within the WISS 2025 project IDA-Lab (20102-F1901166-KZP and 20204-WISS/225/1972019), the ARO (W911NF-21-1-0125), the ONR (N00014-23-1-2159), and the CZ Biohub.

References

A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, Ł. Kaiser, and I. Polosukhin, 'Attention is all you need,' in Proc. Neural Information Processing Systems (NeurIPS) , 2017.

C. Shorten and T. M. Khoshgoftaar, 'A survey on image data augmentation for deep learning,' Journal of Big Data , vol. 6, 2019.

A. Dosovitskiy, L. Beyer, A. Kolesnikov, D. Weissenborn, X. Zhai, T. Unterthiner, M. Dehghani, M. Minderer, G. Heigold, and S. G. et al. , 'An image is worth 16x16 words: Transformers for image recognition at scale,' in Proc. International Conference on Learning Representations (ICLR) , 2020.

H. Touvron, M. Cord, M. Douze, F. Massa, A. Sablayrolles, and H. Jégou, 'Training data-efficient image transformers & distillation through attention,' in Proc. International Conference on Machine Learning (ICML) , 2021.

H. Touvron, M. Cord, and H. Jégou, 'DeiT III: Revenge of the ViT,' in Proc. European Conference on Computer Vision (ECCV) , 2022.

S. Ioffe and C. Szegedy, 'Batch normalization: Accelerating deep network training by reducing internal covariate shift,' in Proc. International Conference on Machine Learning (ICML) , 2015.

T. Salimans and D. P. Kingma, 'Weight normalization: A simple reparameterization to accelerate training of deep neural networks,' in Proc. Conf. Neural Information Processing Systems (NeurIPS 2016, Barcelona, Spain) , Neural Information Processing Systems Foundation, 2016.

J. L. Ba, J. R. Kiros, and G. E. Hinton, 'Layer normalization,' Computing Research Repository (CoRR) in arXiv , 2016.

M. Kumar, M. Dehghani, and N. Houlsby, 'Dual patchnorm,' Trans. on Machine Learning Research , 2023.

F. Dangel, 'Backpropagation Beyond the Gradient,' Ph.D. dissertation, University of Tübingen, 2023.

S. J. Wright, Numerical optimization . 2006.

J. Martens and R. Grosse, 'Optimizing neural networks with kronecker-factored approximate curvature,' in Proc. International Conference on Machine Learning (ICML) , 2015.

F. Dangel, L. Tatzel, and P. Hennig, 'ViViT: Curvature access through the generalized Gauss-Newton's low-rank structure,' Transactions on Machine Learning Research , 2022.

F. Petersen, T. Sutter, C. Borgelt, D. Huh, H. Kuehne, Y. Sun, and O. Deussen, 'ISAAC Newton: Input-based approximate curvature for Newton's method,' in Proc. International Conference on Learning Representations (ICLR) , 2023.

F. Dangel, F. Kunstner, and P. Hennig, 'Backpack: Packing more into backprop,' in Proc. International Conference on Learning Representations (ICLR) , 2020.

F. Petersen, C. Borgelt, T. Sutter, H. Kuehne, O. Deussen, and S. Ermon, 'Newton losses: Using curvature information for learning with differentiable algorithms,' in Conference on Neural Information Processing Systems (NeurIPS) , 2024.

A. E. Hoerl and R. W. Kennard, 'Ridge regression: Biased estimation for nonorthogonal problems,' Technometrics , vol. 12, no. 1, pp. 55-67, 1970.

D. Calvetti and L. Reichel, 'Tikhonov regularization with a solution constraint,' SIAM Journal on Scientific Computing , vol. 26, no. 1, pp. 224-239, 2004.

A. Paszke, S. Gross, F. Massa, A. Lerer, J. Bradbury, G. Chanan, T. Killeen, Z. Lin, N. Gimelshein, L. Antiga, et al. , 'Pytorch: An imperative style, high-performance deep learning library,' Proc. Neural Information Processing Systems (NeurIPS) , 2019.

J. Bradbury, R. Frostig, P. Hawkins, M. J. Johnson, C. Leary, D. Maclaurin, G. Necula, A. Paszke, J. VanderPlas, S. Wanderman-Milne, and Q. Zhang, JAX: Composable transformations of Python+NumPy programs , 2018. [Online]. Available: http://github.com/google/jax .

D. Kingma and J. Ba, 'Adam: A method for stochastic optimization,' in Proc. International Conference on Learning Representations (ICLR) , 2015.

A. Krizhevsky, V. Nair, and G. Hinton, 'Cifar-10 (canadian institute for advanced research),' 2009. [Online]. Available: http://www.cs.toronto.edu/~kriz/cifar.html .

K. He, X. Zhang, S. Ren, and J. Sun, 'Deep residual learning for image recognition,' in Proc. International Conference on Computer Vision and Pattern Recognition (CVPR) , 2016.

F. N. Iandola, S. Han, M. W. Moskewicz, K. Ashraf, W. J. Dally, and K. Keutzer, 'Squeezenet: Alexnet-level accuracy with 50x fewer parameters and< 0.5 mb model size,' Computing Research Repository (CoRR) in arXiv , 2016.

A. G. Howard, M. Zhu, B. Chen, D. Kalenichenko, W. Wang, T. Weyand, M. Andreetto, and H. Adam, 'Mobilenets: Efficient convolutional neural networks for mobile vision applications,' Computing Research Repository (CoRR) in arXiv , 2017.

M. Sandler, A. Howard, M. Zhu, A. Zhmoginov, and L.-C. Chen, 'Mobilenetv2: Inverted residuals and linear bottlenecks,' in Proc. International Conference on Computer Vision and Pattern Recognition (CVPR) , 2018.

X. Zhang, X. Zhou, M. Lin, and J. Sun, 'Shufflenet: An extremely efficient convolutional neural network for mobile devices,' in Proc. International Conference on Computer Vision and Pattern Recognition (CVPR) , 2018.

N. Ma, X. Zhang, H.-T. Zheng, and J. Sun, 'Shufflenet v2: Practical guidelines for efficient cnn architecture design,' in Proc. European Conference on Computer Vision (ECCV) , 2018.

K. Simonyan and A. Zisserman, 'Very deep convolutional networks for large-scale image recognition,' Computing Research Repository (CoRR) in arXiv , 2014.

G. Huang, Z. Liu, L. Van Der Maaten, and K. Q. Weinberger, 'Densely connected convolutional networks,' in Proc. International Conference on Computer Vision and Pattern Recognition (CVPR) , 2017.

C. Szegedy, W. Liu, Y. Jia, P. Sermanet, S. Reed, D. Anguelov, D. Erhan, V. Vanhoucke, and A. Rabinovich, 'Going deeper with convolutions (2014),' Computing Research Repository (CoRR) in arXiv , 2014.

C. Szegedy, V. Vanhoucke, S. Ioffe, J. Shlens, and Z. Wojna, 'Rethinking the inception architecture for computer vision,' in Proc. International Conference on Computer Vision and Pattern Recognition (CVPR) , 2016.

C. Szegedy, S. Ioffe, V. Vanhoucke, and A. Alemi, 'Inception-v4, inception-resnet and the impact of residual connections on learning,' in AAAI Conference on Artificial Intelligence , 2017.

F. Chollet, 'Xception: Deep learning with depthwise separable convolutions,' in Proc. International Conference on Computer Vision and Pattern Recognition (CVPR) , 2017.

S. Xie, R. Girshick, P. Dollár, Z. Tu, and K. He, 'Aggregated residual transformations for deep neural networks,' in Proc. International Conference on Computer Vision and Pattern Recognition (CVPR) , 2017.

J. Hu, L. Shen, and G. Sun, 'Squeeze-and-excitation networks,' in Proc. International Conference on Computer Vision and Pattern Recognition (CVPR) , 2018.

B. Zoph, V. Vasudevan, J. Shlens, and Q. V. Le, 'Learning transferable architectures for scalable image recognition,' in Proc. International Conference on Computer Vision and Pattern Recognition (CVPR) , 2018.

S. Zagoruyko and N. Komodakis, 'Wide residual networks,' Computing Research Repository (CoRR) in arXiv , 2016.

G. Huang, Y. Sun, Z. Liu, D. Sedra, and K. Q. Weinberger, 'Deep networks with stochastic depth,' in Proc. European Conference on Computer Vision (ECCV) , 2016.

J. Deng, W. Dong, R. Socher, L.-J. Li, K. Li, and L. Fei-Fei, 'Imagenet: A large-scale hierarchical image database,' in Proc. International Conference on Computer Vision and Pattern Recognition (CVPR) , 2009.

M.-E. Nilsback and A. Zisserman, 'Automated flower classification over a large number of classes,' in 2008 Sixth Indian conference on computer vision, graphics & image processing , IEEE, 2008.

J. Krause, J. Deng, M. Stark, and L. Fei-Fei, 'Collecting a large-scale dataset of fine-grained cars,' 2013.

N. Wadia, D. Duckworth, S. S. Schoenholz, E. Dyer, and J. Sohl-Dickstein, 'Whitening and second order optimization both make information in the dataset unusable during training, and can reduce or prevent generalization,' in Proc. International Conference on Machine Learning (ICML) , 2021.

S. Ren, K. He, R. Girshick, and J. Sun, 'Faster R-CNN: Towards real-time object detection with region proposal networks,' in Proc. Neural Information Processing Systems (NeurIPS) , 2015.

M. Everingham, L. Van Gool, C. K. I. Williams, J. Winn, and A. Zisserman, The PASCAL Visual Object Classes Challenge 2007 (VOC2007) Results . [Online]. Available: http://www.pascal-network.org/challenges/ VOC/voc2007/workshop/index.html .

A Theory

Lemma 1. The solution ∆ W ⋆ i of Equation 7 is

Proof. We would like to solve the optimization problem

A necessary condition for a minimum of the functional

is that ∇ ∆ W i F (∆ W i ) vanishes:

It follows for the optimal ∆ W ⋆ i that minimizes F (∆ W i )

Lemma 2. Using TrAct does not change the set of possible convergence points compared to vanilla (full batch) gradient descent. Herein, we use the standard definition of convergence points as those points where no update is performed because the gradient is zero.

Proof. First, we remark that only the training of the first layer is affected by TrAct. To show the statement, we show that ( i ) a zero gradient for GD implies that TrAct also performs no update and that ( ii ) TrAct performing no update implies zero gradients for GD.

( i ) In the first case, we assume that gradient descent has converged, i.e., the gradient wrt. first layer weights is zero ∇ W L ( W ) = 0 . We want to show that, in this case, our proposed update is also zero, i.e., ∆ W ⋆ = 0 . Using the definition of ∆ W ⋆ from Equation 9, we have

which shows this direction.

( ii ) In the second case, we have ∆ W ⋆ = 0 and need to show that this implies ∇ W L ( W ) = 0 . For this, we can observe that ( xx ⊤ / b + λ · I n ) -1 is PD (positive definite) by definition and

( xx ⊤ / b + λ · I n ) also exists. If ∆ W ⋆ = 0 , then 0 = ∆ W ⋆ = ∆ W ⋆ ( xx ⊤ b + λ · I n ) (17) = -η ·∇ z L ( z ) · x ⊤ · ( xx ⊤ b + λ · I n ) -1 ( xx ⊤ b + λ · I n ) = -η · ∇ z L ( z ) · x ⊤ = -η · ∇ W L ( W ) , (18)

which also shows this direction.

Overall, we showed that if gradient descent has converged according to the standard notion of a zero gradient, then our update has also converged and vice versa.

B Additional Results

We display additional results in Figures 10, 11, and 12 as well as in Tables 6 and 7.

As an additional experiment, in order to verify the applicability of TrAct beyond training / pre-training, we train Faster R-CNN models [44] on PASCAL VOC2007 [45] using a VGG-16 backbone [29]. However, Faster R-CNN uses a pretrained vision encoder where the first 4 layers are frozen. In order to enable TrAct, as TrAct only affects the training of the first layer, we unfreeze these first layers when training the object detection head. The mean average precision (mAP) on test data for the vanilla model versus TrAct training are shown in Table 5.

We can observe that TrAct performs better than the vanilla method by about 1.1%. We would like to point out that, while is TrAct especially designed for speeding up pretraining or training from scratch, i.e., when actually learning the first layer, we find that it also helps in finetuning pretrained models. Here, a limitation is of course that TrAct requires actually training the first layer.

0

20

40

60

80

100

60

61

62

63

64

65

66

67

68

Test accuracy

=0.05)

=0.1)

=0.2)

SGD   | vanilla

0

20

40

60

80

100

60

61

62

63

64

65

66

67

68

Test accuracy

=0.05)

=0.1)

=0.2)

Adam | vanilla

Epochs

Epochs

0

10

20

30

40

50

60

70

80

90

60

62

64

66

68

70

72

Top-1 accuracy

=0.05)

=0.1)

=0.2)

Adam | vanilla

0

10

20

30

40

50

60

70

80

90

82

84

86

88

90

92

94

Top-5 accuracy

=0.05)

=0.1)

=0.2)

Adam | vanilla

Epochs

Epochs

0

10

20

30

40

50

60

70

80

90

64

66

68

70

72

74

Top-1 accuracy

=0.05)

=0.1)

=0.2)

Adam | vanilla

0

10

20

30

40

50

60

70

80

90

82

84

86

88

90

92

94

Top-5 accuracy

=0.05)

=0.1)

=0.2)

Adam | vanilla

Epochs

Epochs

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: We address all claims made in the abstract and introduction in the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: All assumptions are pointed out in the work. Limitations are discussed.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: All assumptions are provided. Proofs are provided in the SM.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: We discuss all experimental parameters necessary for reproduction.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: The code will be made publicly available at github.com/Felix-Petersen/tract.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: All training details are either explicitly discussed in the main paper or inherited from the references.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: Standard deviations are reported and correctly defined, described as such, and utilize Bessel's correction. In some experiments, we display each seed's accuracy. In the ImageNet ViT experiments, running multiple seeds was not feasible due to compute cost.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: We describe the hardware for each experiment and a provide runtime analysis section.

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: The research conducted in this paper is conforming with the NeurIPS Code of Ethics. No animal were harmed in the execution of this research.

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [Yes]

Justification: The proposed method reduces training cost of vision models or improves vision models under the same computational budget.

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: All utilized assets (data sets, models, and code bases) are cited.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [NA]

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Peer reviews:
comments: SUMMARY:
This paper presents TrAct, a new and novel training strategy to  modify the optimization behavior of the first layer. Basically, it can achieve faster convergence and achieve better classification performance across different models. The effectiveness of TrAct is demonstrated across a range of 50 experimental setups on the various benchmarks, underscoring the method's capability.

STRENGTHS:
1. The technical elaboration of the proposed method is clear. 
2. I really like the motivation for the proposed method, which is straightforward.
3. The evaluations conducted on the provided benchmarks provide evidence of the effectiveness of the proposed methods. However, there are some concerns regarding the experimental results, which will be further discussed in the weaknesses section.

WEAKNESSES:
Weaknesses

- What I am curious about is whether Tract is also generally applicable in other visual tasks, such as detection, segmentation, pose estimation, etc. The author could simply do some experiments on FasterCNN to verify the applicability of TrAct in visual dense prediction. If applicable, this will further enhance the impact of the proposed method.
- Why is there a significant difference in the performance of TrAct between SGD and Adam? For example, in Figure 2, TrAct seems to converge faster on SGD, even surpassing the performance of the baseline 800 epoch with 100 epochs of training time. However, on Adam, performance improvement seems to be rather minor. Is there any theoretical explanation for this?

QUESTIONS:
Shown as above

LIMITATIONS:
The paper has concluded this in the checklist
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 4
DATE: 2024-07-14

comments: SUMMARY:
The paper introduces TrAct, a training strategy that modifies the optimization behaviour of the first layer. The proposed first-layer optimization enables a slightly better results for training the model at lesser number of epochs. TrAct is being demonstrated for a wider setup on image classification and have shown to be consistent.

STRENGTHS:
1. TrAct enables faster convergence or can achieve slightly better performance for the same number of epochs.

2. The paper demonstrates applicability of TrAct in various possible scenarios covering a wide range of settings.

WEAKNESSES:
1. A small overhead (due to additional parameters) in computational time when training with TrAct.

QUESTIONS:
1. What happens when only the last layer of a pre-trained model is fine-tuned? This is a setup used in evaluating many self-supervised  methods. Wouldn't the first layer optimization using TrAct create bias for a new data distrbution?

LIMITATIONS:
Not explicitly stated.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-25

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nWhen training vision models, the update of the weights of the first layer is proportional to the input pixel values. This can make the model learn images with high contrast faster and damage learning efficiency. To reduce this dependency, the paper proposed to optimize the first layer embedding (before activation) directly. This perspective leads to an optimization on a new way to update the weights and results in a light weight modification to a variety of vision models. The method was demonstrated on image classification problems.\n\nSTRENGTHS:\nThe authors identified the fundamental problem of training vision models. Compared to modifying the input or the model architecture, targeting the update of the first layer weight is a more direct approach and easy to understand the impact. They formulated an optimization problem and clearly listed out the derivation for the solution to the update of the weights.\nThe proposed method is easy to implement.\n\nWEAKNESSES:\nThere should be more elaboration on the intuition of the first layer embedding optimization. I can understand the purpose and the conceptual procedure can be viewed as activation applied to the first layer, but how this formulation reduces the dependency on the inputs without hurting the training is not obvious.\nIn the introduction, the authors argued that they bridged the gap between the “embedding” layer in language models and the first layers in vision models. However, it is not clear to me the connection of the proposed work to the “embedding” layer update in the language models.\n\nQUESTIONS:\nI believe the proposed method is general enough for vision tasks other than image classification. It would be interesting to see the performance on segmentation, detection, and even image generation. Specifically, I would like to see if the modification on the embedding hurts the performance of the tasks requiring more pixel-by-pixel understanding.\n\nLIMITATIONS:\nThe authors stated that limitations are discussed as all assumptions are pointed out in the work. However, the assumptions listed are fairly broad and not specific to the proposed work. I’d like to see more discussion on the limitations imposed when modifying the first layer embedding.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-26"}