[Paper Review] Fine-Tuning Language Models via Epistemic Neural Networks

Ian Osband, Seyed Mohammad Asghari+4 more|arXiv (Cornell University)|Nov 3, 2022

Topic Modeling4 citations

TL;DR

This paper proposes augmenting large language models like BERT with an epinet—a small auxiliary network that estimates epistemic uncertainty to prioritize informative training examples during fine-tuning. By using uncertainty-based active learning, the method achieves BERT's full performance on GLUE benchmarks using 2x fewer labeled examples, while also improving final accuracy over standard fine-tuning.

ABSTRACT

Language models often pre-train on large unsupervised text corpora, then fine-tune on additional task-specific data. However, typical fine-tuning schemes do not prioritize the examples that they tune on. We show that, if you can prioritize informative training data, you can achieve better performance while using fewer labels. To do this we augment a language model with an epinet: a small additional network that helps to estimate model uncertainty and forms an extit{epistemic neural network} (ENN). ENNs are neural networks that can know what they don't know. Using an epinet to prioritize uncertain data, we can fine-tune BERT on GLUE tasks to the same performance while using 2x less data than training without prioritization. We also investigate performance in synthetic neural network generative models designed to build understanding. In each setting, using an epinet outperforms heuristic active learning schemes.

Motivation & Objective

• To improve data efficiency in fine-tuning large language models by prioritizing informative training examples.
• To address the limitation of standard fine-tuning, which treats all training data equally without considering uncertainty.
• To enable language models to 'know what they don’t know' by estimating epistemic uncertainty through a learnable epinet.
• To demonstrate that epistemic prioritization outperforms heuristic active learning methods like entropy or margin-based selection.
• To show that epinet-augmented models achieve better final performance and lower data requirements than baseline fine-tuning.

Proposed method

• An epinet is trained alongside the base language model to estimate epistemic uncertainty, distinguishing it from aleatoric uncertainty.
• The epinet uses variance or BALD (Bayesian Active Learning by Disagreement) as uncertainty metrics to prioritize uncertain, informative examples.
• During fine-tuning, the model selects training data based on epistemic uncertainty, focusing on examples where the model is least confident.
• The epinet is trained using a combination of Monte Carlo dropout and model ensembling to improve uncertainty estimation.
• The method is applied to BERT on GLUE benchmark tasks and validated on a synthetic neural network generative model (Neural Testbed).
• Performance is evaluated by measuring data efficiency—how many labels are needed to match baseline performance—using geometric mean across tasks.

Experimental results

Research questions

• RQ1 Can epistemic uncertainty estimation via an epinet improve data efficiency in fine-tuning large language models?
• RQ2 Does prioritizing uncertain examples using epistemic uncertainty lead to better performance with fewer labels compared to heuristic active learning methods?
• RQ3 How does the epinet compare to ensemble-based or dropout-based uncertainty estimation in terms of data efficiency and computational cost?
• RQ4 Can epistemic prioritization improve not only sample efficiency but also final model accuracy?
• RQ5 Is the benefit of the epinet due to the uncertainty estimation or simply the added network capacity?

Key findings

• On GLUE benchmark tasks, fine-tuning BERT with epistemic prioritization achieves the same performance as standard fine-tuning using 2x fewer labeled examples.
• The epinet-based method improves final model accuracy beyond the baseline, even when using less data.
• Epistemic prioritization outperforms heuristic methods like entropy and margin-based selection, which fail to match baseline performance on MNLI and other tasks.
• Using an epinet with uniform data selection performs significantly worse than epistemic prioritization, proving the benefit comes from uncertainty-driven selection, not just architecture.
• Ensemble and dropout-based uncertainty estimation also improve data efficiency, but at higher computational cost than the lightweight epinet.
• The choice between variance and BALD as uncertainty metrics has minimal impact, as both yield similar performance gains.