Putting Registers to Work: Task Registers for Token Pruning in Vision Transformers

Resumo Técnico: Colocando Registradores para Trabalhar: Registradores de Tarefa para Poda de Tokens em Vision Transformers

Declaração do Problema

Os Vision Transformers (ViTs) exibem um escalonamento computacional quadrático com o número de tokens, o que motiva a poda de tokens para reduzir os custos de inferência. Embora os critérios de poda (ex: escores de atenção, cobertura espacial) sejam frequentemente desenvolvidos para classificação de imagens, os backbones de ViT pré-treinados são cada vez mais adaptados para tarefas de predição densa, como segmentação semântica e detecção de objetos.

Os autores identificam uma lacuna crítica: políticas de poda projetadas para uma tarefa não se transferem necessariamente de forma eficaz para outras. Através de experimentos de "sonda" controlados, o artigo demonstra que:

1. Divergência de Classificação de Critérios: Tarefas de segmentação e detecção classificam os critérios de poda de forma diferente. Por exemplo, a poda baseada em atenção causa quedas significativas de precisão nas camadas iniciais para classificação, enquanto a cobertura espacial é mais estável. Inversamente, o critério que minimiza o erro de detecção frequentemente apresenta um desempenho ruim para segmentação.
2. Sensibilidade à Profundidade: A eficácia de um critério de poda varia significamente através das profundidades da rede e difere entre as tarefas.
3. Conflitos de Recuperação: Tarefas de predição densa requerem a reconstrução de características nas posições removidas. O artigo descobre que a segmentação e a detecção preferem endpoints de recuperação opostos: a segmentação beneficia-se ao restaurar o deslocamento completo da característica (diferença entre os tokens removidos e os retidos), enquanto a detecção apresenta melhor desempenho usando apenas o token "substituto" (stand-in) retido.

Métodos existentes tipicamente aplicam uma política de poda única e estática ou utilizam seletores agnósticos à tarefa, falhando em considerar essas variações específicas de tarefa e dependentes de profundidade.

Metodologia: Poda Adaptativa à Tarefa (TAP)

O artigo propõe a Poda Adaptativa à Tarefa (TAP), um framework que utiliza Registradores de Tarefa para controlar dinamicamente a poda de tokens, a alocação de orçamento e a recuperação de características.

Mecanismo Central:

• Registradores de Tarefa: Em vez de usar registradores meramente como armazenamento agnóstico à tarefa para artefatos de características (como no trabalho anterior de Darcet et al., 2024), o TAP introduz um registrador inicial treinável por tarefa. Durante uma passagem direta (forward pass), apenas o registrador associado à tarefa atual está ativo.
• Estado Evolutivo: O registrador ativo evolui juntamente com os tokens da imagem através do backbone. Seu estado em qualquer camada condiciona três decisões fundamentais:
  1. Seleção de Tokens: O registrador computa um escore para cada patch candidato reutilizando as projeções de query e key existentes do bloco. Os tokens são classificados, e um subconjunto é retido para auto-atenção.
  2. Alocação Exata de Orçamento: Uma leitura linear compartilhada, condicionada ao estado do registrador, determina o número inteiro exato de tokens a serem removidos em cada camada de poda. Isso garante que o orçamento global de remoção seja atendido exatamente sem exigir uma perda de taxa (rate loss).
  3. Escalonamento de Recuperação: Para tarefas densas, o registrador controla um fator de escala ($\alpha$) que determina quanto do deslocamento de característica armazenado (diferença entre o token removido e seu substituto) é restaurado na leitura.

Variantes de Implementação:

• TAP-J (Fine-tuning Conjunto): Cada tarefa possui um backbone, um registrador de tarefa e leituras (readouts) separados. Isso permite especialização total.
• TAP-F (Base Congelada): Um único backbone congelado é compartilhado entre as tarefas. Registradores específicos da tarefa, atualizações de baixo posto (LoRA) e cabeças são treinados conjuntamente. Isso reduz o armazenamento de parâmetros enquanto mantém a adaptabilidade.

Recuperação de Características:
Tokens removidos são correspondidos deterministicamente ao token retido com a maior similaridade de cosseno no espaço de chaves (o "substituto"). A diferença de característica (deslocamento) é armazenada. Nos pontos de leitura densa, a reconstrução é computada como $\hat{x}_i = x_{\text{substituto}} + \alpha \cdot \delta_i$, onde $\alpha$ é aprendido por tarefa e profundidade. Tarefas de classificação pulam este passo de recuperação.

Principais Contribuições

1. Framework Controlado: Os autores desenvolvem um framework que congela checkpoints sem poda e aplica critérios de redução sem parâmetros camada por camada para isolar o impacto das escolhas de poda entre classificação, segmentação e detecção. Isso revela o desalinhamento na classificação de critérios e preferências de recuperação.
2. Registradores de Tarefa como Controladores: O artigo redefine o papel dos registradores de armazenamento passivo para controladores ativos e específicos de tarefa da computação esparsa. O registrador coordena a seleção de tokens, a distribuição do orçamento de remoção através da profundidade e o escalonamento de recuperação.
3. Poda e Recuperação Unificadas: O TAP trata a remoção de tokens e a recuperação de características como uma operação unificada. Ele associa tokens descartados aos seus substitutos e escala a restauração dos deslocamentos com base no registrador da tarefa, mantendo as características reconstruídas fora do encoder para manter a esparsidade.

Resultados Experimentais

Os autores avaliam o TAP no ImageNet-1K (classificação), ADE20K (segmentação) e COCO (detecção/segmentação de instância) usando um backbone ViT-B com $\rho = 0.5$ (50% de tokens retidos).

• Desempenho:
  • TAP-J alcança as menores quedas de precisão entre os métodos de fine-tuning total. No ADE20K, atinge 47.0 mIoU (apenas 0.2 de queda em relação ao modelo sem poda) com 1.30× de throughput do encoder. No COCO, atinge 53.7 box AP (0.3 de queda) com 1.32× de throughput. No ImageNet-1K, atinge 83.2 top-1 (0.4 de queda).
  • TAP-F (backbone compartilhado) segue ligeiramente atrás do TAP-J, mas supera significativamente regras estáticas e outros baselines, alcançando 45.8 mIoU no ADE20K e 52.5 box AP no COCO com overhead de parâmetros mínimo (~1.2M de parâmetros específicos de tarefa).
• Eficiência: O TAP-J fornece acelerações no encoder de 1.26× a 1.32× mantendo uma precisão competitiva. O overhead do controlador (query do registrador, leitura de orçamento, contabilidade) representa apenas ~0.7% do tempo de execução do encoder.
• Memória: O TAP reduz o pico de memória de inferência em 9–18% em comparação com baselines densos, sendo que a segmentação apresenta os maiores ganhos devido à redução nos tamanhos de ativação nos blocos intermediários.
• Estudos de Ablação: Remover o registrador de tarefa (usando um compartilhado) ou fixar a alocação/recuperação para valores estáticos degrada significativamente o desempenho, confirmando a necessidade de controle adaptativo à tarefa e à profundidade.

Significância e Alegações

O artigo alega que a redução de tokens não pode ser tratada como uma regra única e transferível. À medida que as representações do ViT evoluem de estruturas locais para semânticas de alto nível, o valor dos tokens muda tanto com a profundidade quanto com a tarefa específica que lê o backbone.

A significância do TAP reside em sua capacidade de adaptar a computação esparsa às necessidades específicas de diferentes tarefas visuais usando um backbone compartilhado. Ao colocar os registradores de tarefa evolutivos para trabalhar, o TAP:

• Resolve conflitos nas classificações de critérios de poda entre tarefas.
• Otimiza a distribuição do orçamento de remoção através da profundidade da rede.
• Condiciona a recuperação de características aos requisitos específicos de leituras densas vs. esparsas.

Os autores posicionam o TAP como um passo em direção a sistemas de visão unificados que podem servir múltiplas tarefas através de um backbone compartilhado com armazenamento específico de tarefa mínimo, equilibrando o desempenho preditivo com o throughput do encoder de forma mais eficaz do que políticas de poda estáticas. O artigo permanece modesto, observando que o TAP-J fica 0.1 ponto atrás do melhor método exclusivo de classificação (DiffRate), provavelmente porque a classificação depende de um token de classe e pula a recuperação densa, deixando um dos mecanismos do TAP subutilizado para essa tarefa específica. Sugere-se que trabalhos futuros explorem orçamentos condicionados pela imagem e múltiplos substitutos.