CULMFiT: Calibrated ULMFiT for Medical Dialogue

• CULMFiT is a calibrated extension of ULMFiT that integrates label smoothing, temperature scaling, and self-distillation to improve confidence estimates.
• It employs a composite loss and post-hoc recalibration methods to mitigate overconfidence and reduce expected calibration error in clinical dialogue tasks.
• Empirical results show improved BLEU-1, METEOR, and perplexity alongside significant reductions in calibration error across dedicated medical and public dialogue datasets.

CULMFiT is a calibrated extension of Universal LLM Fine-tuning (ULMFiT) designed for medical dialogue systems, integrating label smoothing (LS) in fine-tuning, temperature scaling (TS) for post-hoc calibration, and self-distillation (SD) with optimal temperature selection. The method targets the critical requirement of reliable confidence estimates and robust uncertainty quantification in clinical NLP applications, where conventional deep neural networks often exhibit overconfidence and poor calibration. CULMFiT demonstrates systematic improvements in BLEU-1, METEOR, and perplexity, with a significant reduction in expected calibration error (ECE) relative to baseline and established methods, as validated on both domain-specific and large public medical dialogue datasets (Ao et al., 2021).

1. Motivation for Calibration in Medical Dialogue Systems

Calibration is essential in clinical conversational AI because miscalibrated models tend to be over-confident in wrong predictions—a critical failure mode in healthcare. Properly calibrated models allow for credible uncertainty quantification, support deferral to human experts when confidence is low, and facilitate trustworthy integration into safety-critical workflows. While ULMFiT enhances downstream performance via transfer learning, its softmax probabilities generally require explicit calibration measures before deployment in sensitive domains such as medical diagnostics (Ao et al., 2021).

2. Label Smoothing and CULMFiT Formulation

CULMFiT augments ULMFiT with label smoothing during fine-tuning. The standard cross-entropy (CE) loss for a single example, given ground-truth $y$ (one-hot) and model prediction $\hat{p}$ over $C$ classes, is

$$L_\mathrm{CE} = -\sum_{i=1}^C y^{(i)} \log \hat{p}^{(i)}.$$

Label smoothing replaces one-hot $y^{(i)}$ with a smoothed target $$y^{(i)}_{LS} = (1-\alpha)y^{(i)} + \frac{\alpha}{C}$$, typically with $\alpha=0.1$. The loss becomes

$$L_{LS} = -\sum_{i=1}^C y^{(i)}_{LS} \log \hat{p}^{(i)},$$

and the fine-tuned CULMFiT objective is

$$L_\mathrm{CULMFiT} = -\sum_{i=1}^C y_{LS}^{(i)} \log p_\theta^{(i)}(x),$$

where $p_\theta(x)$ is the model’s softmax output. This regularization directly encourages smoother confidence estimates, mitigating overconfidence inherent in ULMFiT’s original training (Ao et al., 2021).

3. Temperature Scaling and Knowledge Distillation

Temperature scaling (TS) is introduced post-hoc to recalibrate confidence scores. Given model logits $\hat{p}$0,

$\hat{p}$1

where $\hat{p}$2 is the temperature parameter. For knowledge distillation (KD), the softened distribution from a teacher model $\hat{p}$3 is matched to the student model $\hat{p}$4 via a temperature-scaled Kullback–Leibler divergence: $\hat{p}$5 where the factor $\hat{p}$6 maintains gradient scaling for $\hat{p}$7 (Ao et al., 2021).

4. Self-Distillation Protocol and Loss Composition

Self-distillation in CULMFiT involves training the model to convergence, generating teacher logits, and then applying KD with the same model as both teacher and student. Two strategies are adopted:

• Fixed-T SD: Uses a constant $\hat{p}$8 (evaluated over $\hat{p}$9; $C$0 chosen on validation).
• Optimal-T SD: $C$1 is selected by grid search over a held-out set (e.g., $C$2 for the backpain dataset).

The total training objective combines CE loss, LS regularization, and KD: $C$3 with $C$4. This compositional loss ensures contributions from both primary task learning and calibration-focused regularization (Ao et al., 2021).

5. Model Architecture and Optimization

The underlying architecture is an encoder stack of a standard Transformer (6 layers, hidden size $C$5, 8 attention heads) augmented with two fully connected layers for token generation. Pre-training utilizes ULMFiT weights from a large English corpus, followed by fine-tuning on medical dialogue data. Hyperparameters include:

• Label smoothing: $C$6
• Fixed temperature: $C$7
• Optimal TS: $C$8 in $C$9 depending on model/data
• Loss weights: $$L_\mathrm{CE} = -\sum_{i=1}^C y^{(i)} \log \hat{p}^{(i)}.$$0, $$L_\mathrm{CE} = -\sum_{i=1}^C y^{(i)} \log \hat{p}^{(i)}.$$1
• Optimizer and schedule: Adam, learning rate $$L_\mathrm{CE} = -\sum_{i=1}^C y^{(i)} \log \hat{p}^{(i)}.$$2, batch size $$L_\mathrm{CE} = -\sum_{i=1}^C y^{(i)} \log \hat{p}^{(i)}.$$3 (Ao et al., 2021).

6. Datasets, Evaluation Metrics, and Results

Evaluations are conducted on:

• Consultation backpain dataset: 1,000 train and 200 validation patient–clinician pairs, spanning sleep, mental health, exercise, nutrition, and environment determinants.
• MedDialog corpus: 0.8M+ public medical conversations, with an 80/20 train/validation split.

Metrics emphasize both linguistic quality and calibration:

• BLEU-1: Unigram overlap
• METEOR
• Perplexity (uncertainty)
• Expected Calibration Error (ECE): 15-bin, main calibration metric

Key Results

Model	BLEU-1 (Backpain)	ECE (Backpain)	BLEU-1 (MedDialog)	ECE (MedDialog)
Transformer	0.4292	0.370	0.3387	0.261
ULMFiT	0.4321	0.376	0.3609	0.352
CULMFiT (LS)	0.4632	0.367	0.3765	0.373
Fine-tune (TS)	0.4415	0.288	0.3747	0.058
SD + fixed TS ($$L_\mathrm{CE} = -\sum_{i=1}^C y^{(i)} \log \hat{p}^{(i)}.$$4)	0.4236	0.196	-	-
SD + optimal TS ($$L_\mathrm{CE} = -\sum_{i=1}^C y^{(i)} \log \hat{p}^{(i)}.$$5)	0.4473	0.179	-	-

CULMFiT with label smoothing raises BLEU-1 by approximately 3–4% and lowers perplexity compared to both Transformer and ULMFiT. TS fine-tuning achieves a drastic ECE reduction—e.g., from $$L_\mathrm{CE} = -\sum_{i=1}^C y^{(i)} \log \hat{p}^{(i)}.$$60.37 to 0.06 on MedDialog—at minimal cost to accuracy. Self-distillation with optimal $$L_\mathrm{CE} = -\sum_{i=1}^C y^{(i)} \log \hat{p}^{(i)}.$$7 attains the lowest ECE while retaining most accuracy gains (Ao et al., 2021).

7. Significance and Implications

Integration of label smoothing, temperature scaling, and self-distillation in CULMFiT systematically improves robustness, linguistic quality, and especially reliability of model-generated confidence estimates. Reliable calibration is imperative for practical deployment in clinical dialogue systems, given the risks associated with overconfidence in incorrect predictions. A plausible implication is that the CULMFiT methodology—explicitly targeting calibration with multiple complementary strategies—is a broadly applicable paradigm for other safety-critical NLP applications requiring trustworthy confidence and uncertainty quantification (Ao et al., 2021).