ACL: Aligned Contrastive Learning Improves BERT and Multi-exit BERT Fine-tuning

TL;DR

This work addresses the conflict between cross-entropy and supervised contrastive learning in BERT fine-tuning by proposing Aligned Contrastive Learning (ACL), which couples label embeddings with sample representations via ACL-Embed and stabilizes optimization with ACL-Grad. It further extends the approach to multi-exit BERT through cross-layer ACL (ACL-CL), using the deepest exit as a teacher to guide shallower exits. Empirical results on GLUE show that ACL improves standard fine-tuning over CE and CE+SCL across most tasks, and ACL-CL delivers substantial gains for multi-exit BERT, achieving better quality-speed tradeoffs without extra pre-training. Altogether, ACL provides a robust, alignment-aware framework for supervised contrastive learning with practical benefits for both single-path and early-exit NLP models.

Abstract

Despite its success in self-supervised learning, contrastive learning is less studied in the supervised setting. In this work, we first use a set of pilot experiments to show that in the supervised setting, the cross-entropy loss objective (CE) and the contrastive learning objective often conflict with each other, thus hindering the applications of CL in supervised settings. To resolve this problem, we introduce a novel \underline{A}ligned \underline{C}ontrastive \underline{L}earning (ACL) framework. First, ACL-Embed regards label embeddings as extra augmented samples with different labels and employs contrastive learning to align the label embeddings with its samples' representations. Second, to facilitate the optimization of ACL-Embed objective combined with the CE loss, we propose ACL-Grad, which will discard the ACL-Embed term if the two objectives are in conflict. To further enhance the performances of intermediate exits of multi-exit BERT, we further propose cross-layer ACL (ACL-CL), which is to ask the teacher exit to guide the optimization of student shallow exits. Extensive experiments on the GLUE benchmark results in the following takeaways: (a) ACL-BRT outperforms or performs comparably with CE and CE+SCL on the GLUE tasks; (b) ACL, especially CL-ACL, significantly surpasses the baseline methods on the fine-tuning of multi-exit BERT, thus providing better quality-speed tradeoffs for low-latency applications.

Figures (6)

• Figure 1: Our proposed ACL framework. ACL differs from the vanilla SCL objective in that it also aligns the semantic representations of input samples with the label embeddings through the contrastive learning objective. Intuitively, ACL considers the label embeddings as anchors to encourage more apparent separations of samples of different labels. Although we show a binary classification case for simplicity, the loss is generally applicable to any multi-class classification setting.
• Figure 2: The distribution of gradient angles after the BERT model is fine-tuned on RTE with the CE objective (in blue) or the CE+SCL objective (in orange).
• Figure 3: The layer-score curves for different multi-exit BERT training methods, with the RoBERTa backbone. TinyBERT is represented as a dot since it is a static compression method.
• Figure 4: The layer-score curves for different multi-exit BERT training methods, with the RoBERTa backbone. The $x$-axis is the depth of the exit (or the number of layers before entering this exit), the $y$-axis is the performance metrics following GLUE Wang2018GLUEAM. The performance of TinyBERT is represented by a dot since it is a static model compression method.
• Figure 5: The layer-score curves for different multi-exit BERT training methods, with the BERT-base backbone.