Consist-Retinex: One-Step Noise-Emphasized Consistency Training Accelerates High-Quality Retinex Enhancement

1. Formulação do Problema

O aprimoramento de imagens em baixa luminosidade (LLIE) é crítico para aplicações que vão desde a fotografia móvel até a navegação autônoma. Embora a teoria Retinex (decompor uma imagem em Reflexão $R$ e Iluminação $L$) forneça um quadro robusto para separar estrutura da iluminação, abordagens generativas modernas de alta qualidade (como modelos de Difusão) dependem tipicamente de amostragem iterativa (por exemplo, 1000 passos). Isso cria um gargalo para a implantação em tempo real sob orçamentos estritos de latência.

Modelos de Consistência (CMs) oferecem uma solução ao permitir a geração em um único passo através da auto-consistência ao longo de trajetórias de fluxo de probabilidade. No entanto, aplicar diretamente Modelos de Consistência padrão ao aprimoramento baseado em Retinex enfrenta dois desafios críticos:

1. Incompatibilidade de Extremos: A inferência em um único passo opera no nível máximo de ruído ($\sigma_{max}$), mas os cronogramas de treinamento padrão (amostragem log-uniforme) raramente visitam essa região de alto ruído (apenas ~0,5% das amostras), levando a previsões instáveis no ponto de inferência.
2. Ambiguidade Condicional: CMs padrão impõem auto-consistência, mas não garantem que a saída se alinhe com o alvo condicional correto (a imagem aprimorada) quando dados de verdade fundamental pareados estão disponíveis. Um modelo pode ser auto-consistente, mas ancorado ao valor errado.

2. Metodologia: Consist-Retinex

Os autores propõem o Consist-Retinex, um quadro que integra a decomposição Retinex com uma estratégia de treinamento de consistência especializada projetada para inferência em um único passo.

A. Arquitetura e Decomposição

• Rede de Decomposição Transformer Retinex (TDN): O pipeline começa com um módulo TDN que decompõe a imagem de entrada em baixa luminosidade ($I_l$) em mapas de reflexão em baixa luminosidade ($R_l$) e iluminação ($L_l$). Da mesma forma, processa imagens em luz normal para gerar alvos ($R_n, L_n$).
• Modelos de Consistência Duplamente Condicionais: Dois modelos de consistência separados, $f^R_\theta$ e $f^L_\theta$, são treinados. Eles mapeiam componentes em luz normal com ruído ($R_n + \sigma\epsilon, L_n + \sigma\epsilon$) para componentes aprimorados, condicionados aos mapas correspondentes em baixa luminosidade ($R_l, L_l$).
• Reconstrução: A imagem final aprimorada é reconstruída via multiplicação elemento a elemento: $\hat{I} = \hat{R} \odot \hat{L}$.

B. Inovações Principais

A estratégia de treinamento aborda a incompatibilidade de extremos e a ambiguidade condicional através de dois mecanismos chave:

1. Amostragem Adaptativa com Ênfase em Ruído:

   • A amostragem log-uniforme padrão é ineficiente para inferência em um único passo.
   • Os autores introduzem uma estratégia de amostragem mista: 95% das amostras de treinamento são extraídas de um intervalo de alto ruído $[0.95\sigma_{max}, \sigma_{max}]$, enquanto 5% são extraídas do intervalo completo.
   • Isso garante que o modelo receba supervisão densa especificamente no nível de ruído usado durante a inferência ($\sigma_{max}$).

2. Função de Perda de Duplo Objetivo:
   A perda total combina dois termos complementares:

   • Perda de Consistência Temporal ($L_{consist}$): Impõe auto-consistência ao longo de uma grade discreta de Karras (treinamento de consistência padrão) para garantir coerência da trajetória.
   • Perda de Alinhamento de Ponto Fixo de Verdade Fundamental ($L_{fixed}$): Ancora a previsão do modelo no extremo de alto ruído diretamente aos componentes de verdade fundamental pareados ($R_n, L_n$). Isso impede que o modelo convirja para um valor auto-consistente, mas incorreto.
   • Fórmula: $L_{total} = \lambda_{consist} L_{consist} + \lambda_{fixed} L_{fixed}$.

C. Análise Teórica

O artigo fornece justificativa teórica para o design:

• Limite de Erro de Extremo: Prova que o erro de inferência em um único passo é controlado pelo erro de ponto fixo de alto ruído mais um termo de incompatibilidade.
• Propagação de Ancoragem: Demonstra que, se o extremo de alto ruído estiver corretamente ancorado à verdade fundamental, a consistência temporal propaga essa correção para trás ao longo de toda a trajetória.
• Alocação de Amostras: Mostra matematicamente que enfatizar a amostragem de alto ruído reduz o número total de amostras de treinamento necessárias para cobrir o regime de inferência em quase 200 vezes em comparação com a amostragem uniforme.

3. Contribuições Principais

1. Formulação: Um quadro inovador de treinamento condicional de consistência de duplo objetivo especificamente para componentes decompostos em Retinex, combinando consistência de trajetória com alinhamento de verdade fundamental pareada.
2. Estratégia de Amostragem: Desenvolvimento de uma regra de amostragem adaptativa com ênfase em ruído que foca a supervisão perto do extremo de inferência em um único passo, validada através de extensos estudos de ablação.
3. Insights Teóricos: Fornecimento de limites de erro de extremo e análise de alocação de amostras que explicam por que a ancoragem de alto ruído e a consistência temporal são complementares para aprimoramento em um único passo.
4. Eficiência: Alcançamento de desempenho state-of-the-art com inferência em um único passo, reduzindo drasticamente tanto os passos de amostragem (1 vs. 1000) quanto os custos de treinamento (1/8 das amostras efetivas de baselines de difusão).

4. Resultados Experimentais

O método foi avaliado em benchmarks pareados (LOL, VE-LOL-L) e conjuntos de dados não pareados (DICM, VV).

• Desempenho no VE-LOL-L: O Consist-Retinex alcançou o melhor PSNR (25,51 dB) e SSIM (0,880) entre todos os métodos comparados sob inferência em um único passo, superando modelos de difusão multi-passo (por exemplo, Diff-Retinex++) e outras abordagens baseadas em consistência.
• Desempenho no LOL: Permaneceu competitivo (22,21 dB PSNR), superando significativamente outros métodos de um único passo e aproximando-se da qualidade de baselines de difusão de 1000 passos.
• Eficiência:
  • Inferência: ~12,5 FPS (0,08s por imagem) em uma única GPU, representando uma redução de 1000x nos passos de amostragem em comparação com baselines de difusão.
  • Treinamento: Requer apenas 200K iterações (1,6M amostras efetivas) em comparação com 800K iterações (12,8M amostras) para baselines de difusão, alcançando uma redução de 8x no custo de treinamento para a etapa de restauração.
• Estudos de Ablação: Confirmaram que remover tanto a consistência temporal quanto o alinhamento de verdade fundamental leva ao colapso do desempenho, e que a amostragem com ênfase em ruído é crítica para alto PSNR.

5. Significado

O Consist-Retinex preenche a lacuna entre restauração generativa de alta qualidade e implantação em tempo real. Ao resolver os desafios específicos de aplicar modelos de consistência ao aprimoramento de imagens fatorizado (Retinex), demonstra que a inferência em um único passo pode alcançar qualidade comparável a modelos de difusão iterativos sem a latência associada. A amostragem com ênfase em ruído e a função de perda de duplo objetivo propostas oferecem uma estratégia generalizável para treinar modelos de consistência em tarefas de restauração condicional onde a inferência ocorre na extremidade do cronograma de ruído. Este trabalho abre caminho para a implantação de aprimoramento de baixa luminosidade de alta fidelidade em dispositivos com restrições de latência, como telefones celulares e veículos autônomos.