Leveraging Conditional Mutual Information to Improve Large Language Model Fine-Tuning For Classification

TL;DR

This work reframes transformer-based LLMs as high-order Markov source models and integrates Conditional Mutual Information (CMI) into fine-tuning for classification. It introduces FPDCs and CMI-based objectives to either minimize CMI for improved standalone performance or maximize CMI to enhance knowledge distillation to a student model. Empirical results on GLUE show that minimizing CMI yields gains for BERT-base across multiple tasks, while maximizing CMI for teacher models improves KD outcomes for DistilBERT on a majority of tasks, with notable gains on WNLI and QQP. Overall, the study demonstrates a principled, information-theoretic pathway to strengthen both single-model performance and knowledge transfer in LLM contexts, suggesting broader applicability beyond classification.

Abstract

Although large language models (LLMs) have demonstrated remarkable capabilities in recent years, the potential of information theory (IT) to enhance LLM development remains underexplored. This paper introduces the information theoretic principle of Conditional Mutual Information (CMI) to LLM fine-tuning for classification tasks, exploring its promise in two main ways: minimizing CMI to improve a model's standalone performance and maximizing CMI to enhance knowledge distillation (KD) for more capable student models. To apply CMI in LLM fine-tuning, we adapt the recently proposed CMI-constrained deep learning framework, which was initially developed for image classification, with some modification. By minimizing CMI during LLM fine-tuning, we achieve superior performance gains on 6 of 8 GLUE classification tasks compared to BERT. Additionally, maximizing CMI during the KD process results in significant performance improvements in 6 of 8 GLUE classification tasks compared to DistilBERT. These findings demonstrate CMI's adaptability for optimizing both standalone LLMs and student models, showcasing its potential as a robust framework for advancing LLM fine-tuning. Our work bridges the gap between information theory and LLM development, offering new insights for building high-performing language models.

Figures (2)

• Figure 1: Relationship between LLMs and IT.
• Figure 2: An illustration of a transformer-based LLM architecture (GPT-1 architecture GPT1).