Reg4Pru: Regularisation Through Random Token Routing for Token Pruning

Resumo Técnico: Reg4Pru – Regularização por meio de Roteamento Aleatório de Tokens para Poda de Tokens

Problema
Os Vision Transformers (ViTs) tornaram-se fundamentais para modelos de visão modernos devido às suas capacidades de escalabilidade e generalização. No entanto, sua dependência de mecanismos de autoatenção resulta em uma complexidade computacional que escala quadraticamente ($O(N^2)$) com o número de tokens $N$. Este gargalo é particularmente agudo na segmentação de imagens médicas, onde são necessários inputs de alta resolução para preservar detalhes diagnósticos críticos.

Embora os métodos de poda de tokens (remoção de tokens redundantes) e as estratégias de reativação (reintrodução de tokens podados durante a inferência para previsões densas pixel a pixel) ofereçam um caminho para a eficiência, os autores observam uma degradação significativa de desempenho nas camadas profundas. Especificamente, quando os tokens podados são reativados durante a inferência, eles exibem representações drasticamente diferentes de seu estado durante o treinamento. Esse desvio de distribuição entre treino e teste leva à instabilidade e a uma baixa performance de predição densa, causando, de forma contraintuitiva, uma queda no desempenho à medida que a profundidade do modelo aumenta.

Metodologia
O artigo propõe o Reg4Pru, uma técnica de regularização em tempo de treinamento projetada para mitigar a instabilidade associada à poda e reativação de tokens. O mecanismo central envolve o roteamento aleatório de tokens durante a fase de treinamento para simular os efeitos da poda sem descartar permanentemente os tokens.

1. Roteamento Aleatório de Tokens: Durante a passagem direta (forward pass) do treinamento, o modelo seleciona um intervalo de roteamento aleatório $[l, n]$ entre blocos intermediários do transformer. Um subconjunto aleatório de tokens ($\tau_r$) é escolhido para pular o cálculo dos blocos de $l$ a $n$, enquanto os tokens restantes ($\tau_k$) são processados normalmente.
2. Fluxo de Gradiente e Representação: Os tokens roteados ignoram blocos específicos, cortando efetivamente seus gradientes em relação a essas camadas. No entanto, ao randomizar tanto o subconjunto de tokens roteados quanto o intervalo de roteamento $[l, n]$ a cada etapa, diferentes tokens contribuem com gradientes para diferentes blocos ao longo das iterações. Isso força o modelo a aprender representações invariantes à profundidade, garantindo que os tokens possam ser reativados em qualquer profundidade durante a inferência sem sofrerem de incompatibilidade de representação.
3. Estratégia de Inferência: Na inferência, uma rede de política separada prevê quais tokens devem ser podados com base em um cronograma hierárquico de razão de manutenção (keep-ratio). Os tokens podados são temporariamente removidos para economizar computação e são reativados no bloco final para gerar máscaras de segmentação densas.
4. Funções de Perda: O objetivo de treinamento combina uma perda padrão baseada em razão (para impor orçamentos de tokens) e, opcionalmente, uma perda de política informada ($L_{pol}$) que utiliza máscaras de segmentação de verdade fundamental (ground truth) para guiar as decisões de poda. O mecanismo de roteamento atua como um regularizador ao lado dessas perdas.

Principais Contribuições

• Análise de Instabilidade por Profundidade: Os autores identificam e caracterizam uma instabilidade específica na poda de tokens para predição densa, onde o desempenho degrada em camadas mais profundas devido à reativação de tokens que não foram expostos à profundidade total da rede durante o treinamento.
• Framework Reg4Pru: Eles introduzem uma estratégia de roteamento de tokens em tempo de treinamento com limites aleatórios. Diferente de métodos de roteamento anteriores que utilizam limites fixos ou políticas de roteamento aprendidas, o Reg4Pru randomiza o intervalo de roteamento para aproximar a distribuição de características em todas as profundidades de poda, thereby regularizando o modelo.
• Validação Empírica: O método é validado no conjunto de dados de segmentação de vasos sanguíneos retinianos FIVES, demonstrando que o roteamento aleatório estabiliza significativamente o desempenho em comparação com o roteamento de limites fixos ou sem roteamento.

Resultos
Experimentos foram conduzidos no dataset FIVES (imagens de 2048$\times$2048, redimensionadas para quadrantes para testes de throughput) usando uma arquitetura EoMT inicializada com DINOv2-small.

• Ganhos de Desempenho: Comparado a um modelo podado treinado sem roteamento, o Reg4Pru melhorou a Precisão Média (AP) em um valor absoluto de 46% (de 25,60% para 71,67% para a configuração $L_{ratio}$).
• Eficiência: A configuração Reg4Pru alcançou um aceleração relativa de 29% no tempo de execução (wall-clock time) em comparação com o baseline não podado, mantendo alta qualidade de segmentação.
• Estabilidade: A análise de profundidade (Figura 5) mostrou que, enquanto os modelos podados sem roteamento sofriam uma queda acentuada de AP em blocos mais profundos, o Reg4Pru manteve o desempenho estável. O roteamento de limites fixos mitigou parte da instabilidade, mas falhou em evitar quedas de desempenho após o término do intervalo fixo, confirmando a necessidade de randomização.
• Comparação: O Reg4Pru superou os métodos de Token Merging (ToMe) em termos de AP ao comparar configurações de throughput similares, particularmente quando o ToMe foi aplicado sobre as camadas MHSA e MLP.

Significância e Alegações
O artigo afirma que o Reg4Pru é um regularizador valioso para estratégias de redução de tokens em tarefas de predição densa. Os autores argumentam que a instabilidade observada na poda é resultado de um desvio de distribuição de treino-teste em relação às representações dos tokens em diferentes profundidades. Ao randomizar o intervalo de roteamento durante o treinamento, o Reg4Pru força o modelo a aprender representações que são robustas ao serem ignoradas em várias profundidades, efetivamente preenchendo a lacuna entre o treinamento e a inferência.

Os autores mantêm a modéstia quanto à generalização de seus achados, observando que os experimentos foram conduzidos em um único conjunto de dados médicos de alta resolução. Eles sugerem que, embora o método seja eficaz para segmentação, avaliações adicionais são necessárias para verificar sua aplicabilidade em outras tarefas de predição densa (ex: estimativa de profundidade) e seu potencial como um mecanismo de regularização geral quando combinado com técnicas como dropout ou stochastic depth. Eles também reconhecem que, embora tenham explorado políticas informadas, as políticas baseadas em razão produzem atualmente resultados superiores em sua configuração específica.