Resumo Técnico: FTerViT: Vision Transformer Totalmente Ternário

Declaração do Problema

Os Vision Transformers (ViTs) são classificadores de imagem poderosos, mas geralmente inadequados para dispositivos de classe de microcontrolador devido ao seu grande consumo de memória. Um ViT compacto padrão, como o DeiT-Small, requer aproximadamente 88,3 MB em precisão FP32, enquanto as unidades de microcontrolador (MCUs) típicas oferecem apenas alguns megabytes de RAM externa. Essa incompatibilidade impede a implantação de aplicações de visão sempre ativas e de baixo poder em dispositivos de borda sem descarregamento para a nuvem.

Embora a quantização ternária (constrangendo os pesos a $\{-1, 0, +1\}$) ofereça uma compressão teórica de 16$\times$ sobre o FP32, os métodos existentes mais avançados apenas ternarizam as camadas do codificador do transformador. Eles mantêm a incorporação de patches, os parâmetros de LayerNorm e as cabeças de classificador em precisão completa (FP32) ou inteiros de bits inferiores (INT8). Em modelos compactos, esses componentes restantes de precisão completa tornam-se o gargalo dominante, consumindo até 38% do tamanho do modelo, apesar de representarem uma pequena fração dos parâmetros. Além disso, tentativas anteriores de totalmente ternarizar esses componentes sensíveis (por exemplo, TerViT) resultaram em quedas significativas de precisão (até 22,4 pontos percentuais), em grande parte devido à conhecida sensibilidade das primeira e última camadas à quantização.

Metodologia

Os autores propõem o FTerViT, o primeiro Vision Transformer totalmente ternarizado, onde todas as matrizes de pesos e parâmetros de normalização são constrangidos a $\{-1, 0, +1\}$. A abordagem consiste em dois pilares principais: operadores ternários novel e um protocolo de treinamento especializado.

1. Primitivas Ternárias

Para eliminar gargalos FP32, os autores substituem componentes padrão por módulos ternários dedicados:

• TernaryBitLinear: Usado para camadas totalmente conectadas. Ele quantiza os pesos para $\{-1, 0, +1\}$ usando escalonamento por média absoluta por tensor ($absmean$) e arredondamento limitado. As ativações são quantizadas para inteiros de 8 bits usando escalonamento por máximo absoluto por token.
• TernaryBitConv2d: Especificamente projetado para a camada de incorporação de patches. Ele aplica escalonamento por canal para lidar com magnitudes heterogêneas de filtros, abordando a sensibilidade da camada de incorporação de patches, que trabalhos anteriores identificaram como o componente mais crítico.
• TernaryLayerNorm: Ternariza os parâmetros afins aprendíveis ($\gamma$ e $\beta$) do LayerNorm usando o mesmo esquema $absmean$. Estatísticas por token (média e variância) e vieses permanecem em FP32, pois sua contagem de parâmetros é negligenciável, mas sua ternarização destruiria informações posicionais.

2. Protocolo de Treinamento

O modelo é treinado usando uma estratégia de Distilação de Conhecimento (KD) em duas fases:

• Fase 1 (Treinamento): O estudante ternário é inicializado a partir de um professor FP32 pré-treinado e treinado usando Distilação Consciente de Quantização (QAD) com uma taxa de aprendizado de $10^{-4}$ (decaimento cosseno) por 250 épocas. A função de perda utiliza apenas o termo de divergência KL forward entre os logits do estudante e do professor, evitando a entropia cruzada, que pode entrar em conflito com o KL em redes de baixo bit.
• Fase 2 (Recuperação): Uma fase de ajuste fino leve onde a taxa de aprendizado é reiniciada em $10^{-5}$ (decaimento cosseno) por 10 épocas. Esta fase permite a recalibração de escalas e a recuperação da precisão perdida durante a saturação inicial da Fase 1.

Contribuições Principais

1. Primeiro ViT Totalmente Ternário: Os autores demonstram que os componentes mais frágeis dos ViTs (incorporação de patches, LayerNorms e cabeça de classificador) podem ser ternarizados com sucesso para $\{-1, 0, +1\}$ quando treinados com distilação de conhecimento, empurrando os limites do que era anteriormente considerado quantizável.
2. Operadores Novos: Introdução do TernaryBitConv2d com escalonamento por canal e do TernaryLayerNorm para lidar com a sensibilidade específica de camadas não codificadoras.
3. Estratégia de Treinamento: Uma estratégia de treinamento simples, porém eficaz, em duas fases (treinamento com alta taxa de aprendizado seguida por recuperação com baixa taxa de aprendizado) que permite convergência mais rápida e precisão final significativamente maior em comparação com abordagens anteriores de ViT ternário.
4. Implementação em Dispositivo: A primeira implementação C autônoma de um Vision Transformer ternário em um microcontrolador ESP32-S3 de dois núcleos com restrições de recursos (8 MB de PSRAM).

Resultados Experimentais

O método foi avaliado no ImageNet-1K e em outros benchmarks:

• Precisão e Compressão:

  • FTerViT-DeiT-III-S (384$\times$384): Alcança 82,43% de precisão top-1 em 6,09 MB (compressão de 14,6$\times$), perdendo apenas 2,42 pontos percentuais em comparação com seu professor FP32. Isso supera os métodos anteriores de ViT ternário (por exemplo, TerViT) em até 8 pontos percentuais.
  • FTerViT-DeiT-III-S (224$\times$224): Alcança 79,64% de top-1 em 5,81 MB.
  • FTerViT-DeiT-S (224$\times$224): Alcança 77,47% de top-1 em 5,81 MB, superando o TerViT-DeiT-S em 3,27 pontos percentuais.
  • Benchmarks CIFAR: Variantes DeiT-Tiny alcançam 97,43% (CIFAR-10) e 86,01% (CIFAR-100) com apenas 1,53 MB de armazenamento.

• Análise de Sensibilidade de Componentes:

  • Análises de primeira ordem de Taylor e traço de Hessian confirmam que LayerNorms e incorporações de patches dominam a importância do modelo (34–39% e 4,6%, respectivamente), apesar de ocuparem <0,2% dos parâmetros. Isso explica por que métodos anteriores os deixaram em precisão superior e por que sua ternarização bem-sucedida é crítica.

• Implantação em Dispositivo (ESP32-S3):

  • O modelo de 5,81 MB cabe inteiramente nos 8 MB de PSRAM do ESP32-S3, deixando ~4,5 MB livres para buffers de câmera e LCD.
  • Desempenho: Uma passagem direta leva 21,06 segundos.
  • Eficiência: Comparado ao FP32, a inferência ternária empacotada de 2 bits é 1,71$\times$ mais rápida (804 ms vs. 1376 ms para uma camada específica) e reduz o consumo de energia acima do estado ocioso em 55% (59 mJ vs. 130 mJ).

Significado

O artigo afirma que o FTerViT reduz substancialmente a lacuna entre a precisão padrão do ViT e as restrições estritas de memória da implantação em escala de microcontrolador. Ao ternarizar com sucesso toda a arquitetura, incluindo as incorporações de patches e camadas de normalização anteriormente "intocáveis", os autores permitem a implantação de modelos de visão de alta precisão em MCUs comerciais e de baixo custo (10 dólares). O trabalho valida que, com a estratégia de treinamento correta (KD de mesma arquitetura e recuperação em duas fases), a sensibilidade desses componentes pode ser gerenciada, tornando os ViTs totalmente ternários uma realidade prática para IA de borda. Os autores observam que a lacuna de precisão em relação ao FP32 escala inversamente com a capacidade do modelo, com modelos maiores (DeiT-III-S384) sofrendo menos perda relativa de precisão do que os menores.