Geometry-Guided Layerwise FFN Width Allocation in Transformers

Resumo Técnico: Alocação de Largura de FFN por Camadas Guiada pela Geometria em Transformers

Declaração do Problema
Nas arquiteturas Transformer padrão, a largura oculta das Redes de Alimentação Direta (FFNs) é tipicamente constante em todas as camadas, apesar das evidências de que as funções das camadas e a dinâmica das representações mudam com a profundidade. Embora trabalhos recentes (Tapered Language Models) tenham demonstrado que a realocação de capacidade das camadas mais profundas para as mais rasas melhora o desempenho, os métodos existentes dependem de cronogramas monotônicos projetados à mão (ex: decaimentos de cosseno). Esses formatos prescritos não se adaptam às demandas computacionais específicas de um determinado modelo ou distribuição de dados. O artigo questiona se uma alocação de largura dependente dos dados e específica por camada pode ser derivada diretamente do comportamento geométrico das ativações durante um passo de propagação direta (forward pass), em vez de depender de um prior arquitetônico fixo.

Metodologia
Os autores propõem um framework que trata as camadas FFN como transportadoras de nuvens de representação de tokens. A metodologia central envolve três estágios:

1. Medição de Trabalho Geométrico:
   O artigo define "trabalho geométrico" como a mudança induzida por uma camada FFN em uma nuvem de representações de tokens. Para garantir que a métrica reflita a contribuição real da FFN, e não artefatos de escala residual, os autores avaliam três estatísticas de trabalho distintas sob três frameworks geométricos:

   • Geometria Bruta (Raw Geometry): Utiliza distância Euclidiana em estados não normalizados. É sensível à escala e serve como um controle.
   • Geometria Esférica: Utiliza estados unitários normalizados e distância angular.
   • Geometria Hiperbólica: Mapeia estados unitários normalizados em uma bola de Poincaré, atuando como uma reponderação não linear das distâncias angulares.

   Para cada geometria, três estatísticas de trabalho são computadas:

   • Deslocamento (Shift): O deslocamento médio dos tokens (preservando a correspondência).
   • Distorção Métrica: Medida via distância de Gromov-Wasserstein (GW), quantificando como as relações par a par são preservadas ou alteradas.
   • Mudança Topológica: Medida via distância de 2-Wasserstein entre diagramas de homologia persistente de grau um (filtração de Vietoris–Rips), capturando a criação ou destruição de loops.

2. Surrogate de Aproximação e Otimização:
   Os autores introduzem um modelo substituto (surrogate) para mapear o trabalho geométrico para larguras ideais de camada. Eles postulam que o erro de aproximação de uma camada FFN escala com sua largura de acordo com uma dimensão intrínseca ($d_{int}$) e uma sensibilidade de perda ($\lambda$).

   • Premissas: (1) Os efeitos de perda são localmente aditivos; (2) O erro escala como $w^{-a}$ onde $a$ depende da dimensão intrínseca; (3) O trabalho geométrico medido é um proxy para a dificuldade de aproximação.
   • Otimizador: Um problema de otimização restrita é resolvido para minimizar a soma dos erros de aproximação ponderados sujeitos a um orçamento total de parâmetros fixo. Isso resulta em uma solução de forma fechada para a largura ótima $w_\ell$ de cada camada, incorporando limites inferiores e expoentes dependentes da profundidade.

3. Construção do Cronograma (Schedule):
   Um modelo de referência (largura uniforme) é executado para coletar nuvens de ativação. O perfil de trabalho geométrico é estimado, suavizado e mapeado para um cronograma de largura usando o otimizador derivado. Este cronograma é então aplicado para treinar um novo modelo do zero.

Principais Contribuições

1. Definição de Trabalho Geométrico: O artigo formaliza três estatísticas complementares de trabalho de propagação direta (deslocamento, distorção GW, mudança topológica) sob métricas brutas, esféricas e hiperbólicas, distinguindo entre magnitude bruta e mudanças estruturais normalizadas por escala.
2. Solver de Orçamento Exato: Deriva um otimizador contínuo único para alocação de largura de camada sob um orçamento de parâmetros fixo, lidando com limites inferiores e dimensões intrínsecas variáveis entre as camadas.
3. Diagnósticos Abrangentes: O estudo avalia a sensibilidade de escala, a aditividade de perda local, a dimensão intrínseca e a correlação entre o trabalho geométrico e a sensibilidade da camada baseada em perturbação em sete modelos pré-treinados (Llama 3.1/3.2, Gemma 2/3, Mistral).
4. Validação Empírica: O artigo compara cronogramas derivados da geometria contra largura uniforme, um decaimento de cosseno desenhado à mão e um controle "anti-topológico" através de três escalas de modelos (128M, 256M e 440M).

Resultos

• Sensibilidade de Escala: O trabalho geométrico bruto (Euclidiano) rastreia amplamente o crescimento da norma residual, enquanto métricas normalizadas (esférica/hiperbólica) revelam um perfil "front-loaded" onde as primeiras camadas exibem o maior trabalho geométrico.
• Sensibilidade vs. Alocação: Embora o trabalho de Gromov-Wasserstein mostre uma correlação mais forte com a sensibilidade da camada baseada em perturbação, a alocação topológica/hiperbólica resultou no menor erro de validação médio nos treinamentos. Isso sugere que sensibilidade e dificuldade de aproximação são fatores distintos.
• Desempenho de Treinamento:
  • Nos experimentos de 128M e 256M, os cronogramas de geometria normalizada geralmente superaram a largura uniforme e o controle anti-topológico, embora os resultados sejam exploratórios devido ao número limitado de sementes (seeds).
  • No experimento de cinco sementes de 440M, o cronograma topológico/hiperbólico alcançou uma redução média de perda de validação de 0.019 em relação à largura uniforme. Isso é aproximadamente 6,3 vezes a redução alcançada pelo decaimento de cosseno do estilo TLM (0.003).
  • O controle anti-topológico/bruto teve um desempenho pior que a largura uniforme, confirmando que a direção da alocação importa.

Significância e Alegações
O artigo afirma que o trabalho geométrico fornece um método principiado e baseado em dados para substituir constantes de largura de FFN globais por cronogramas dependentes de camada. Estabelece que:

• A capacidade da FFN não deve ser assumida como uniforme; as atualizações das camadas mudam de papel conforme a profundidade.
• Decaimentos monotônicos fixos não são a única alternativa eficaz; perfis derivados de ativação podem capturar estruturas não monotônicas e alcançar ganhos de desempenho maiores do que cronogramas desenhados à mão sob condições idênticas.
• Mudanças geométricas normalizadas (especificamente topologia/hiperbólica) revelam um perfil "early-heavy" transferível que melhora a eficiência do modelo.

Os autores mantêm modéstia em suas conclusões, observando que os resultados são comparações de médias descritivas até que a incerteza de diferença pareada seja totalmente reportada. Eles não alegam ter provado a superioridade universal de uma geometria específica (ex: topologia vs. GW), mas sim que a mudança geométrica normalizada é um sinal útil para alocação de capacidade. O trabalho posiciona a alocação derivada de ativação como uma extensão prática do tapering, motivando trabalhos futuros para calibrar diretamente o trabalho geométrico à demanda de aproximação por camada.