Bridging the Training-Deployment Gap: Gated Encoding and Multi-Scale Refinement for Efficient Quantization-Aware Image Enhancement

Título: Superando la Brecha entre Entrenamiento e Implementación: Codificación con Puertas y Refinamiento Multi-Escala para Mejora de Imágenes Eficiente con Conciencia de Cuantización

1. El Problema

Los modelos de mejora de imágenes basados en aprendizaje profundo para dispositivos móviles enfrentan un desafío fundamental: equilibrar la alta calidad de salida con la velocidad de procesamiento requerida por el hardware móvil.

• Discrepancia de Precisión: Aunque los modelos modernos pueden transformar fotos de baja calidad en imágenes de alta fidelidad, su rendimiento se degrada significativamente al convertirlos a formatos de baja precisión (como INT8) para su ejecución en teléfonos móviles.
• Desajuste Entrenamiento-Despliegue: Los modelos se optimizan en dominios de punto flotante de alta precisión (FP32/FP16), pero se ejecutan en entornos cuantizados (8 bits).
• Sensibilidad de las Tareas de Nivel de Píxel: A diferencia de tareas de clasificación, las tareas de restauración de imágenes (Deep ISP) tienen distribuciones de activación altamente sensibles, asimétricas y con colas largas. La cuantización estándar (Post-Training Quantization - PTQ) introduce artefactos visuales severos, como desplazamientos de color, bandas y distorsiones de textura, debido a la pérdida de información en valores extremos (outliers) críticos para la fidelidad visual.

2. Metodología Propuesta

Los autores proponen un marco unificado que combina una arquitectura híbrida eficiente con una estrategia de entrenamiento consciente de la cuantización (QAT).

A. Arquitectura de Red Híbrida:
El modelo utiliza una estructura jerárquica de tres escalas diseñada para preservar características visuales finas:

• Bloques de Codificador con Puertas (Gated Encoder Blocks): Inspirados en arquitecturas de codificación descendente, utilizan un diseño de doble rama con un mecanismo de puerta (producto de Hadamard). Esto permite un flujo de información no lineal donde una rama actúa como una "máscara suave" para la otra. Se preservan las tres salidas de la puerta ($x_a, x_g, x_b$) y se envían al decodificador mediante conexiones de salto multicanal para recuperar detalles espaciales que de otro modo se perderían.
• Refinamiento Multi-Escala: Se aplica un operador de refinamiento (basado en bloques residuales tipo UNet) en múltiples niveles de la red (escalas S/2, S/4 y en el cuello de botella S/8). Esto mejora la discriminabilidad de las características y preserva tanto la estructura de iluminación global como los detalles de textura local.
• Fusión de Características Multi-Rama (Decodificador): El decodificador agrega información de múltiples niveles semánticos concatenando tres flujos: características semánticas de alto nivel, características refinadas del codificador y características de salto crudas. Una capa de convolución $1\times1$ reduce la dimensionalidad antes de un bloque de refinamiento final para asegurar una reconstrucción estable.

B. Estrategia de Entrenamiento: QAT (Quantization-Aware Training):

• Simulación de Cuantización: En lugar de entrenar en FP32 y cuantizar después, el modelo se entrena simulando los efectos de la cuantización de 8 bits durante el proceso de entrenamiento.
• Fake Quantization: Se utilizan operadores de "cuantización falsa" (FakeQuant) que mapean pesos y activaciones a niveles discretos en el paso forward, pero mantienen el cálculo en FP32.
• Estimador Straight-Through (STE): Para permitir el retropropagación de gradientes a través de las operaciones de redondeo no diferenciables, se utiliza un estimador que pasa los gradientes sin alterar el valor de los pesos subyacentes.
• Función de Pérdida: Se emplea una combinación de pérdida PSNR (para fidelidad de píxeles), similitud coseno (para integridad estructural) y una pérdida "Outlier-Aware" (para estabilizar el aprendizaje ante valores extremos).

3. Contribuciones Clave

1. Identificación de Desafíos Específicos: Se destaca la sensibilidad de las distribuciones de activación en tareas de mejora de imágenes (Deep ISP) frente a la cuantización, diferenciándolas de las tareas de visión de alto nivel.
2. Marco de Optimización Consistente con el Despliegue: Introducción de QAT específicamente adaptado para Deep ISP, minimizando la degradación de calidad al pasar a inferencia de 8 bits.
3. Arquitectura Eficiente y Escalable: Diseño de una red con bloques de codificador con puertas y refinamiento multi-escala que equilibra la eficiencia computacional con la calidad perceptual.
4. Validación en Dispositivos Reales: Demostración de que el modelo mantiene una baja sobrecarga computacional y es viable para su despliegue en dispositivos móviles comerciales estándar (ej. Snapdragon 8 Gen 2).

4. Resultados

Los experimentos se realizaron en el conjunto de datos DPED (iPhone vs. DSLR) y en el desafío Mobile AI 2026.

• Rendimiento Cuantitativo: El modelo propuesto obtuvo el segundo lugar en el desafío de mejora de imágenes RGB de Mobile AI 2026, con un PSNR de 21.82 dB y un SSIM de 0.7653.
• Comparativa de Cuantización (QAT vs. PTQ):
  • La cuantización posterior al entrenamiento (PTQ) en 8 bits causó una caída drástica en la calidad (PSNR ~20.57 dB) y artefactos visuales severos (cambio de color rosa en el cielo).
  • El modelo entrenado con QAT en INT8 recuperó significativamente la calidad, alcanzando un PSNR de 21.05 dB y un SSIM de 0.725, acercándose mucho al rendimiento del modelo FP32.
• Velocidad y Latencia:
  • En hardware dedicado (Qualcomm Hexagon Tensor Processor - QNN HTP), la inferencia INT8 con QAT logró una reducción de latencia del 72% en comparación con la base FP16 (de 151 ms a ~41.8 ms).
  • El modelo ofrece un equilibrio óptimo entre precisión y velocidad, siendo viable para aplicaciones en tiempo real.
• Estudios de Ablación: Se demostró que los bloques de refinamiento son críticos para la fidelidad, y que la configuración óptima de canales es $c=32$, ofreciendo el mejor equilibrio entre calidad y costo computacional.

5. Significado e Impacto

Este trabajo es significativo porque cierra la brecha crítica entre el diseño teórico de arquitecturas de alta fidelidad y las restricciones prácticas de implementación en el borde (edge computing).

• Paradigma de Despliegue: Establece que para tareas de nivel de píxel sensibles, el entrenamiento debe ser consciente de la cuantización desde el inicio, no como un paso posterior.
• Viabilidad Móvil: Demuestra que es posible lograr mejoras de imagen de alta calidad (cercanas a DSLR) en teléfonos móviles sin sacrificar la velocidad de inferencia ni la eficiencia energética, gracias a la combinación de una arquitectura eficiente y QAT.
• Reproducibilidad: El código está disponible públicamente, fomentando el avance en la optimización de modelos de visión para dispositivos móviles.

En resumen, el artículo presenta una solución robusta y escalable para la mejora de imágenes en móviles, superando las limitaciones de la cuantización tradicional mediante un diseño arquitectónico inteligente y un entrenamiento adaptado al hardware final.