Conditioning Tree-Based Diffusions and Flows for Probabilistic Tabular Regression

Resumen Técnico: Condicionamiento de Difusiones y Flujos Basados en Árboles para la Regresión Tabular Probabilística

Planteamiento del Problema
La regresión probabilística en datos tabulares requiere modelos que combinen la fuerza predictiva de los árboles de gradiente aumentado con distribuciones predictivas completas adecuadas para la calibración, la toma de decisiones y la cuantificación de la incertidumbre. Aunque los modelos de difusión basados en árboles (por ejemplo, Treeffuser) han surgido como una solución al entrenar regresores LightGBM para estimar el score de una densidad condicional perturbada por ruido, actualmente heredan valores predeterminados de diseño de los entornos de difusión neuronal. Estos valores predeterminados —que incluyen la trayectoria de ruido, la parametrización, la distribución de entrenamiento, las características y el muestreador— son a menudo subóptimos para los aprendices de árboles basados en histogramas. El artículo argumenta que estos valores predeterminados actúan como una restricción vinculante: el problema específico de regresión supervisada que resuelve un ensamble de gradiente aumentado está determinado por cómo se configuran estas elecciones de condicionamiento.

Metodología: DiffGBM
El autor introduce DiffGBM, una extensión de Treeffuser que trata las elecciones de condicionamiento de la difusión basada en árboles como ejes explícitos y ajustables en lugar de valores predeterminados fijos. La metodología se desarrolla a lo largo de dos ejes principales:

1. Modificaciones del Lado del Score (El Espacio "Score-Flex"):
   El autor propone una receta conjuntamente ajustable para la difusión basada en score sobre una superficie LightGBM compartida. Las modificaciones clave incluyen:

   • Residualización de la Media Condicional: Primero se ajusta un predictor de media $\hat{\mu}(x)$ mediante validación cruzada con LightGBM. La difusión se aplica al residuo $y - \hat{\mu}(x)$ en lugar de a $y$, desplazando la ubicación condicional fuera del objetivo de regresión del score.
   • Precondicionamiento EDM: Siguiendo los principios de EDM, el modelo predice un objetivo de denoisado $D_\theta(y_t, t)$ escalado a varianza unitaria, reconstruyendo el score algebraicamente. Esto estabiliza el objetivo de regresión para los árboles, evitando la necesidad de extrapolar magnitudes de predicción de ruido a través de distintos niveles de ruido.
   • Muestreo de Tiempo y Características de Log-Sigma: En lugar de un muestreo de tiempo uniforme, la distribución de entrenamiento extrae $\log \sigma(t) \sim \mathcal{N}(-1.2, 1.2^2)$. Crucialmente, se añade una característica explícita de $\log \sigma(t)$ al vector de entrada, permitiendo que las divisiones (splits) de los árboles se condicionen directamente al régimen de ruido.
   • Ponderación de Pérdida y Resolución de Histograma: Estos se exponen como dimensiones ajustables para controlar dónde se emplea la capacidad del modelo.

2. Flujo de Coincidencia de Gauss (Flow Matching - FM):
   El autor implementa un entrenador de coincidencia de flujo (flow-matching) de trayectoria gaussiana para $p(y|x)$.

   • Aprendizaje del Campo de Velocidad: En lugar de aprender un score directamente, el modelo aprende un campo de velocidad $v_\theta$ mediante árboles de gradiente aumentado.
   • Recuperación Algebraica del Score: Para cualquier trayectoria gaussiana, el score condicional puede recuperarse algebraicamente a partir de la velocidad (Lema 1), lo que permite al modelo entrar en muestreadores de interpolante estocástico si se desea.
   • Eficiencia del Muestreador: Este enfoque permite el uso de muestreadores de ODE deterministas de pocos pasos (específicamente un ODE Heun de 5 pasos), lo que reduce significativamente la latencia de inferencia en comparación con los SDE estocásticos.
   • Elección de Trayectorias: El autor compara trayectorias lineales, trigonométricas y de preservación de varianza (VP). Observa que, aunque las trayectorias VP mantienen las entradas cerca de la escala unitaria, la trayectoria óptima depende del conjunto de datos y de la interacción con la residualización.

Contribuciones Clave

1. Coincidencia de Flujo para la Predicción Condicional Basada en Árboles: Un entrenador de FM de trayectoria gaussiana que aprende un campo de velocidad directamente con árboles de gradiente aumentado, ofreciendo el muestreador más económico y el punto de operación mejor calibrado entre las variantes de DiffGBM.
2. Una Receta de Score Conjuntamente Ajustable: La exposición de las modificaciones del lado del score (residualización, precondicionamiento EDM, muestreo de log-sigma, características de ruido explícitas, ponderación de pérdida y resolución de histograma) como ejes ajustables. Este espacio de "score-flex" representa la receta publicada como un caso especial, pero permite la optimización específica para cada conjunto de datos.
3. El Condicionamiento de la Regresión como el Cuello de Botella Compartido: El artículo postula que las ganancias de rendimiento en ambos ejes provienen de un mejor ajuste del problema supervisado inducido a la resolución finita de división de los histogramas. La restricción no son las propiedades a nivel de población de los objetivos, sino la asignación finita de bins requerida para servir todos los niveles de ruido simultáneamente.

Resultados Experimentales
El método fue evaluado en once bancos de pruebas estándar de regresión tabular (nueve de UCI, dos de scikit-learn) utilizando un protocolo de ajuste anidado (fold 0 para ajuste, folds 1–5 para evaluación) con un presupuesto igualado de 40 ensayos.

• Precisión (Score-Flex): La receta de score-flex ajustada conjuntamente logró el mejor puntaje de habilidad CRPS agregado (0.725 frente a 0.699 de la línea base publicada) y CRPS normalizado (1.106 frente a 1.248). Superó a la línea base publicada en cada uno de los conjuntos de datos (Wilcoxon pareado $p < 10^{-3}$).
• Calibración y Velocidad (Flow Matching): La variante DiffGBM-FM fue la fila de DiffGBM mejor calibrada (menor error de cobertura en los niveles 50/90% y mayor tasa de paso PIT) y el muestreador más económico. Fue 5.2× más rápida que la línea base publicada (61.3s total frente a 196.3s) manteniendo un CRPS competitivo (0.707).
• Estudios de Ablación:
  • Cambiar de Euler SDE a Heun ODE para el modelo de score sin otros cambios redujo la precisión, indicando que el muestreador no es una mejora de velocidad gratuita.
  • El precondicionamiento EDM proporcionó el mayor salto único en el rendimiento del lado del score.
  • La residualización fue crítica para el Flow Matching; sin ella, el rendimiento de FM se degradó significamente.
  • Los muestreadores estocásticos ($\epsilon > 0$) no dominaron de forma pareto al modelo de los deterministas ($\epsilon = 0$) en el CRPS agregado o en la latencia, aunque pudieron mejorar la cobertura de la cola en regiones de alta incertidumbre.

Significancia y Reivindicaciones
El artículo afirma que los valores predeterminados de diseño heredados de los entornos de difusión neuronal son el principal cuello de botella para la regresión tabular probabilística. Al tratar la trayectoria de ruido, la parametrización y la distribución de entrenamiento como elecciones de condicionamiento ajustables, DiffGBM desbloquea mejoras significativas tanto en precisión como en calibración.

El autor enfatiza que los enfoques de "score-flex" y "flow-matching" son complementarios:

• Score-Flex es la elección de precisión, que domina a medida que los datos crecen y gana en los conjuntos de datos más grandes, aunque incurre en mayor latencia y una cobertura central más laxa debido a su muestreador estocástico.
• Flow Matching es la elección rápida y bien calibrada, que gana en conjuntos de datos pequeños a medianos donde su ODE de pocos pasos residualizado es más fuerte.

El artículo concluye que ninguna receta única es universal; la configuración óptima depende de las características del conjunto de datos (por ejemplo, los conjuntos de datos casi deterministas como ct_slices prefieren un mayor número de bins de histograma y la predicción de ruido sobre EDM). El trabajo establece que condicionar el problema supervisado inducido a la resolución finita de las divisiones de histograma es la clave para desbloquear el potencial de los modelos de difusión y de flujo basados en árboles.