💬 NLP

CLSGen: A Dual-Head Fine-Tuning Framework for Joint Probabilistic Classification and Verbalized Explanation

The paper introduces CLSGen, a novel dual-head fine-tuning framework that enables Large Language Models to generate reliable quantitative probabilities for binary classification while preserving their ability to produce coherent verbalized explanations, thereby overcoming the limitations of catastrophic forgetting and linguistic collapse found in traditional discriminative approaches.

Original authors: WonJin Yoon, Kangyu Zhu, Ian Bulovic, Autumn Sehy, Yanjun Gao, Dmitriy Dligach, Majid Afshar, Timothy A. Miller

Published 2026-04-14

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: WonJin Yoon, Kangyu Zhu, Ian Bulovic, Autumn Sehy, Yanjun Gao, Dmitriy Dligach, Majid Afshar, Timothy A. Miller

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you have a brilliant, highly educated doctor (the Large Language Model, or LLM) who can read thousands of medical pages and write a beautiful, detailed story about a patient's condition. This doctor is great at explaining why they think a patient is sick.

However, there's a problem: when you ask this doctor, "How sure are you that this patient will survive?" they can't give you a number. They might say, "I feel pretty good about it," but they can't say, "I am 85% sure."

If you try to force this doctor to become a statistician and give you exact numbers, something strange happens. They start forgetting how to speak properly. Their sentences turn into gibberish like "Pom Pomuppy Pom Pom." They lose their ability to explain things because they are so focused on the math that they forget how to write. This is what the paper calls "Linguistic Collapse."

The authors of this paper, CLSGen, wanted to fix this. They created a new training method that teaches the doctor to do two things at once without losing their mind:

Give a precise probability (a number).
Write a clear, logical explanation (a story).

Here is how they did it, using some simple analogies:

1. The Problem: The "One-Track Mind"

Think of the LLM as a talented actor who is great at improvising a play. If you tell them, "Just memorize these facts and recite them," they might get so focused on the facts that they forget how to act. They stop speaking in sentences and start spitting out random words.

The Paper's Discovery: When they tried to train the model just to guess "Yes" or "No" (like a simple math problem), the model's ability to write good explanations vanished almost immediately.

2. The Solution: The "Dual-Head" Hat

The authors built a special training framework called CLSGen. Imagine the model is wearing a special hat with two heads:

Head A (The Statistician): This head looks at the patient's notes and calculates a number (e.g., "85% chance of survival").
Head B (The Storyteller): This head looks at the same notes and writes a paragraph explaining why that number is what it is.

The magic of CLSGen is that it trains both heads together. It doesn't let the Statistician head take over and silence the Storyteller. They work as a team. If the Storyteller writes a confusing reason, the Statistician gets a "ding" (a penalty), and vice versa. This keeps the model balanced.

3. The Training Data: The "Quality Control" Filter

To teach this dual-headed model, they needed examples where the number and the story matched perfectly.

The Process: They used a super-smart AI to generate thousands of practice cases.
The Filter: They set up a rule: "If the AI writes a story that says 'The patient will die' but the number it gives is '100% chance of living,' throw that example in the trash."
They only kept the perfect examples where the story and the number agreed. This ensured the model learned that the explanation must match the prediction.

4. The Results: The Best of Both Worlds

When they tested this new method on real medical data (predicting if a patient would survive 30 or 90 days after leaving the hospital), it worked wonders:

Accuracy: It was better at guessing the right outcome than previous methods.
Reliability: It gave a clear number (probability) that actually matched the story it told.
Clarity: The explanations were readable and made sense. The model didn't turn into gibberish; it stayed a good writer.

Why Does This Matter?

In high-stakes fields like medicine or law, you can't just have a "black box" that says "Yes" or "No." You need to know how sure the system is (the probability) and why it thinks that (the explanation).

Before this paper, you had to choose: either a smart explainer with no numbers, or a number-cruncher that couldn't speak. CLSGen gives you both, ensuring that the AI remains a helpful, transparent partner rather than a confusing, broken machine.

1. Problem Statement

The paper addresses a critical bottleneck in deploying Large Language Models (LLMs) for high-stakes decision-making (e.g., clinical prognosis, drug discovery): the inability to provide reliable, quantitative probability estimates without sacrificing generative capabilities.

The Trade-off: Traditional fine-tuning of decoder-only LLMs for classification (using a standard Sigmoid layer and Binary Cross-Entropy loss) often leads to catastrophic forgetting or linguistic collapse. In this phenomenon, the model loses its ability to generate coherent text, rendering it useless for providing verbalized explanations, which are essential for interpretability.
Limitations of Alternatives:
- Inference-only methods (e.g., prompting for verbalized probabilities or self-consistency) lack the reliability and calibration of fine-tuned models.
- Standard fine-tuning destroys the model's generative "reasoning" ability.
- Reinforcement Learning (RL) approaches (RLHF, DPO) are complex, unstable, and do not inherently solve the joint optimization of probability and text generation.

2. Methodology: CLSGen Framework

The authors propose CLSGen, a novel framework designed to jointly optimize for binary classification accuracy and high-quality text generation (explanations). The framework consists of three core components:

A. Dual-Head Architecture

Instead of replacing the language modeling head, CLSGen utilizes a decoder-only model with two parallel output heads sharing a common backbone:

Generation Head (LM Head): A standard causal language modeling head trained to generate a structured sequence containing a textual rationale (explanation) followed by a verbalized label (e.g., "Reasoning: ... Classification: 1").
Classification Head (CLS Head): A specialized head that extracts the hidden representation ( $h_n$ ) of the last token of the input prefix. This vector is passed through a two-layer MLP with a Sigmoid activation to output a direct binary probability.

Training Objective: The model is trained to minimize a combined loss function:
$\mathcal{L}_{total} = \mathcal{L}_{gen} + \mathcal{L}_{cls}$

$\mathcal{L}_{gen}$ : Causal language modeling loss (Teacher Forcing) for generating the explanation and label.
$\mathcal{L}_{cls}$ : Binary Cross-Entropy loss for the classification head.
Optimization: Both the Classification Head and the LoRA (Low-Rank Adaptation) parameters of the backbone are optimized end-to-end, allowing gradients from both objectives to shape the shared representations.

B. Data Construction Strategy

To ensure the training data aligns the generated explanations with ground-truth labels, the authors employ a weakly supervised, iterative data generation strategy:

An LLM ( $M_{data}$ ) is prompted to generate a prediction and an explanation for a given input.
The output is parsed. If the generated label matches the ground truth ( $y_i$ ) and the explanation is valid, the sample is retained.
If the output is invalid or the label is wrong, the process is retried (up to $K$ times). Only the first valid, correct instance is added to the augmented training set.
This ensures the model learns to generate explanations that logically support the correct classification.

C. Training Setup

Base Model: Llama-3.1-8B-Instruct.
Technique: LoRA for efficient fine-tuning.
Datasets:
- LCD Benchmark: Predicting 30/90-day mortality from discharge notes (long-context, low prevalence).
- CliniFact: Verifying clinical research claims against literature abstracts.

3. Key Contributions

Novel Architecture: Introduction of a dual-head fine-tuning framework that prevents linguistic collapse by maintaining the generative head while adding a dedicated classification head.
Joint Optimization: Demonstrates that classification probability and verbalized explanations can be learned simultaneously without the instability of Reinforcement Learning.
New Evaluation Metrics:
- AUROC-Alignment: Measures the correlation between the classification head's probability and the verbalized label in the generated text.
- Rationale-Label Inconsistency (RLI): Uses an LLM-as-a-judge to verify if the generated explanation logically implies the predicted label.
Empirical Validation of "Linguistic Collapse": The paper provides preliminary evidence showing that standard classification-only fine-tuning causes generative capabilities to vanish (parsability drops to near zero) after a few epochs, justifying the need for CLSGen.

4. Experimental Results

The framework was evaluated against inference-only baselines (Verbalized Probability, Label Prediction, Self-Consistency) and a traditional CNN baseline.

Classification Performance:
- CLSGen achieved the highest AUROC across all datasets (e.g., 0.8586 on LCD 30-day vs. 0.7455 for Self-Consistency and 0.8582 for CNN).
- It outperformed baselines in F1-score (when using optimized thresholds) and maintained superior parsability (>97%), proving it retains structural constraints.
Consistency & Alignment:
- AUROC-Alignment: Scores ranged from 0.95 to 0.98, indicating the classification head's probability is highly consistent with the verbalized label in the text.
- Internal Consistency: The Rationale-Label Inconsistency (RLI) was extremely low (<0.5% for LCD, <5% for CliniFact), meaning the generated explanations almost always logically matched the predicted label.
- Readability: LLM-as-a-judge scores indicated high readability and clinical utility.

5. Significance and Conclusion

CLSGen offers a practical solution for deploying LLMs in high-stakes domains where both a confidence score and a human-readable justification are required.

Solves the "Black Box" Problem: It provides not just a prediction, but a reliable probability and a coherent explanation, enhancing trust and interpretability.
Prevents Capability Loss: By avoiding the "linguistic collapse" associated with standard fine-tuning, it preserves the LLM's core strength: text generation.
Scalability: The framework is compatible with existing decoder-only models and uses efficient fine-tuning techniques (LoRA), making it feasible for resource-constrained environments.

The authors conclude that CLSGen successfully bridges the gap between discriminative performance (probability estimation) and generative capabilities (explanation), setting a new standard for fine-tuning LLMs in clinical and scientific applications.