Train Faster, Perform Better: Modular Adaptive Training in Over-Parameterized Models

Summary

The authors introduce a novel modification of the neural tangent kernel (NTK) called modular NTK (mNTK). It is essentially the decomposition of the NTK as a sum of tangent kernels for the disjoint modules that make up the network. This allows for module-level analysis of training dynamics. They find that principle eigenvalue of mNTKs tends to be orders of magnitude larger than the others and that the variation of these principle eigenvalues is a synchronous across modules. Since the directions associated with these principle eigenvalues dominates gradient updates during training, they can be utilized to selectively update subsets of the modules to better allocate computational resources during training. The authors further characterize modules with smaller max eigenvalues as being more likely to learn "nuisance" features. They introduce a novel optimization schema called modular adaptive training (MAT) that stops back propagation to modules whose principle eigenvalue falls below a threshold or if variation of the principle eigenvalue falls below a threshold. They demonstrate that MAT produces models with equivalent performance to vanilla trained model while requiring fewer FLOPs during training. Furthermore, they show that MAT trained models tend to generalize better.

Strengths

- Paper is well-written. - Method is novel and theoretically well motivated. - Supported by experiments. - The method is very general and has the potential for high impact.

Weaknesses

For the language models, the paper only looks at perplexity scores of the pretrained model for text. I'm not the most up to date on this, but I think typically practitioners are more concerned about the performance of the model when fine-tuned on down stream tasks. I think comparing performance on (vanilla) fine-tuned downstream tasks of MAT vs vanilla pretrained models would strengthen the results. Tying a bit into the previous point, papers such as RoBERTa indicate that continuing to train pretrained models after the train perplexity has mostly stopped decreasing can lead to significant boosts in performance when fine-tuned on downstream tasks. If the low $\lambda_\textrm{max}$ heads mostly learn diverse features, then this might be negatively impacted by MAT pretraining. If they mostly learn noise, then there shouldn't be much negative impact on this. Success of MAT in this setting (as indicated by increased performance on downstream tasks at a given computational cost) would greater strengthen the method.

Questions

- Line 221 mentions how the hyperparameter $\alpha$ is relatively stable throughout training. This doesn't really make sense to me because how could a constant not be stable throughout training. Is this a typo? I think you may have meant to say something like $\lambda_\alpha$ is stable throughout training. - Have you explored using MAT for the process of fine-tuning a pretrained model? Typically, fine-tuning takes far less time than pretraining, but it can still take a while to converge if the target task has many examples.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

4 excellent

Contribution

4 excellent

Limitations

Yes

Summary

This paper introduces modular adaptive training based on the largest eigenvalue of the NTK. In particular, based on the assumption that the principal eigendirection matters the most for generalization, the authors propose to dynamically turn off gradient updating on certain modules if they have small principal eigenvalues. Experiments are conducted on NLP and CV tasks to show the superiority of MAT.

Strengths

It is very important to reduce computation resources for large model training. This paper focuses on an important topic, and the approach is principled. The presentation is very organized and the delivery is clear.

Weaknesses

See questions.

Questions

1. In Figure 2(a), can we also see a picture showing the distribution of all eigenvalues? Now even the 16th eigenvalue seems pretty large, is there very small eigenvalue? 2. In definition 1, can you write down the size of each matrix? Like $J\in\mathbb{R}^{nk\times p}$. 3. In figure 3, I wonder if there is a phenomenon of benign overfitting in this BERT task. If you keep training beyond 100 epochs to make preplexity even smaller, will you see a double descent on validation? 4. In figure 2(a), $\lambda_1$ of the 12th layer is still large close to 100 epoch. If you keep training will that eigenvalue eventually decays to $0$? What is special about this layer? 5. Around line 220, you mentioned the computation cost of backprop is reduced. Say at some point my first layer is in the information modules but the last layer is in the nuisance modules. In this case, the computation cost is as large as all modules are active, is this right? Because you need to backprop through the later layers. 6. Can you give more intuition on eq(6)? If my $\lambda_1(\Theta_1)$ and $\lambda_1(\Theta_0)$ are small, but $\lambda_1(\Theta_t)$ and $\lambda_1(\Theta_{t-1})$ are both large, then this equation can still be satisfied. Do you want to stop gradient in this scenario? 7. In algorithm 1, how large did you use for sample number $S$? How accurate is this estimation on the eigenvalues? 8. When you compute FLOPs, do you include the computations used for eigendecomposition? 9. Can we see the downstream performance of the MAT-trained BERT on a downstream task?

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

4 excellent

Contribution

3 good

Limitations

N/A

Summary

In this work the authors propose a method for analyzing modules of a neural network, in particular their utility in test-time generalization, via a modular neural tangent kernel (mNTK). The mNTK is eigenspectrum the NTK of the network modules, e.g. different attention heads. In this work it was found that the larger the largest eigenvalue for a module's mNTK, (I will call this lam-max-mNTK for convenience) the more that module contributes features which are useful for generalizability. For example when, during training, a neural network's lam-max-mNTKs plateauing on a low value is indicative that the network has overfit. The authors also demonstrate that the modules' lam-mat-mNTKs evolve asynchronously during training, indicating that optimization of certain modules is more important during different epochs during training. Using this intuition the authors propose modular adaptive training (MAT), which weights modules' learning rates during training in proportion to the lam-max-mNTK of that module. The authors find that this improves training, achieving a minimum loss faster (around 1-5% faster). These experiments, including the analysis of generalizability, were done on BERT, although MAT was also performed on VGG, where the computational savings were even greater.

Strengths

The paper is clearly written and pleasant to read. The proposed analysis is interesting and would be of general interest to the deep learning community. The computational improvements are nice.

Weaknesses

The mNTK isn't super novel, as it is very similar to spectral analysis work which as been performed before, although I do think it is sufficiently novel enough for to ML venues. The performance improvements, while nice, aren't large enough (at least for BERT) that I would expect MAT to see a large deployment in the future considering implementation seems somewhat involved. Perhaps more work is needed to make the improvement more substantial. The experimental results are somewhat limited. In particular its hard to conclude that the observations in Section 2.2 generalize to neural networks in general considering this was just applied to one network. Also I think concluding that two regions in Figure 4 correspond to "info space" and "nuisance space" seems like a bit of a large jump and again, its just for one network. The split is interesting, but can we really conclude that the small eigenvalues are nuisances?

Questions

Do we have some more evidence that your interpretation of mNTK is correct? Is there some reason we are more interested in languange models? Do the VGG lam-max-mNTK results align with those you got for BERT?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

2: You are willing to defend your assessment, but it is quite likely that you did not understand the central parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

4 excellent

Contribution

3 good

Limitations

none

Summary

Inspired by the Neural Tangent Kernel literature, this paper analyzes the training dynamics of over-parameterized neural networks from the perspective of modularity. In particular, the paper proposes to decompose the Jacobians of the NTK into those of modules. Here, the concept of "module" is a bit ill-defined, referring to some coherent computation block. The paper proposes to analyze the NTK within each module and in particular look at the eigenspectrum. It shows that empirically the first eigenvalues dominate within each module and that those eigenvalues are very different across modules. This provides an avenue for focusing mostly on those modules with a high MNTK eigenvalue. Results on different architectures and datasets show noticeable speedups. Finally, a connection between optimizing only high eigenvalue parts of the NTK and generalization is cited and corresponding empirical results are provided in this direction as well.

Strengths

- The paper reads well - It's an interesting topic and potentially impactful in improving our understanding of training dynamics.

Weaknesses

- The biggest limitation for me is that the method of computing the NTK and then computing its eigenvalues seems quite expensive. Speedup experiments are performed using flops, which bypasses the actual running time (which would have to include the algorithm itself), defeating the purpose if it takes longer than the speed gain. - As presented, the algorithm mostly claims speedups by not updating some modules, but if we need to update one of the earliest layers we still need to perform most backpropagation computations. Furthermore, at best, this algorithms provides a 2x (or 1.5x depending on how you count) improvement since it basically only affects the backward step. - I found the connection of eigenvalues of mNTKs and generalization a bit week both conceptually/theoretically and in terms of empirical results. It's enough to show promise, but the connection with the distance to initialization is only tangential and improvements in generalization are not necessarily directly linked to the NTK eigenspectrum in the experiments. Minor suggestion: I would put % of heads instead of number of heads in figure 5 to make it more immediately parseable.

Questions

- I'm a bit surprised at the lack of role of interactions between Jacobians or similar terms across modules in the entire derivation. Given how little constraints there are in what constitutes a module, I would've expected to see some effect. Could you explain this in more detail? - As mentioned under weaknesses, I would really appreciate a deeper comment on the entire computation time when accounting for the algorithm itself.

Rating

4: Borderline reject: Technically solid paper where reasons to reject, e.g., limited evaluation, outweigh reasons to accept, e.g., good evaluation. Please use sparingly.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

2 fair

Presentation

3 good

Contribution

2 fair

Limitations

- The added compute should've been mentioned more clearly in my opinion. - No ethics review necessary.