Soft Equivariance Regularization for Invariant Self-Supervised Learning

O artigo propõe a Regularização de Equivariância Suave (SER), um método plug-in que desacopla a imposição de invariância e equivariância em diferentes camadas de redes de aprendizado auto-supervisionado, melhorando significativamente o desempenho em tarefas de classificação, robustez a perturbações e detecção de objetos sem a necessidade de cabeças auxiliares ou rótulos de transformação.

O essencial

Imagine que você está ensinando um aluno muito inteligente a reconhecer objetos em fotos, mas ele nunca viu uma foto antes. O método tradicional de ensino (chamado de Aprendizado Auto-Supervisionado) funciona assim: você mostra ao aluno a mesma foto, mas com alguns "truques" aplicados, como cortar um pedaço, mudar a cor ou girar levemente. O aluno é treinado para dizer: "Ei, essa foto cortada e essa foto colorida são a mesma coisa!".

O problema é que, ao focar apenas em dizer "é a mesma coisa", o aluno aprende a ignorar detalhes importantes. Ele perde a noção de como a imagem mudou. Se você girar a foto 90 graus, ele sabe que é o mesmo objeto, mas não aprende como o objeto se moveu no espaço. Isso é ruim se você quiser usar esse aluno para tarefas que exigem precisão espacial, como dirigir um carro autônomo ou detectar objetos em um vídeo.

Aqui entra a ideia do papel: SER (Regularização de Equivariância Suave).

A Analogia do "Chefe" e do "Estagiário"

Para entender o SER, vamos imaginar uma empresa com dois funcionários: um Chefe (a camada final da rede neural) e um Estagiário (uma camada intermediária).

1. O Problema Antigo (Tudo no Chefe):
   Antes, os pesquisadores tentavam ensinar duas coisas ao mesmo tempo para o Chefe:

   • "Ignore as mudanças de cor e corte" (Invariância).
   • "Entenda como a imagem girou ou foi espelhada" (Equivariância).
   • O Resultado: O Chefe ficava confuso. Tentar ser "cego" para mudanças e ao mesmo tempo "super atento" a elas no mesmo lugar causava um conflito. O Chefe aprendia a girar bem, mas perdia a capacidade de reconhecer o objeto em testes simples. Era como pedir para um juiz ser ao mesmo tempo imparcial e tendencioso na mesma decisão.

2. A Solução do SER (Dividir as Tarefas):
   O SER propõe uma mudança inteligente na estrutura da empresa:

   • O Estagiário (Camada Intermediária): Recebe a tarefa de observar os detalhes espaciais. Ele é treinado para entender: "Se a foto girou 90 graus, minha representação interna também deve girar 90 graus de forma previsível". Ele mantém a estrutura da imagem (quem está onde).
   • O Chefe (Camada Final): Continua fazendo apenas o que ele faz de melhor: ignorar as mudanças de cor e corte para dizer "Isso é um gato, não importa onde ele esteja". Ele não é forçado a entender a rotação; ele apenas recebe a informação já organizada pelo Estagiário.

Como funciona na prática?

O SER usa um truque matemático elegante:

• Ele pega um lote de fotos e divide em dois grupos.
• Um grupo é usado para o treinamento normal (o "Chefe" aprende a ser invariante).
• O outro grupo usa apenas transformações geométricas reversíveis (como girar 90 graus ou espelhar, sem cortar a foto). Nesse grupo, o sistema força o "Estagiário" a entender que, se a foto girou, a representação interna deve girar junto.
• O Pulo do Gato: O sistema não precisa de um professor externo dizendo "isso é uma rotação de 90 graus". Ele descobre sozinho, usando matemática pura, como as peças se encaixam.

Por que isso é incrível?

1. Sem custo extra: O método é tão leve que aumenta o tempo de treinamento em menos de 1%. É como adicionar um tempero extra ao prato sem precisar comprar novos ingredientes.
2. Melhor em tudo: Ao separar as tarefas, o modelo fica melhor em reconhecer objetos (como em testes de classificação) e muito mais robusto quando a imagem está borrada, distorcida ou em condições ruins (como chuva ou neblina).
3. Funciona em qualquer lugar: Eles testaram em vários modelos famosos e todos melhoraram.

Resumo em uma frase

O SER é como ensinar um aluno a ter dois modos de pensar: um modo "estagiário" que entende a geometria e o movimento do mundo, e um modo "chefe" que foca apenas em identificar o objeto, evitando que as duas tarefas se atrapalhem e resultando em uma inteligência artificial mais inteligente e robusta.

Resumo técnico

Título: Soft Equivariance Regularization for Invariant Self-Supervised Learning (SER)

1. O Problema

A Aprendizagem Auto-supervisionada (SSL) visual tradicional baseia-se no princípio de invariância: o modelo aprende representações que permanecem constantes (invariantes) sob transformações semânticas preservadoras (como cortes aleatórios e jitter fotométrico). Embora eficaz para reconhecimento de objetos, essa invariância forte pode suprimir estruturas dependentes de transformações (como orientação, reflexão ou escala), que são cruciais para:

• Robustez a perturbações geométricas.
• Transferência de aprendizado sensível ao espaço (ex: detecção de objetos).

Métodos anteriores que tentam incorporar equivariância (onde a representação muda de forma previsível e estruturada sob transformações) frequentemente impõem objetivos de invariância e equivariância na mesma representação final. No entanto, a representação final em arquiteturas como ViTs (Vision Transformers) é tipicamente "colapsada espacialmente" (ex: token [CLS]), o que a torna mal alinhada com ações de grupos espaciais.

Os autores observam empiricamente um trade-off: ao forçar a regularização de equivariância em camadas mais profundas (próximas à saída), os escores de equivariância aumentam, mas a precisão na avaliação linear no ImageNet-1k degrada consistentemente. Isso sugere que impor ambos os objetivos na mesma camada final é subótimo.

2. Metodologia: Soft Equivariance Regularization (SER)

O SER é um regularizador "plug-and-play" que desacopla onde a invariância e a equivariância são aplicadas, sem alterar a arquitetura base ou o objetivo principal de SSL.

Principais Componentes:

• Desacoplamento de Camadas (Layer-Decoupling):

  • Invariância: Mantida na representação final (embedding) através do objetivo SSL original (ex: MoCo-v3, DINO, Barlow Twins).
  • Equivariância: Aplicada suavemente em uma representação intermediária espacial (mapa de tokens espaciais), antes do colapso espacial final.
  • O token [CLS] é inserido apenas após a camada regularizada para preservar a estrutura espacial dos tokens intermediários.

• Ações de Grupo Analíticas no Espaço de Características:

  • O método utiliza ações de grupo geométricas invertíveis ($\rho_g$) especificadas analiticamente no espaço de características, como rotações de 90°, flips horizontais e escalas anisotrópicas.
  • Essas ações são aplicadas diretamente aos mapas de tokens, evitando a necessidade de módulos auxiliares de aprendizado de ação ou previsão de códigos de transformação por amostra.

• Estratégia de Partição de Batch e Aumentação:

  • Como cortes aleatórios (random crop) não formam um grupo (são não invertíveis), o SER divide o mini-batch em duas sub-partes:
    1. $b_1$ (Invariância): Segue a política de aumento padrão (incluindo cortes) para calcular a perda de invariância base.
    2. $b_2$ (Equivariância): Segue uma política modificada ($T_{eq}$) que desabilita cortes, mas mantém jitter fotométrico e amostra transformações geométricas invertíveis do grupo $G$.
  • A perda de equivariância é calculada apenas em $b_2$, alinhando os mapas de tokens intermediários através da transformação relativa $g = g_2 g_1^{-1}$.

• Função de Perda:

  • A perda total é $L = L_{inv1} + L_{inv2} + \lambda L_{equiv}$.
  • $L_{equiv}$ é uma perda contrastiva (NT-Xent) aplicada patch-a-patch entre os mapas de tokens alinhados e os originais, sem necessidade de cabeçalhos de previsão de transformação adicionais.

3. Contribuições Principais

1. Identificação do Trade-off: Demonstração empírica de que forçar equivariância na camada final prejudica o desempenho de classificação, motivando o design desacoplado.
2. SER (Soft Equivariance Regularization): Um método escalável que promove equivariância em representações intermediárias espaciais, mantendo a invariância na saída final.
3. Simplicidade e Eficiência: Não requer módulos auxiliares complexos, não aprende códigos de transformação por amostra e adiciona apenas 1.008x de custo computacional (FLOPs) em relação à base.
4. Princípio Geral de Desacoplamento: A técnica de mover o objetivo de equivariância para camadas intermediárias melhora não apenas o SER, mas também métodos base-line existentes (como EquiMod e AugSelf), sugerindo um princípio de design universal para combinar invariância e equivariância.

4. Resultados Experimentais

Os experimentos foram realizados no pré-treinamento de ViT-S/16 no ImageNet-1k, com comparações justas (número de visualizações/ views igualado).

• Avaliação Linear no ImageNet-1k:

  • O SER melhora o MoCo-v3 em +0.84% de Top-1 (atingindo 69.28% vs 68.44% da base) em configuração estritamente igualada de 2 views.
  • Consistentemente supera o MoCo-v3, DINO e Barlow Twins quando o SER é adicionado.
  • Supera outros métodos de SSL equivariante (AugSelf, STL, EquiMod) em cenários de views igualados.

• Robustez e Perturbações:

  • ImageNet-C (Corrupções): Melhoria de +1.11% no Top-1 médio.
  • ImageNet-P (Perturbações Geométricas): Melhoria de +1.22% no Top-1 médio.
  • O método demonstra maior robustez a ruídos, desfoque e distorções geométricas.

• Transferência Espacial (Detecção de Objetos):

  • Em detecção de objetos no COCO com backbone congelado, o SER alcança +1.7 mAP de melhoria, evidenciando a preservação de informações espaciais úteis para tarefas sensíveis à localização.

• Ablação:

  • A melhor performance é alcançada quando a regularização de equivariância e a inserção do token [CLS] ocorrem em uma camada intermediária (ex: 3ª camada do ViT), confirmando a hipótese do "ponto ideal" (sweet spot).

5. Significado e Impacto

O trabalho propõe uma mudança de paradigma na combinação de invariância e equivariância no aprendizado auto-supervisionado. Em vez de tentar resolver o conflito de objetivos na camada final, o SER resolve-o através de uma arquitetura de camadas desacopladas.

• Eficiência: É uma solução leve que não exige reengenharia pesada de modelos ou grandes custos computacionais.
• Generalidade: Funciona como um módulo complementar para diversas arquiteturas de SSL (contrastivas e não-contrastivas).
• Aplicabilidade: Melhora significativamente a robustez e a capacidade de transferência para tarefas que exigem sensibilidade geométrica e espacial, preenchendo uma lacuna deixada pelos métodos puramente invariantes.

O código é disponibilizado publicamente, facilitando a adoção e verificação independente.