Resumo Técnico: MoECodec

Declaração do Problema

A compressão de imagens orientada para máquinas emergiu como uma área de pesquisa crítica para suportar tarefas de visão computacional de jusante (downstream), tais como classificação, detecção e segmentação. No entanto, as abordagens existentes enfrentam limitações significativas:

1. Designs Específicos para Tarefas: A otimização de ponta a ponta para tarefas específicas frequentemente exige pares de codificador-decodificador dedicados ou arquiteturas de múltiplos ramos, levando a um overhead substancial de parâmetros e complexidade de implantação.
2. Adaptação Baseada em Transferência: Métodos que adaptam codecs base orientados para humanos usando pequenos parâmetros treináveis (ex: adaptadores) permanecem externamente anexados e dependem de heurísticas específicas de tarefas, em vez de serem intrinsecamente aprendidos dentro de um modelo unificado.
3. Padrões de Computação Estáticos: Ambas as linhas de trabalho aplicam majoritariamente transformações uniformes em todos os tokens de imagem. Isso falha em considerar o fato de que diferentes regiões da imagem exibem necessidades de importância semântica e complexidade marcadamente distintas para a percepção de máquina, levando a uma alocação ineficiente de recursos.

Metodologia

Os autores propõem o MoECodec, um framework de compressão de imagem unificado que integra a Mistura de Especialistas (Mixture-of-Experts - MoE) por token em modelos de Compressão de Imagem Aprendida (LIC) baseados em transformers. A inovação central é substituir as camadas de Rede Feed-Forward (FFN) padrão em blocos de transformer por camadas MoE, permitindo computação dinâmica ao nível de token condicionada ao conteúdo de entrada e aos objetivos da tarefa.

Componentes Arquiteturais Principais

1. Integração de MoE por Token:

   • O MoECodec substitui as camadas FFN tanto no codificador quanto no decodificador de backbones de LIC baseados em transformer (ex: TIC, TinyLIC, TCM).
   • Os tokens de entrada são roteados adaptativamente para redes de especialistas especializadas, permitindo que o modelo suporte objetivos perceptivos heterogêneos dentro de um único modelo unificado.

2. Estratégia de Roteamento (Expert-Choice + Regularização Espacial):

   • Desafio: Modelos de LIC baseados em Transformer frequentemente utilizam tokenização de grão fino (ex: patches de 1×1). O roteamento de escolha de token (Token-Choice - TC) ingênuo leva a atribuições de especialistas ruidosas e fragmentadas (artefatos de sal e pimenta) devido à falta de contexto espacial.
   • Solução: Os autores adotam o roteamento de Escolha de Especialista (Expert-Choice - EC), onde cada especialista seleciona ativamente seus top-$k$ tokens preferidos de toda a sequência. Isso inverte a perspectiva de decisão, permitindo que os especialistas considerem o contexto global e garantam uma utilização equilibrada sem perdas auxiliares de balanceamento de carga.
   • Regularização de Variância Total Espacial (TV): Para abordar adicionalmente as inconsistências espaciais locais no roteamento EC, um termo de regularização de TV espacial é adicionado à perda de treinamento. Isso penaliza variações de alta frequência nos mapas de afinidade de especialistas espaciais, incentivando atribuições de especialistas coerentes e suavemente segmentadas.

3. Arquitetura de Especialista Leve (Group Shuffle MLP):

   • Desafio: Substituir FFNs por $E$ especialistas independentes normalmente introduz um overhead de parâmetros de $E\times$, o que é proibitivo para modelos de compressão onde as FFNs dominam a contagem de parâmetros. Simplesmente reduzir o tamanho do especialista degrada as capacidades de agregação de canais.
   • Solução: Os autores propõem o Group Shuffle MLP (GShMLP). Esta arquitetura particiona os canais de entrada em $G$ grupos e aplica projeções lineares independentes dentro de cada grupo. Um embaralhamento de canais (channel shuffle) livre de parâmetros é inserido entre as duas camadas de projeção para restaurar a interação entre canais. Este design reduz significativamente a complexidade de parâmetros enquanto mantém a agregação de canais eficaz.

Principais Contribuições

1. Framework Unificado: O MoECodec é o primeiro framework de LIC unificado que integra MoE por token em codecs baseados em transformer, permitindo que um único modelo otimize conjuntamente a qualidade de reconstrução e múltiplas tarefas de visão de máquina.
2. Abordagem de Desafios Específicos de LIC: O artigo identifica e resolve dois desafios específicos da introdução de MoE em LIC:
   • Descontinuidade Espacial: Resolvida via roteamento Expert-Choice combinado com Regularização de Variância Total Espacial.
   • Overhead de Parâmetros: Resolvido via a arquitetura leve de Group Shuffle MLP (GShMLP).
3. Alocação Dinâmica de Computação: O framework estabelece um paradigma de computação dinâmica onde objetivos de tarefas heterogêneas são acomodados intrinsecamente, permitindo que o modelo foque a computação em regiões relevantes para a tarefa.

Resultados Experimentais

Experimentos extensivos demonstram melhorias consistentes sobre fortes baselines tanto em reconstrução quanto em tarefas de máquina.

• Reconstrução de Imagem: Aplicado a três backbones baseados em transformer (TIC, TinyLIC, TCM), o MoECodec alcançou melhorias de BD-Rate de 11,54%, 8,12% e 6,42% respectivamente no dataset Kodak (PSNR), enquanto simultaneamente reduziu o total de parâmetros.
• Tarefas de Jusante de Máquina (Machine Downstream Tasks): Usando o TIC como base, o MoECodec superou os baselines Full Finetune, TransTIC e AdaptICMH:
  • Classificação (ImageNet-1K): Melhoria de 7,74% no BD-Rate.
  • Detecção de Objetos (COCO): Melhoria de 8,98% no BD-Rate.
  • Segmentação de Instância (COCO): Melhoria de 9,72% no BD-Rate.
• Eficiência: O MoECodec reduziu os parâmetros ativados de 7,51M para 6,25M e alcançou uma latência de inferência de 145ms em uma RTX 4090, superando o TIC original (155ms) em 8,4%.
• Estudos de Ablação:
  • A mudança de roteamento Token-Choice para Expert-Choice melhorou o BD-Rate de 19,57% para 15,90%.
  • A adição de regularização de TV espacial reduziu ainda mais o BD-Rate para 15,44%.
  • O GShMLP proporcionou um melhor compromisso entre eficiência e desempenho do que a redução das razões de expansão de FFN.
  • Posicionar as camadas MoE tanto no codificador quanto no decodificador rendeu o melhor desempenho.

Significância e Alegações

O artigo afirma que o MoECodec oferece uma nova abordagem para o treinamento de codecs multi-tarefa de ponta a ponta, afastando-se de paradigmas de treinamento customizados estáticos e específicos para tarefas. Ao permitir a alocação de computação dinâmica e esparsamente ativada, o framework permite que um único modelo se adapte para focar recursos computacionais em regiões semanticamente importantes para tarefas específicas. Os autores asseguram que esta abordagem equilibra com sucesso o compromisso entre custo do modelo (parâmetros e latência) e desempenho, tanto na percepção humana (reconstrução) quanto em tarefas de visão de máquina, demonstrando que a granularidade ao nível de token é eficaz para computação heterogênea em compressão de imagem multi-tarefa.