Flow-Based Generative Modeling for Optimizing Sampling Policies in Compressed Sensing Applications

Resumen Técnico: Modelado Generativo Basado en Flujos para la Optimización de Políticas de Muestreo en Computación Compresiva

Planteamiento del Problema
Las aplicaciones modernas en el procesamiento de señales y la imagenología médica a menudo requieren la adquisición de señales de alta dimensión bajo estrictas restricciones de recursos. La teoría de muestreo clásica dicta que la reconstrucción precisa requiere mediciones proporcionales a la dimensión ambiental de la señal, un requisito que frecuentemente es impracticable. La Computación Compresiva (CS, por sus siglas en inglés) aborda esto permitiendo la recuperación de señales dispersas a partir de menos mediciones, siempre que el operador de medición satisfaga condiciones específicas (por ejemplo, incoherencia). Sin embargo, la calidad de la reconstrucción depende altamente de la elección del esquema de muestreo. La búsqueda de un patrón de submuestreo óptimo es un problema combinatorio NP-duro porque la máscara de muestreo es discreta y no diferenciable, lo que impide una optimización directa de extremo a extremo. Además, el espacio de búsqueda crece exponencialmente con el número de ubicaciones de los sensores, haciendo que la exploración exhaustiva sea infactible. Los enfoques generativos existentes, como el Submuestreo Probabilístico Profundo (DPS) o el Submuestreo de Difusión Activa (ADS), a menudo producen máscaras fijas o sufren de una alta latencia computacional, lo que limita su adaptabilidad y utilidad clínica en tiempo real.

Metodología
Los autores proponen un marco generativo basado en flujos y consciente de la tarea que reformula el paradigma convencional de entrenamiento de Flow Matching (FM). En lugar de modelar directamente la variedad de los datos, el modelo opera en la variedad del espacio de muestreo para aprender máscaras de submuestreo óptimas.

• Flow Matching en la Variedad de Muestreo: El marco utiliza una ecuación diferencial ordinaria (ODE) neuronal para transportar puntos desde una distribución de ruido gaussiano estándar ($p_0$) hacia una distribución objetivo de operadores de muestreo ($p_1$). El modelo aprende un campo vectorial dependiente del tiempo $v_\phi(x, t)$ que define una trayectoria generativa.
• Optimización Condicionada por Tarea: A diferencia del FM estándar que depende de un acoplamiento lineal entre el ruido y los datos, este marco acopla el flujo generativo con un resolvedor de tareas de CS de flujo descendente ($f_\theta$). El objetivo de entrenamiento minimiza una función de pérdida compuesta que combina la pérdida de regresión FM estándar ($L_{FM}$) y una pérdida específica de la tarea ($L_{task}$), ponderadas por parámetros de incertidumbre aprendibles ($\sigma_1, \sigma_2$). Esto permite optimizar conjuntamente la política de muestreo y el rendimiento de la tarea.
• Puerta Sigmoide Diferenciable: Para abordar la no diferenciabilidad de las máscaras binarias discretas, la salida continua del modelo de flujo ($x_t$) se pasa a través de una puerta sigmoide personalizada y diferenciable. Esta puerta emplea una función de pendiente dependiente del tiempo $s(t)$ y un umbral de percentil $\tau$ para producir máscaras suaves ($g(x_t) \in (0, 1)$). A medida que el paso de tiempo $t$ aumenta, la transición se vuelve más nítida, regularizando efectivamente la entropía del proceso de generación de la máscara.
• Modos de Inferencia: El marco soporta dos modos operativos:
  1. De Aleatorio a Optimizado: Se toman mediciones aleatorias iniciales, se extraen características y el modelo de flujo genera una máscara optimizada para las mediciones subsecuentes. Esto es adecuado para escenarios donde las mediciones son económicas o las señales son estáticas (por ejemplo, MNIST, CelebA).
  2. Detección Adaptativa: Primero se adquieren datos de k-space de baja frecuencia (LF). Un extractor de características condiciona el modelo de flujo para determinar el patrón de muestreo óptimo para las bandas de alta frecuencia restantes. Este modo está diseñado para mediciones costosas, específicamente la aceleración de RM.

Contribuciones Clave

• Nuevo Paradigma: El artículo introduce una prueba de concepto donde el Flow Matching se entrena dentro de la variedad de los operadores de muestreo en lugar del espacio de la señal, diseñado específicamente para tareas de CS de flujo descendente.
• Política Continua y Adaptativa: El marco proporciona una parametrización continua y de tiempo continuo del esquema de muestreo, ofreciendo una evolución suave desde estructuras gruesas hasta refinamientos específicos de la tarea, a diferencia de los métodos de generación de máscaras estáticas.
• Optimización Conjunta: El método permite la optimización de extremo a extremo tanto de la estrategia de muestreo como de la tarea de reconstrucción/clasificación, aprovechando el prior de la tarea para guiar el proceso generativo.
• Marco Unificado: Ofrece un enfoque flexible aplicable a diversas tareas inversas, incluyendo clasificación de imágenes, reconstrucción de imágenes y aceleración de RM, sin requerir el reentrenamiento para diferentes operadores de medición.

Resultados Experimentales
El marco fue evaluado en MNIST, CelebA y el conjunto de datos fastMRI knee:

• Clasificación de MNIST: El enfoque basado en flujos mejoró la precisión de la clasificación en más de un 20% en comparación con clasificadores entrenados con máscaras aleatorias. Aunque quedó por detrás del modelo A-DPS en ciertas proporciones de muestreo fijas, demostró una mayor versatilidad a través de un rango de proporciones (1% a 8%) sin necesidad de un entrenamiento específico para cada proporción.
• Reconstrucción de Imágenes (CelebA y MNIST): El método alcanzó un rendimiento de vanguardia en la reconstrucción de imágenes. En el conjunto de datos CelebA, alcanzó una Relación Señal-Ruido de Pico (PSNR) de 25.17 dB con una tasa de submuestreo del 5%. Superó al método de vanguardia ADS, particularmente en tasas de muestreo agresivas (bajas).
• Aceleración de RM: En los experimentos de fastMRI, el marco logró un Índice de Similitud Estructural (SSIM) de 0.7163 con una aceleración de 8×, superando a todos los modelos de referencia (LOUPE, DPS, ADS, A-DPS) en este factor de aceleración elevado. Con una aceleración de 4×, su desempeño fue comparable al de otros métodos.
• Eficiencia Computacional: El tiempo de inferencia para generar una máscara de muestreo fue de aproximadamente 127–130 ms en una NVIDIA GeForce RTX 4070. Esto es significativamente más rápido que el tiempo de adquisición por línea de k-space en RM (500–2000 ms) y drásticamente más rápido que el tiempo de adquisición por línea reportado por el método ADS (3040 ms), lo que sugiere la viabilidad para el despliegue clínico en tiempo real.

Significancia y Reivindicaciones
Los autores afirman que este trabajo establece la viabilidad fundamental de utilizar modelos generativos basados en flujos condicionados por tareas para optimizar las políticas de muestreo en computación compresiva. El estudio destaca que:

1. El condicionamiento por tarea es efectivo: Integrar un prior de tarea de flujo descendente directamente en el proceso de entrenamiento de Flow Matching mejora significativamente el desempeño en problemas inversos.
2. Restricciones de la variedad de muestreo: El espacio de muestreo posee restricciones estructurales intrínsecas (binariedad, dispersión) que pueden aprovecharse para mejorar la eficiencia del entrenamiento en comparación con el modelado de la variedad de la imagen natural.
3. Viabilidad Clínica: La baja carga computacional y la naturaleza adaptativa del marco lo convierten en un candidato prometedor para aplicaciones de imagenología médica en el mundo real, particularmente para reducir los tiempos de escaneo de RM y mejorar la calidad de los datos diagnósticos.
4. Generalizabilidad: El enfoque ofrece una estrategia unificada y flexible para diseñar esquemas de detección impulsados por datos y tareas, aplicables a una amplia gama de problemas inversos más allá de las tareas específicas probadas.

El artículo se posiciona como un estudio de prueba de concepto, reconociendo que, si bien logra resultados de vanguardia en tareas de reconstrucción específicas, la contribución principal es la demostración de un nuevo y flexible marco para aprender políticas de muestreo que pueden adaptarse a diversos problemas inversos.