Resumo Técnico: Um Framework Unificado para Vision Transformers Equivariantes a Subgrupos Discretos de O(2)

Declaração do Problema

Os Vision Transformers (ViTs) estabeleceram-se como arquiteturas dominantes para o reconhecimento visual, no entanto, os modelos padrão falham em codificar explicitamente as simetrias planares inerentes a muitos domínios de visão. Embora redes neurais convolucionais (CNNs) grupo-equivariantes tenham incorporado com sucesso simetrias (como translações e rotações) como vieses indutivos, a transição para arquiteturas baseadas em Transformers requer um novo framework teórico. As propostas existentes de ViTs equivariantes têm sido limitadas a grupos de simetria específicos (por exemplo, simetria de reflexão ou $D_4$) ou a estratégias de implementação específicas. Há uma necessidade de um framework unificado que generalize essas abordagens para quaisquer subgrupos discretos do grupo de roto-reflexão $O(2)$, fornecendo uma maneira sistemática de construir, comparar e analisar ViTs equivariantes independentemente de um grupo de simetria específico.

Metodologia

Fundamentação Teórica

Os autores adotam uma abordagem de teoria das representações para construir ViTs equivariantes a um grupo finito $G \leq O(2)$. O framework baseia-se na decomposição dos espaços de características em representações irredutíveis (irreps) de $G$.

• Espaços de Características: As características dos tokens são modeladas como elementos de uma representação $G$ de $V$, decomposta como uma soma direta de irreps: $V = \bigoplus_{\rho \in \hat{G}} \mathbb{R}^{C_\rho} \otimes V_\rho$, onde $C_\rho$ é a multiplicidade.
• Grades de Patchification: A entrada é definida em uma grade de patchification $G$ $H_0 = U + H$, onde $U$ é um patch base e $H$ é um conjunto de translações, ambos estáveis sob a ação de $G$. Isso permite que o modelo lide com várias estruturas de grade (por exemplo, quadrada para $D_4$, hexagonal para $D_6$).

Componentes Arquiteturais Principais

O framework constrói análogos equivariantes dos componentes padrão de um ViT:

1. Camadas Lineares: Baseando-se no Lema de Schur, camadas lineares equivariantes são construídas como operações diagonais em blocos atuando nas dimensões de canal de cada irrep. Para irreps do tipo complexo (comuns em subgrupos de $O(2)$), as camadas lineares utilizam uma base que abrange uma subálgebra isomórfica a $\mathbb{C}$.
2. Embedding de Patch & Codificação Posicional: O embedding de patch é um mapa afim $G$-equivariante do espaço da imagem para o espaço de características do token. As codificações posicionais são implementadas como vetores aprendíveis $G$-invariantes adicionados aos tokens.
3. Não-linearidades: Para introduzir não-linearidade enquanto preserva a equivariância, os autores propõem uma construção "tipo Fourier". As características são mapeadas da base da irrep para um $G$-set $X$ (um espaço homogêneo), uma não-linearidade pontual $\sigma$ é aplicada, e o resultado é mapeado de volta para a base da irrep. A escolha do $G$-set $X$ (composto por espaços homogêneos $G/H$) determina a expressividade da não-linearidade.
4. Auto-atenção (Self-Attention): Os autores propõem uma construção geral para auto-atenção equivariante ao computar escores de atenção usando um produto interno $G$-invariante e projetando em subespaços estáveis sob $G$ (cabeças de atenção). Dois mecanismos específicos são definidos:
   • Atenção por Irrep: As cabeças são decompostas por irrep.
   • Atenção Acoplada: As cabeças são formadas pela mistura de irreps.
     Os autores provam que, no cenário de cabeça única, a parametrização equivariante é tão expressiva quanto a auto-atenção padrão restrita a mapas $G$-equivariantes.
5. Cabeça de Classificação: Para tarefas invariantes (como classificação de imagens), um token de classe é anexado. Para garantir que a saída final seja invariante, os componentes não triviais do token de classe são "invariantizados" (convertidos em normas) antes da camada de classificação linear final.

Resultados Teóricos

• Teorema de Embedding: Os autores provam que, se $H \leq G$, a classe de ViTs $G$-equivariantes se incorpora naturalmente na classe de ViTs $H$-equivariantes. Isso implica que impor um grupo de simetria maior restringe a classe de funções (reduzindo a expressividade), mas aumenta o viés indutivo.
• Expressividade da Atenção: A camada de auto-atenção $G$-equivariante de cabeça única pode representar qualquer mapa $G$-equivariante alcançável pela auto-atenção comum, desde que o mapa seja $G$-equivariante.

Resultados Experimentais

O framework foi avaliado no conjunto de dados de imagens aéreas PatternNet (38 classes) sob regimes de escassez de dados (10%, 40% e 100% dos dados de treinamento) para subgrupos de $D_4$ (grade quadrada) e $D_6$ (grade hexagonal).

• Impacto da Equivariância: Resultados preliminares sob orçamentos de parâmetros correspondentes indicam que modelos equivariantes (especificamente aqueles que impõem invariância para classificação) podem melhorar a precisão do reconhecimento em comparação com baselines não-equivariantes, particularmente em regimes de baixos dados. Isso apoia a hipótese de que a equivariância atua como uma estratégia eficaz de aumento de dados.
• Mecanismos de Atenção: Comparações entre mecanismos de atenção "por irrep" e "acoplada" mostraram diferenças de desempenho dependendo do regime de dados e da capacidade do modelo, embora nenhum mecanismo tenha sido universalmente superior em todos os cenários.
• Configurações de Não-linearidade: Os experimentos analisaram como a escolha das configurações de espaço homogêneo (o $G$-set $X$ usado nas não-linearidades) afeta o desempenho, destacando que as escolhas arquiteturais na camada de não-linearidade são críticas para a eficácia do modelo.

Contribuições Principais e Significância

• Framework Unificado: O artigo fornece uma construção geral para ViTs equivariantes a qualquer subgrupo discreto de $O(2)$, subsumindo trabalhos anteriores de ViTs equivariantes a $D_1$ (reflexão) e $D_4$ como casos especiais.
• Garantias Teóricas: Estabelece garantias rigorosas de expressividade, provando que camadas equivariantes não perdem poder expressivo em relação às camadas padrão quando restritas a mapas equivariantes, e formalizando o compromisso entre a força da simetria e a expressividade do modelo.
• Compatibilidade Hexagonal: O framework estende-se com sucesso para a patchification hexagonal, permitindo modelos $D_6$-equivariantes que respeitam simetrias de rotação de seis dobras, que são relevantes para imagens aéreas.
• Alegações Modestas: Os autores enfatizam que os experimentos foram desenhados para investigar escolhas arquiteturais em cenários controlados de pequena escala, em vez de otimizar o desempenho em benchmarks de larga escala. Os resultados são apresentados como evidência preliminar motivando estudos mais profundos sobre como grupos de simetria discretos moldam o reconhecimento visual baseado em transformers.

O trabalho visa separar a estrutura de representação teórica dos detalhes de implementação específicos de cada grupo, oferecendo uma base comum para construir e analisar ViTs equivariantes através de diferentes grupos de simetria planar.