Enhancing Layer Interaction Using Key-Correlated Layer Attention

Resumo Técnico: Melhorando a Interação entre Camadas Usando Atenção de Camada Correlacionada por Chave

Definição do Problema
Avanços recentes no design de redes neurais profundas (DNN) utilizaram a Atenção de Camada (Layer Attention) para melhorar as interações entre camadas, tratando as saídas das camadas como tokens de uma sequência. No entanto, a atenção de camada padrão sofre de complexidade computacional quadrática ($O(D^2)$) em relação à profundidade da rede $D$, juntamente com um crescimento significativo de memória devido à necessidade de armazenar em cache as saídas de todas as camadas precedentes.

Para abordar esses problemas, trabalhos anteriores introduziram a Atenção de Camada Recorrente (RLA) e aproximações lineares como a MRLA-L. Embora a RLA reduza o overhead computacional ao armazenar em cache as representações de Chave ($K$) e Valor ($V$), ela funciona como um mecanismo estático, onde as características das camadas iniciais permanecem fixas uma vez geradas, limitando as atualizações dinâmicas de informação. Além disso, aproximações lineares como a MRLA-L tentam reduzir a complexidade aproximando linearmente a saída da atenção. Os autores observam que essa abordagem faz com que as contribuições das camadas anteriores sofram uma atenuação exponencial (conforme mostrado em sua análise dos coeficientes de decaimento $\beta_l$), tornando difícil estabelecer dependências de longo alcance entre camadas eficazes. Adicionalmente, métodos de linearização existentes frequentemente falham em abordar o gargalo de memória do armazenamento em cache das saídas de camadas anteriores ou os desafios específicos de aplicar atenção linear a estruturas de DNN hierárquicas.

Metodologia: Atenção de Camada Correlacionada por Chave (KCLA)
A proposta Atenção de Camada Correlacionada por Chave (KCLA) é um novo método de atenção de camada linearizado, projetado para alcançar complexidade computacional linear enquanto preserva atualizações dinâmicas e a modelagem de dependências de longo alcance. A metodologia é construída sobre três insights e mecanismos centrais:

1. Observação de Correlação de Chave: A análise empírica revela que as representações de Chave em diferentes camadas na atenção de camada padrão exibem uma alta similaridade de cosseno. Aproveitando isso, a KCLA aproxima a saída da atenção de camada padrão fatorando a transformação linear do Valor ($V$) e aproximando as representações de Chave como sendo direcionalmente correlacionadas. Isso permite a agregação das saídas das camadas em um estado recorrente ($U^t$) sem armazenar o histórico completo das matrizes $K$ e $V$.
2. Restauração de Softmax Dinâmica: Um desafio crítico na linearização da atenção é a perda da operação softmax, que é essencial para normalizar os escores de atenção e permitir ajustes dinâmicos de distribuição. A KCLA introduz um mecanismo de atenção multi-cabeça modificado para restaurar essa capacidade:
   • Ela utiliza múltiplas cabeças com diferentes fatores de escala (temperaturas) derivados de uma sequência logaritmicamente uniforme.
   • Utiliza uma transformação linear aprendível $\Phi$ para mapear a norma da Consulta ($\|Q^t\|$) e o ângulo de correlação ($\cos \theta$) em pesos específicos de cada cabeça.
   • Isso permite que o modelo ajuste dinamicamente a distribuição de atenção em cada camada, simulando efetivamente o papel do parâmetro de temperatura do softmax sem computar a matriz completa $N \times N$.
3. Adaptação de Estágio Cruzado: Para acomodar DNNs com tamanhos de mapa de características e dimensões de canal variados entre estágios, a KCLA implementa um procedimento de adaptação específico (Algoritmo 1). Isso envolve a normalização das características acumuladas, aplicação de uma convolução $1 \times 1$ com stride para downsampling e o recálculo dos fatores de normalização para unir as disparidades dimensionais entre os estágios.

Principais Contribuições

• Observação e Prova: Os autores demonstram que as representações de Chave em diferentes camadas na atenção de camada padrão são altamente correlacionadas em direção, fornecendo a base teórica para a aproximação linear.
• Nova Linearização: A KCLA é proposta como um método de atenção de camada linearizado que evita a natureza estática da RLA e a atenuação exponencial da MRLA-L. Ela alcança complexidade computacional linear e complexidade espacial fixa, independente da profundidade da rede.
• Aproximação de Softmax Dinâmica: Um novo método é introduzido para aproximar a operação softmax usando escala multi-cabeça e transformações aprendíveis, permitindo o ajuste dinâmico e adaptativo das distribuições de atenção em cada camada.

Resultados Experimentais
Os autores avaliaram a KCLA em três tarefas de visão fundamentais: classificação de imagens, detecção de objetos e segmentação de imagens médicas.

• Classificação de Imagens: No CIFAR-100, a KCLA alcançou o melhor desempenho em profundidades de ResNet de 20, 56 e 110, superando a MRLA-L, DLA-L e a atenção de camada padrão. No ImageNet-1K, a integração da KCLA no ResNet-50 e ResNet-101 melhorou a precisão Top-1 em +1,7% e +1,4%, respectivamente. Ela manteve uma vantagem leve, porém consistente, sobre a MRLA-L, utilizando significativamente menos parâmetros do que o método dinâmico DLA-L.
• Detecção de Objetos: No conjunto de dados COCO2017 usando Faster R-CNN e Mask R-CNN, a KCLA demonstrou um desempenho competitivo. Com ResNet-101, ela superou a arquitetura mais complexa DLA-L em diversas métricas (AP$_{75}$, AP$_S$, AP$_M$) utilizando 2,2 milhões de parâmetros a menos.
• Segmentação de Imagens Médicas: Nos datasets ISIC2017, ISIC2018 e Kvasir-SEG, a KCLA superou consistentemente a MRLA-L e o baseline ResNet-UNet, alcançando scores superiores de mIoU e DSC com overhead mínimo de parâmetros.
• Análise de Complexidade: Diferente da MRLA-B, cuja memória e FLOPs crescem quadraticamente com a profundidade, a KCLA exibe crescimento linear tanto no custo computacional quanto no custo de memória.

Significância e Alegações
O artigo afirma que a KCLA consegue preencher a lacuna entre a eficiência da atenção linear e o poder de representação da atenção padrão. Ao aproximar diretamente o mecanismo de atenção de camada padrão com base nas correlações de Chave, a KCLA evita os problemas de representação de característica estática da RLA e a degradação de dependência de longo alcance da MRLA-L.

Os autores enfatizam que seu método mantém a capacidade de modelar dependências de longo alcance entre camadas de forma eficiente, oferecendo um equilíbrio desejável entre eficiência e desempenho. Eles observam que o método proposto para restaurar a função softmax é um fator crítico para alcançar esses resultados. Embora a implementação atual dependa de parâmetros de temperatura pré-definidos, os autores identificam o design de uma atenção de camada linear de multi-cabeça com temperatura auto-adaptável e dinâmica como uma direção para otimização futura. O código está disponível publicamente para facilitar a reprodutibilidade.