Papers
Topics
Authors
Recent
Gemini 2.5 Flash
Gemini 2.5 Flash
119 tokens/sec
GPT-4o
56 tokens/sec
Gemini 2.5 Pro Pro
43 tokens/sec
o3 Pro
6 tokens/sec
GPT-4.1 Pro
47 tokens/sec
DeepSeek R1 via Azure Pro
28 tokens/sec
2000 character limit reached

Understanding Linear Probing then Fine-tuning Language Models from NTK Perspective (2405.16747v1)

Published 27 May 2024 in cs.LG

Abstract: The two-stage fine-tuning (FT) method, linear probing then fine-tuning (LP-FT), consistently outperforms linear probing (LP) and FT alone in terms of accuracy for both in-distribution (ID) and out-of-distribution (OOD) data. This success is largely attributed to the preservation of pre-trained features, achieved through a near-optimal linear head obtained during LP. However, despite the widespread use of LLMs, the exploration of complex architectures such as Transformers remains limited. In this paper, we analyze the training dynamics of LP-FT for classification models on the basis of the neural tangent kernel (NTK) theory. Our analysis decomposes the NTK matrix into two components, highlighting the importance of the linear head norm alongside the prediction accuracy at the start of the FT stage. We also observe a significant increase in the linear head norm during LP, stemming from training with the cross-entropy (CE) loss, which effectively minimizes feature changes. Furthermore, we find that this increased norm can adversely affect model calibration, a challenge that can be addressed by temperature scaling. Additionally, we extend our analysis with the NTK to the low-rank adaptation (LoRA) method and validate its effectiveness. Our experiments with a Transformer-based model on natural language processing tasks across multiple benchmarks confirm our theoretical analysis and demonstrate the effectiveness of LP-FT in fine-tuning LLMs. Code is available at https://github.com/tom4649/lp-ft_ntk.

Definition Search Book Streamline Icon: https://streamlinehq.com
References (60)
  1. Intrinsic dimensionality explains the effectiveness of language model fine-tuning. In Chengqing Zong, Fei Xia, Wenjie Li, and Roberto Navigli, editors, Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Language Processing (Volume 1: Long Papers), pages 7319–7328, Online, August 2021. Association for Computational Linguistics. doi: 10.18653/v1/2021.acl-long.568. URL https://aclanthology.org/2021.acl-long.568.
  2. On exact computation with an infinitely wide neural net. Advances in neural information processing systems, 32, 2019.
  3. The second pascal recognising textual entailment challenge. Proceedings of the Second PASCAL Challenges Workshop on Recognising Textual Entailment, 01 2006.
  4. The fifth pascal recognizing textual entailment challenge. TAC, 7(8):1, 2009.
  5. Proceedings of the 2019 conference of the north american chapter of the association for computational linguistics: Human language technologies, volume 1 (long and short papers). In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers), 2019.
  6. Revisiting parameter-efficient tuning: Are we really there yet? In Yoav Goldberg, Zornitsa Kozareva, and Yue Zhang, editors, Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, pages 2612–2626, Abu Dhabi, United Arab Emirates, December 2022. Association for Computational Linguistics. doi: 10.18653/v1/2022.emnlp-main.168. URL https://aclanthology.org/2022.emnlp-main.168.
  7. BoolQ: Exploring the surprising difficulty of natural yes/no questions. In Jill Burstein, Christy Doran, and Thamar Solorio, editors, Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers), pages 2924–2936, Minneapolis, Minnesota, June 2019. Association for Computational Linguistics. doi: 10.18653/v1/N19-1300. URL https://aclanthology.org/N19-1300.
  8. The pascal recognising textual entailment challenge. In Machine learning challenges workshop, pages 177–190. Springer, 2005.
  9. The commitmentbank: Investigating projection in naturally occurring discourse. In proceedings of Sinn und Bedeutung, volume 23, pages 107–124, 2019.
  10. Pubmed 200k rct: a dataset for sequential sentence classification in medical abstracts, 2017.
  11. Automatically constructing a corpus of sentential paraphrases. In Third international workshop on paraphrasing (IWP2005), 2005.
  12. Ronald A Fisher. The use of multiple measurements in taxonomic problems. Annals of eugenics, 7(2):179–188, 1936.
  13. Deep learning versus kernel learning: an empirical study of loss landscape geometry and the time evolution of the neural tangent kernel, 2020.
  14. The third pascal recognizing textual entailment challenge. In Proceedings of the ACL-PASCAL workshop on textual entailment and paraphrasing, pages 1–9, 2007.
  15. On calibration of modern neural networks. In Doina Precup and Yee Whye Teh, editors, Proceedings of the 34th International Conference on Machine Learning, volume 70 of Proceedings of Machine Learning Research, pages 1321–1330. PMLR, 06–11 Aug 2017. URL https://proceedings.mlr.press/v70/guo17a.html.
  16. Domain-aware fine-tuning: Enhancing neural network adaptability. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 38, pages 12261–12269, 2024.
  17. Towards a unified view of parameter-efficient transfer learning. In International Conference on Learning Representations, 2022. URL https://openreview.net/forum?id=0RDcd5Axok.
  18. Richard Zou Horace He. functorch: Jax-like composable function transforms for pytorch. https://github.com/pytorch/functorch, 2021.
  19. Parameter-efficient transfer learning for NLP. In Kamalika Chaudhuri and Ruslan Salakhutdinov, editors, Proceedings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pages 2790–2799. PMLR, 09–15 Jun 2019. URL https://proceedings.mlr.press/v97/houlsby19a.html.
  20. LoRA: Low-rank adaptation of large language models. In International Conference on Learning Representations, 2022. URL https://openreview.net/forum?id=nZeVKeeFYf9.
  21. Neural tangent kernel: Convergence and generalization in neural networks. Advances in neural information processing systems, 31, 2018.
  22. Lora training in the ntk regime has no spurious local minima, 2024.
  23. Smart: Robust and efficient fine-tuning for pre-trained natural language models through principled regularized optimization. arXiv preprint arXiv:1911.03437, 2019.
  24. Adjusting decision boundary for class imbalanced learning. IEEE Access, 8:81674–81685, 2020.
  25. Adam: A method for stochastic optimization, 2017.
  26. Last layer re-training is sufficient for robustness to spurious correlations. In The Eleventh International Conference on Learning Representations, 2023. URL https://openreview.net/forum?id=Zb6c8A-Fghk.
  27. Do better imagenet models transfer better? In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 2661–2671, 2019.
  28. Fine-tuning can distort pretrained features and underperform out-of-distribution. arXiv preprint arXiv:2202.10054, 2022.
  29. Wide neural networks of any depth evolve as linear models under gradient descent. Advances in neural information processing systems, 32, 2019.
  30. Surgical fine-tuning improves adaptation to distribution shifts. In The Eleventh International Conference on Learning Representations, 2023. URL https://openreview.net/forum?id=APuPRxjHvZ.
  31. The winograd schema challenge. In Thirteenth international conference on the principles of knowledge representation and reasoning, 2012.
  32. Rethinking the hyperparameters for fine-tuning. In International Conference on Learning Representations, 2020. URL https://openreview.net/forum?id=B1g8VkHFPH.
  33. Ro{bert}a: A robustly optimized {bert} pretraining approach, 2020. URL https://openreview.net/forum?id=SyxS0T4tvS.
  34. A kernel-based view of language model fine-tuning. In International Conference on Machine Learning, pages 23610–23641. PMLR, 2023.
  35. Fine-tuning language models with just forward passes. Advances in Neural Information Processing Systems, 36, 2024.
  36. Hidden factors and hidden topics: Understanding rating dimensions with review text. In Proceedings of ACM Conference on Recommender Systems, 2013.
  37. A fast, well-founded approximation to the empirical neural tangent kernel, 2023. URL https://openreview.net/forum?id=HN0ehX-ov5Q.
  38. Obtaining well calibrated probabilities using bayesian binning. In Proceedings of the AAAI conference on artificial intelligence, volume 29, 2015.
  39. SemEval-2016 task 4: Sentiment analysis in Twitter. In Proceedings of International Workshop on Semantic Evaluation (SemEval), 2016.
  40. Fast finite width neural tangent kernel. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato, editors, Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pages 17018–17044. PMLR, 17–23 Jul 2022. URL https://proceedings.mlr.press/v162/novak22a.html.
  41. Pytorch: An imperative style, high-performance deep learning library, 2019.
  42. Scikit-learn: Machine learning in Python. Journal of Machine Learning Research, 12:2825–2830, 2011.
  43. AdapterHub: A framework for adapting transformers. In Qun Liu and David Schlangen, editors, Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing: System Demonstrations, pages 46–54, Online, October 2020. Association for Computational Linguistics. doi: 10.18653/v1/2020.emnlp-demos.7. URL https://aclanthology.org/2020.emnlp-demos.7.
  44. DynaSent: A dynamic benchmark for sentiment analysis. In Proceedings of ACL-IJCNLP, 2021.
  45. Better supervisory signals by observing learning paths, 2022.
  46. How to prepare your task head for finetuning. In The Eleventh International Conference on Learning Representations, 2023. URL https://openreview.net/forum?id=gVOXZproe-e.
  47. Recursive deep models for semantic compositionality over a sentiment treebank. In Proceedings of EMNLP, 2013.
  48. The implicit bias of gradient descent on separable data. In International Conference on Learning Representations, 2018. URL https://openreview.net/forum?id=r1q7n9gAb.
  49. On the theory of transfer learning: The importance of task diversity. Advances in neural information processing systems, 33:7852–7862, 2020.
  50. A closer look at model adaptation using feature distortion and simplicity bias. In The Eleventh International Conference on Learning Representations, 2023. URL https://openreview.net/forum?id=wkg_b4-IwTZ.
  51. Attention is all you need. In I. Guyon, U. Von Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017. URL https://proceedings.neurips.cc/paper_files/paper/2017/file/3f5ee243547dee91fbd053c1c4a845aa-Paper.pdf.
  52. Glue: A multi-task benchmark and analysis platform for natural language understanding. arXiv preprint arXiv:1804.07461, 2018.
  53. Superglue: A stickier benchmark for general-purpose language understanding systems. Advances in neural information processing systems, 32, 2019.
  54. Neural network acceptability judgments. Transactions of the Association for Computational Linguistics, 7:625–641, 2019. doi: 10.1162/tacl_a_00290. URL https://aclanthology.org/Q19-1040.
  55. More than a toy: Random matrix models predict how real-world neural representations generalize. In Proceedings of the 39th International Conference on Machine Learning, 2022.
  56. Transformers: State-of-the-art natural language processing. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing: System Demonstrations, pages 38–45, Online, October 2020. Association for Computational Linguistics. URL https://www.aclweb.org/anthology/2020.emnlp-demos.6.
  57. Revisiting out-of-distribution robustness in nlp: Benchmark, analysis, and llms evaluations. arXiv preprint arXiv:2306.04618, 2023.
  58. The expressive power of low-rank adaptation. In The Twelfth International Conference on Learning Representations, 2024. URL https://openreview.net/forum?id=likXVjmh3E.
  59. Side-tuning: a baseline for network adaptation via additive side networks. In Computer Vision–ECCV 2020: 16th European Conference, Glasgow, UK, August 23–28, 2020, Proceedings, Part III 16, pages 698–714. Springer, 2020.
  60. A comprehensive survey on transfer learning. Proceedings of the IEEE, 109(1):43–76, 2020.
User Edit Pencil Streamline Icon: https://streamlinehq.com
Authors (2)
  1. Akiyoshi Tomihari (3 papers)
  2. Issei Sato (82 papers)
Citations (1)
View on Semantic Scholar

Summary

Neural Tangent Kernel-Based Analysis of Linear Probing then Fine-Tuning for Classification Models

The paper presents an in-depth analysis of the two-stage fine-tuning (FT) method, linear probing then fine-tuning (LP-FT), for classification models, particularly focusing on its application to complex architectures like Transformers. Using neural tangent kernel (NTK) theory, the authors dissect the training dynamics of LP-FT, elucidating the mechanisms that contribute to its observed superior performance over regular FT and linear probing (LP) alone, especially for both in-distribution (ID) and out-of-distribution (OOD) data.

LP-FT first optimizes the linear head of the model during the linear probing stage (LP), then proceeds to the fine-tuning stage (FT) to update the entire model, including the feature extractor and the linear head. This method is shown to reduce feature distortion, a property originally attributed to the feature distortion theory validated for a simpler two-layer linear model. The present paper extends this theoretical understanding to more complex models such as Transformers using the NTK framework, which offers a first-order approximation of changes in model outputs with respect to its parameters.

Key Theoretical Contributions

  1. NTK Decomposition: The NTK matrix is decomposed into two components:
    • Pre-train-effective component: Persistence of pre-trained features, approximated by the inner product of pre-trained features.
    • FT-effective component: Captures dynamic feature changes due to fine-tuning, derived from derivatives of the feature extractor.
  2. Classifier Weight Norm: The authors identify that the norm of the linear head increases significantly during LP. This increase optimizes the classifier for minor feature changes and magnifies the effect of the FT stage. Larger classifier weight norms make LP-FT more efficient in preserving pre-trained features, reducing feature distortion.
  3. Temperature Scaling: An increased norm during LP can adversely affect model calibration, making predicted probabilities overly confident. Applying temperature scaling during inference corrects overconfidence, adjusting predicted probability distributions closer to actual probabilities.
  4. Extension to LoRA: The paper extends its NTK-based analysis to the low-rank adaptation (LoRA) method. It theoretically verifies that LoRA approximates standard FT effectively in terms of the NTK components, ensuring parameter-efficient fine-tuning without compromising performance.

Empirical Validation

Experiments conducted with Transformer-based models validate the theoretical insights across multiple benchmarks, including GLUE, SuperGLUE, and BOSS:

  • Reduced Feature Changes: Empirical results show LP-FT leads to smaller deviations in feature space compared to standard FT. The t-SNE visualizations of features confirm this, showing that the features retain their structure post LP-FT.
  • Performance Metrics: LP-FT consistently exhibits better ID and OOD performance compared to FT and even LoRA in several tasks.
  • NTK Analysis: Kernel regression accuracy and singular value distribution analyses of the NTK matrices underscore the efficacy of the FT-effective component in LP-FT.

Practical Implications and Future Directions

The findings underline the benefits of LP-FT in maintaining robust feature representations, critical for the model's generalization capabilities. Practically, LP-FT's effectiveness in reducing feature distortion makes it a valuable technique for scenarios requiring robust performance across varied and unseen data distributions.

Further exploration may delve into optimizing the initial norm increase during LP to balance calibration and feature retention. Moreover, extending this analysis to other complex model architectures beyond Transformers and different learning paradigms may yield broader insights on the adaptability and effectiveness of LP-FT.

In conclusion, this paper articulates a nuanced understanding of LP-FT through NTK theory, providing a substantial theoretical backing to its empirical successes. Its insights into feature preservation, classifier norm effects, and calibration adjustments offer a robust framework for enhancing and utilizing pre-trained models in diverse and dynamic settings.

File Document Download Save Streamline Icon: https://streamlinehq.com PDF File Document Download Save Streamline Icon: https://streamlinehq.com Markdown