TrAct: Making First-layer Pre-Activations Trainable

We would like to thank you for your time reviewing and for helping us to improve our paper.

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.

Thank you for appreciating the clarity of our technical elaboration, our straightforward motivation, and the effectiveness of our method.

We would appreciate if you could let us know whether we have resolved and successfully answered all of your questions and concerns. If new question or concerns come up, please let us know.

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.

Thank you for this concrete suggestion, which really helps extend our paper. Accordingly, we trained FasterRCNN models. However, we need to mention that FasterRCNN uses a pretrained vision encoder where the first 4 layers are frozen. In order to enable a fair comparison, we unfroze these first two layers when training the object detection head. We trained the models on PASCAL VOC2007, and have the following preliminary results with 2 seeds (measured in test mAP):

We can observe that both seeds with TrAct performed better than the best seed of the vanilla method, and the average improvement is 1.1%. We offer to extend these studies for the camera-ready.

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 were excited to find it also helped in finetuning pretrained models. The limitation would be that TrAct is only applicable when the first layer is actually trained.

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?

This is indeed an interesting question, and there is a theoretical explanation for this: The method is motivated for SGD training as the update that TrAct is intended to perform is returned in form of the gradient, and to execute the update the optimizer should (in theory) be SDG. That it still works very well even for Adam is great, and illustrates the strong robustness and versatility of our method. Your comment inspired us to extend Figure 2 by an experiment where we train everything except for the first layer with Adam, and train the first layer with TrAct and the SGD optimizer. Here, we observed small improvements over optimizing the first layer with Adam, which are shown in the Author Rebuttal PDF. Finally, we want to mention that, while using SGD for the first layer in combination with TrAct can improve performance, for easier adoption as it is often simpler to use the same optimizer for everything. Nevertheless, for advanced users, this is a way to squeeze out even a bit more performance. We will add a respective discussion to the camera-ready.

We would appreciate if you could let us know whether we have resolved and successfully answered all of your questions and concerns. If new question or concerns come up, please let us know.

Thanks to the authors for their rebuttal. I appreciate taking time to answer the various questions presented here. Most of my concerns have been addressed. I specifically appreciate the analysis of transfering to object detection tasks.

I will keep my original rating as accept.

We would like to thank you for your time reviewing and for helping us to improve our paper.

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.

Thank you for appreciating especially the faster convergence of our method, as well as the applicability of our method across a wide range of settings.

Weaknesses:

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

We would like to clarify that the small overhead is only due to the backward of the first layer being slightly more expensive, whereas the forward is not modified. Indeed, we do not introduce additional parameters.

Questions:

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?

During additional fine-tuning experiments, not reported in the paper, when comparing a model pre-trained with TrAct vs. pre-trained without TrAct, and finetuning both models without TrAct, we observe the model that had originally been pre-trained with TrAct performed better. If you would like us to perform this experiment also for only fine-tuning the last layer, we offer to include these results in the camera-ready; however, we strongly assume that the model pre-trained with TrAct still performs better in this case because the performance of the entire model is mostly affected by pretraining, and the first layer being optimized more efficiently, leads to a more efficient optimization of the rest of the network.

We would appreciate if you could let us know whether we have resolved and successfully answered all of your questions and concerns. If new question or concerns came up, please let us know.

We would like to thank you for your time reviewing and for helping us to improve our paper.

The 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. The proposed method is easy to implement.

Thank you very much for appreciating the clarity of the derivation of our method, and the simplicity to implement. We hope that these properties, combined with the training speedups lead to a wide adoption in the community.

There 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. In 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.

In language models, the embedding layer is a lookup table, with an embedding / activation for each token. Thus, given an input token (integer), one embedding is selected (a row of the embedding matrix), and fed forward into the next layer. The training dynamics of the embeddings layer corresponds to updating the embeddings directly wrt. the gradient. The update in a language model, for a token identifier $i$, is $W_i \leftarrow W_i - \eta \cdot \nabla_{z} \mathcal{L}(z)$ where $z = W_i$ is the activation of the first layer and at the same time the $i$th row of the embedding (weight) matrix $W$. Equivalently, we can write $z \leftarrow z - \eta \cdot \nabla_{z} \mathcal{L}(z)$.

In contrast, in vision models, the update is $W \leftarrow W - \eta \cdot \nabla_{z} \mathcal{L}(z) \cdot x^\top$ and our goal is to achieve an update that is close to $z^* \leftarrow z - \eta \cdot \nabla_{z} \mathcal{L}(z)$, which we achieve via our closed form solution.

We will include the extended discussion into the camera-ready.

I 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.

Thank you for this suggestion, which we combined with Reviewer yoFY's suggestion of training a FasterRCNN object detection model. Please see our discussion and preliminary results table in our response to Reviewer yoFY.

The 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.

There are no assumptions that are specific to this work as our method is quite generally applicable as long as the first layer is a fully-connected or convolutional layer. A potential limitation is that the approach is not applicable in settings where the first layer is frozen, e.g., at a random initialization; however, the technique of the first layer being frozen in some models stems from the first layer's training dynamics problem identified in our work and actually training the first layer with our technique could improve performance.

Please let us know if your had a different specific assumption in mind that you would like us to point out in the paper or discuss.

We would appreciate if you could let us know whether we have resolved and successfully answered all of your questions and concerns. If new question or concerns come up, please let us know.