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

Resumo Técnico: Modelagem Generativa Baseada em Fluxo para Otimização de Políticas de Amostragem em Compressed Sensing

Definição do Problema
Aplicações modernas em processamento de sinais e imagem médica frequentemente exigem a aquisição de sinais de alta dimensão sob rigorosas restrições de recursos. A teoria clássica de amostragem dita que a reconstrução precisa requer medições proporcionais à dimensão ambiente do sinal, um requisito que é frequentemente impraticável. O Compressed Sensing (CS) aborda isso ao permitir a recuperação de sinais esparsos a partir de menos medições, desde que o operador de medição satisfaça condições específicas (ex: incoerência). No entanto, a qualidade da reconstrução é altamente dependente da escolha do esquema de amostragem. Encontrar um padrão de subamostragem ideal é um problema combinatório NP-difícil porque a máscara de amostragem é discreta e não diferenciável, impedindo a otimização direta de ponta a ponta (end-to-end). Além disso, o espaço de busca cresce exponencialmente com o número de localizações dos sensores, tornando a exploração exaustiva inviável. Abordagens generativas existentes, como o Deep Probabilistic Subsampling (DPS) ou o Active Diffusion Subsampling (ADS), frequentemente produzem máscaras fixas ou sofrem de alta latência computacional, limitando sua adaptabilidade e utilidade clínica em tempo real.

Metodologia
Os autores propõem uma estrutura generativa baseada em fluxo, orientada por tarefas, que reformula o paradigma convencional de treinamento de Flow Matching (FM). Em vez de modelar o manifold de dados diretamente, o modelo opera no manifold do espaço de amostragem para aprender máscaras de subamostragem ideais.

• Flow Matching no Manifold de Amostragem: A estrutura utiliza uma Equação Diferencial Ordinária (ODE) neural para transportar pontos de uma distribuição Gaussiana padrão ($p_0$) para uma distribuição alvo de operadores de amostragem ($p_1$). O modelo aprende um campo vetorial dependente do tempo $v_\phi(x, t)$ que define uma trajetória generativa.
• Otimização Condicionada por Tarefa: Ao contrário do FM padrão, que depende de um acoplamento linear entre ruído e dados, esta estrutura acopla o fluxo generativo a um resolvedor de tarefa de CS de jusante ($f_\theta$). O objetivo de treinamento minimiza uma função de perda composta que combina a perda de regressão FM padrão ($L_{FM}$) e uma perda específica da tarefa ($L_{task}$), ponderadas por parâmetros de incerteza aprendíveis ($\sigma_1, \sigma_2$). Isso permite que o modelo otimize conjuntamente a política de amostragem e o desempenho da tarefa.
• Portão Sigmoide Diferenciável: Para lidar com a não diferenciabilidade de máscaras binárias discretas, a saída contínua do modelo de fluxo ($x_t$) é passada por um portão sigmoide customizado e diferenciável. Este portão emprega uma função de inclinação dependente do tempo $s(t)$ e um limiar de percentil $\tau$ para produzir máscaras suaves ($g(x_t) \in (0, 1)$). À medida que o passo de tempo $t$ aumenta, a transição torna-se mais nítida, efetivamente regularizando a entropia do processo de geração da máscara.
• Modos de Inferência: A estrutura suporta dois modos operacionais:
  1. Random-to-Optimized (Do Aleatório ao Otimizado): Medições aleatórias iniciais são realizadas, características são extraídas e o modelo de fluxo gera uma máscara otimizada para medições subsequentes. Isso é adequado para cenários onde as medições são baratas ou os sinais são estáticos (ex: MNIST, CelebA).
  2. Adaptive Sensing (Sensoriamento Adaptativo): Dados de k-espaço de baixa frequência (LF) são adquiridos primeiro. Um extrator de características condiciona o modelo de fluxo para determinar o padrão de amostragem ideal para as bandas de alta frequência restantes. Este modo é projetado para medições dispendiosas, especificamente aceleração de RM.

Principais Contribuições

• Novo Paradigma: O artigo introduz uma prova de conceito onde o Flow Matching é treinado dentro do manifold dos operadores de amostragem, em vez do espaço do sinal, adaptado especificamente para tarefas de CS de jusante.
• Política Contínua e Adaptativa: A estrutura fornece uma parametrização contínua e no tempo contínuo do esquema de amostragem, oferecendo uma evolução suave de estruturas grosseiras para refinamentos específicos da tarefa, ao contrário de métodos de geração de máscara estáticos.
• Otimização Conjunta: O método permite a otimização conjunta de ponta a ponta da estratégia de amostragem e da tarefa de reconstrução/classificação, aproveitando o prior da tarefa para guiar o processo generativo.
• Estrutura Unificada: Oferece uma abordagem flexível aplicável a vários problemas inversos, incluindo classificação de imagens, reconstrução de imagens e aceleração de RM, sem exigir o retreinamento para diferentes operadores de medição.

Resultados Experimentais
A estrutura foi avaliada em MNIST, CelebA e no conjunto de dados fastMRI knee:

• Classificação MNIST: A abordagem baseada em fluxo melhorou a precisão da classificação em mais de 20% em comparação com classificadores treinados em máscaras aleatórias. Embora tenha ficado atrás do modelo A-DPS em certas razões de amostragem fixas, demonstrou maior versatilidade em uma gama de razões (1% a 8%) sem necessidade de treinamento específico para cada razão.
• Reconstrução de Imagem (CelebA & MNIST): O método alcançou o estado da arte em reconstrução de imagem. No conjunto de dados CelebA, atingiu uma Relação Sinal-Ruído de Pico (PSNR) de 25,17 dB em uma taxa de subamostragem de 5%. Superou o método de estado da arte ADS, particularmente em taxas de amostragem agressivas (baixas).
• Aceleração de RM: Nos experimentos de fastMRI, a estrutura alcançou um Índice de Similaridade Estrutural (SSSS) de 0,7163 com aceleração de 8×, superando todos os modelos de referência (LOUPE, DPS, ADS, A-DPS) neste fator de aceleração elevado. Com aceleração de 4×, o desempenho foi comparável aos outros métodos.
• Eficiência Computacional: O tempo de inferência para gerar uma máscara de amostragem foi de aproximadamente 127–130 ms em uma NVIDIA GeForce RTX 4070. Isso é significativamente mais rápido que o tempo de aquisição por linha de k-espaço em RM (500–2000 ms) e drasticamente mais rápido que o tempo de aquisição por linha reportado pelo método ADS (3040 ms), sugerindo viabilidade para implantação clínica em tempo real.

Significância e Alegações
Os autores afirmam que este trabalho estabelece a viabilidade fundamental do uso de modelos generativos baseados em fluxo condicionados por tarefas para otimizar políticas de amostragem para compressed sensing. O estudo destaca que:

1. O Condicionamento por Tarefa é Eficaz: Integrar um prior de tarefa de jusante diretamente no processo de treinamento de flow matching melhora significativamente o desempenho em problemas inversos.
2. Restrições do Manifold de Amostragem: O espaço de amostragem possui restrições estruturais intrínsecas (binaridade, esparsidade) que podem ser aproveitadas para melhorar a eficiência do treinamento em comparação com a modelagem do manifold de imagem natural.
3. Viabilidade Clínica: O baixo custo computacional e a natureza adaptativa da estrutura a tornam uma candidata promissora para aplicações reais de imagem médica, particularmente para reduzir o tempo de varredura de RM e melhorar a qualidade dos dados diagnósticos.
4. Generalização: A abordagem oferece uma estratégia unificada e flexível para o design de esquemas de sensoriamento baseados em dados e tarefas, aplicáveis a uma ampla gama de problemas inversos além das tarefas específicas testadas.

O artigo posiciona-se como um estudo de prova de conceito, reconhecendo que, embora alcance resultados de estado da arte em tarefas específicas de reconstrução, a principal contribuição é a demonstração de uma nova e flexível estrutura para aprender políticas de amostragem que podem ser adaptadas a diversos problemas inversos.