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

TL;DR

This paper introduces the use of Conditional Mutual Information (CMI) in fine-tuning large language models for classification, achieving performance improvements on multiple GLUE tasks by minimizing or maximizing CMI during training and knowledge distillation.

Contribution

It adapts the CMI-constrained deep learning framework for LLM fine-tuning, demonstrating its effectiveness in improving both standalone models and student models in classification tasks.

Findings

Minimizing CMI improves BERT performance on 6 of 8 GLUE tasks.

Maximizing CMI enhances DistilBERT performance on 6 of 8 GLUE tasks.

CMI-based fine-tuning outperforms baseline methods in several classification benchmarks.

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.