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

1. El Problema

El artículo aborda una limitación crítica en la aplicación de Modelos de Lenguaje Grande (LLMs) para la toma de decisiones en escenarios de alto riesgo (como la medicina clínica): la incapacidad de proporcionar estimaciones de probabilidad fiables y cuantitativas sin sacrificar su capacidad de generación de texto.

• El Dilema:
  • Los LLMs modernos son excelentes generando explicaciones verbales, pero sus probabilidades internas a menudo no son confiables o están mal calibradas.
  • Los métodos de ajuste fino (fine-tuning) tradicionales para clasificación (usando solo una capa de salida con pérdida de entropía cruzada binaria) suelen provocar un "olvido catastrófico" o un "colapso lingüístico". En este fenómeno, el modelo pierde su capacidad de generar texto coherente y explicaciones, convirtiéndose en un clasificador ciego que no puede justificar sus decisiones.
  • Los métodos de inferencia sin ajuste (como self-consistency o probabilidades verbalizadas) carecen de precisión o son computacionalmente costosos y no ofrecen una estimación de probabilidad directa y robusta.

2. Metodología: CLSGen

Los autores proponen CLSGen, un marco de ajuste fino diseñado específicamente para tareas de clasificación binaria que mantiene simultáneamente la capacidad de estimación de probabilidad y la generación de explicaciones.

Arquitectura y Entrenamiento

El núcleo de CLSGen es una arquitectura de doble cabeza (dual-head) sobre un modelo base de lenguaje (decoder-only):

1. Cabeza de Generación (LM Head): Utiliza el objetivo de modelado de lenguaje causal estándar. Se entrena para generar una secuencia de texto que incluye primero una justificación (explicación) y luego una etiqueta verbalizada (ej. "Clasificación: 1:muerte").
2. Cabeza de Clasificación (CLS Head): Una capa adicional (MLP con función de activación Sigmoid) que toma la representación oculta del último token de entrada (antes de la generación) para predecir directamente la probabilidad binaria.

Función de Pérdida:
El modelo se entrena minimizando una pérdida total combinada:
$$\mathcal{L}_{total} = \mathcal{L}_{gen} + \mathcal{L}_{cls}$$
Donde $\mathcal{L}_{gen}$ es la pérdida de modelado de lenguaje (para mantener la calidad del texto) y $\mathcal{L}_{cls}$ es la pérdida de clasificación (para la precisión probabilística). Se utiliza LoRA (Low-Rank Adaptation) para el ajuste eficiente de los parámetros.

Construcción de Datos

Para entrenar este modelo, se requiere un conjunto de datos enriquecido con explicaciones alineadas con las etiquetas verdaderas.

• Se utiliza un LLM maestro para generar pares de (predicción, explicación) a partir de los datos originales.
• Se emplea una estrategia de muestreo múltiple (ej. $K=5$ intentos): solo se conservan las explicaciones que cumplen con el formato estructurado y cuya predicción verbalizada coincide con la etiqueta real ($y_i$).
• Esto asegura que el modelo aprenda a generar justificaciones lógicas que respalden la etiqueta correcta.

3. Contribuciones Clave

1. Marco de Doble Cabeza: Una arquitectura novedosa que permite la estimación de probabilidades cuantitativas (vía la cabeza de clasificación) sin destruir la capacidad generativa del LLM (vía la cabeza de lenguaje).
2. Prevención del Colapso Lingüístico: Demostración empírica de que el ajuste fino tradicional destruye la capacidad de generación, y propuesta de una solución que preserva la coherencia del texto y la alineación entre la predicción y la explicación.
3. Estrategia de Datos Asistida por LLM: Un método robusto para construir conjuntos de datos de entrenamiento con explicaciones de alta calidad y alineadas con las etiquetas, utilizando validación estricta durante la generación.
4. Métricas de Evaluación Nuevas: Introducción de métricas como AUROC-Alignment (para medir la concordancia entre la probabilidad de la cabeza de clasificación y la etiqueta verbalizada) y el uso de LLM-as-a-Judge para evaluar la consistencia lógica interna.

4. Resultados Experimentales

El modelo CLSGen fue evaluado en dos conjuntos de datos desafiantes:

• LCD Benchmark: Predicción de mortalidad a 30 y 90 días tras el alta hospitalaria (texto clínico largo).
• CliniFact: Verificación de afirmaciones de investigación clínica basadas en resúmenes de literatura.

Hallazgos Principales:

• Rendimiento de Clasificación: CLSGen superó a los métodos de inferencia pura (como self-consistency y predicción de etiquetas) y a los modelos CNN tradicionales en métricas clave como AUROC y F1-score. Por ejemplo, en LCD 30 días, logró un AUROC de 0.8586 frente a 0.7455 de self-consistency.
• Consistencia y Alineación:
  • AUROC-Alignment: Puntuaciones muy altas (0.95 - 0.98), indicando que la probabilidad calculada por la cabeza de clasificación coincide casi perfectamente con la etiqueta generada en el texto.
  • Cohen's Kappa: Mostró un acuerdo sustancial entre la predicción numérica y la explicación textual.
  • Consistencia Lógica (LLM-as-a-Judge): En más del 99.8% de los casos (LCD), la explicación generada era lógicamente consistente con la etiqueta predicha (baja tasa de inconsistencia RLI).
• Calidad de Texto: CLSGen mantuvo una tasa de "outputs parsables" (que siguen el formato) superior al 97%, demostrando que no sufrió colapso lingüístico, a diferencia de los modelos ajustados solo para clasificación.

5. Significado e Impacto

El trabajo de CLSGen es significativo porque resuelve la dicotomía entre precisión cuantitativa y explicabilidad cualitativa en los LLMs aplicados a dominios críticos como la salud.

• Viabilidad Clínica: Permite desplegar LLMs que no solo dicen "el paciente tiene un 85% de riesgo de muerte", sino que también explican por qué (ej. "debido a la hipotensión y la edad avanzada") de manera coherente y confiable.
• Superación de Limitaciones Actuales: Ofrece una alternativa superior a los métodos de inferencia sin entrenamiento (que son inestables) y al ajuste fino tradicional (que destruye la capacidad de generación).
• Interpretabilidad: Al garantizar que la justificación textual y la probabilidad numérica estén alineadas, CLSGen aumenta la confianza del usuario en la toma de decisiones automatizada, un requisito fundamental para la adopción de IA en medicina.

En resumen, CLSGen demuestra que es posible afinar un LLM para tareas de discriminación estricta sin sacrificar su naturaleza generativa, logrando un equilibrio óptimo entre rendimiento estadístico y explicabilidad humana.