Resumo Técnico: RotMoLE – Aprimorando Mistura de Especialistas de Baixo RANK através de Portão Rotacional

Declaração do Problema

Embora os Modelos de Linguagem de Grande Escala (LLMs) utilizem arquiteturas de Mistura de Especialistas (MoE) e técnicas de Ajuste Fino Eficiente em Parâmetros (PEFT) como LoRA para lidar com tarefas diversas, os frameworks existentes de Mistura de Especialistas de Baixo RANK (MoE-LoRA) enfrentam limitações em cenários complexos, multi-tarefa ou multilíngues.

A questão central identificada é a limitação dos mecanismos de portão convencionais em MoE. Os roteadores padrão aplicam apenas um reponderamento escalar (escalonamento) aos especialistas selecionados. Embora isso module a magnitude da contribuição de um especialista, falta-lhe a capacidade de realizar transformações espaciais complexas nas saídas dos especialistas. Consequentemente, o espaço de soluções formado pela agregação ponderada dos especialistas restringe-se a uma combinação linear ao longo de direções espaciais fixas. Essa limitação é particularmente pronunciada quando o número de especialistas candidatos é limitado em relação ao número de tarefas (por exemplo, em cenários com restrições de recursos onde os especialistas devem ser compartilhados entre múltiplos domínios), pois o roteador não consegue distinguir ou especializar suficientemente os especialistas compartilhados além do simples escalonamento de magnitude.

Além disso, implementar transformações complexas (como transformações lineares completas) por meio de módulos de portão geralmente requer um número proibitivo de parâmetros treináveis (por exemplo, matrizes $d \times d$ para a dimensão oculta $d$), o que nega a eficiência paramétrica do PEFT.

Metodologia: RotMoLE

Os autores propõem o RotMoLE (Mistura Rotacionável de Especialistas de Baixo RANK), um framework especializado de MoE-LoRA que introduz um mecanismo de portão rotacional ao lado do portão de escalonamento tradicional.

Fundamentação Teórica

O método baseia-se na observação de que qualquer transformação entre dois vetores não nulos pode ser decomposta em uma operação de escalonamento puro e uma rotação pura dentro do subplano bidimensional (2D) gerado por esses vetores. Ao aproveitar a estrutura de baixo rank do LoRA (onde o rank $r \ll d$), o RotMoLE aplica essas transformações dentro do espaço $r$-dimensional dos adaptadores dos especialistas, em vez da dimensão oculta completa.

Componentes da Arquitetura

A atualização de pesos para o modelo é formulada como:
$$W = W_0 + \sum_{i \in TopK(x)} g_i B_i R_i A_i$$
Onde:

• $W_0$ é a matriz de pesos pré-treinada.
• $A_i$ e $B_i$ são as matrizes de baixo rank para o $i$-ésimo especialista.
• $g_i$ é o valor do portão de escalonamento (peso escalar padrão).
• $R_i$ é a matriz de rotação específica para o $i$-ésimo especialista.

O framework opera em dois modos com base no rank do LoRA $r$:

1. Rank $r = 2$ (Rotação 2D):

   • A rotação é modelada usando um único ângulo aprendível $\theta_i$.
   • A matriz de rotação $R_i$ é uma matriz de rotação 2D padrão:
     $$R_i = \begin{pmatrix} \cos \theta_i & -\sin \theta_i \\ \sin \theta_i & \cos \theta_i \end{pmatrix}$$
   • O ângulo $\theta_i$ é gerado por um portão de rotação $\Theta(x)$, que produz valores em $[-\pi, \pi]$ usando uma normalização baseada em Sigmoid para garantir estabilidade e periodicidade.

2. Rank $r > 2$ (Rotação de Alta Dimensionalidade):

   • Como um único ângulo não pode definir uma rotação em dimensões $>2$, o método constrói um plano de rotação $P_i$ específico para cada especialista.
   • O plano é definido por dois vetores de base:
     1. O embedding de entrada normalizado para o especialista ($e_1 = A_i x / |A_i x|$).
     2. Um vetor centro de cluster aprendível $q_i$ específico para o especialista, que representa o centro agregado das entradas roteadas para aquele especialista. O segundo vetor de base $e_2$ é derivado da componente de $q_i$ ortogonal a $e_1$.
   • A matriz de rotação $R_i$ é construída projetando a rotação 2D (definida por $\theta_i$) sobre esse plano $r$-dimensional gerado por $e_1$ e $e_2$.

Eficiência Paramétrica

O RotMoLE mantém alta eficiência paramétrica. Os parâmetros adicionais introduzidos são:

• Para $r=2$: $d \times n$ parâmetros (para os pesos do portão de rotação).
• Para $r>2$: $(d + r) \times n$ parâmetros (para o portão de rotação e os vetores de centro de cluster $q$).
  Isso é significativamente menor do que os $d^2 \times n$ parâmetros necessários para transformações lineares completas, mantendo a sobrecarga comparável aos próprios módulos de especialistas.

Contribuições Principais

1. Identificação de Limitações: O artigo destaca que os roteadores MoE existentes estão limitados ao escalonamento de norma, carecendo da capacidade de realizar transformações espaciais complexas necessárias para cenários diversos com contagens limitadas de especialistas.
2. Framework RotMoLE: A proposta de um framework especializado de MoE-LoRA que integra um portão de rotação, permitindo rotações 2D dentro do espaço de baixo rank dos especialistas para expandir o espaço de soluções das combinações ponderadas de especialistas.
3. Validação Empírica: Extensos experimentos em cenários de aprendizado conjunto multi-tarefa e multilíngue demonstrando que o RotMoLE supera as linhas de base convencionais de MoE-LoRA e outras variantes recentes.

Resultados Experimentais

Os autores avaliaram o RotMoLE em Llama-3.2-3B, Qwen-2.5-3B e GLM-4-9B através de vários benchmarks:

• Aprendizado Conjunto Multi-Tarefa (QA & GLUE):
  • Em uma mistura de três tarefas de QA (CommonsenseQA, OpenBookQA, SocialIQA), o RotMoLE superou consistentemente a linha de base MoE-LoRA, alcançando melhorias médias de precisão variando de 1% a 10% (média ~5%).
  • As melhorias foram mais pronunciadas quando o número de especialistas era limitado (por exemplo, $n=2$), validando a hipótese de que a rotação ajuda a distinguir especialistas compartilhados.
  • Ganhos semelhantes foram observados nas tarefas GLUE (SST-2, MRPC, QNLI).
• Aprendizado Conjunto Multilíngue:
  • Em Geração de Texto Multilíngue (MTG) e XNLI (Inferência em Linguagem Natural), o RotMoLE alcançou pontuações médias e por idioma mais altas em comparação com as linhas de base.
  • A análise de convergência indicou que o RotMoLE exibe melhores propriedades anti-sobreajuste, convergindo mais lentamente inicialmente, mas alcançando melhor desempenho final sem a saturação precoce observada nas linhas de base.
• Comparação com Variantes:
  • O RotMoLE superou outras variantes de MoE-LoRA focadas em roteamento direto (MoCLE, MoLA, AdaMoLE) em termos de precisão geral e por tarefa.
• Estudos de Ablação:
  • Experimentos confirmaram que os ganhos de desempenho devem-se à estrutura rotacional e não simplesmente ao aumento da contagem de parâmetros. Uma linha de base equivalente em tamanho usando um portão MLP performou pior, sugerindo que o mecanismo de rotação é mais eficaz em utilizar parâmetros extras para diversidade.

Significado e Alegações

O artigo alega que o RotMoLE fornece uma função de transformação mais poderosa para especialistas de baixo rank além do simples escalonamento. Ao introduzir um portão de rotação, o framework expande significativamente o espaço de soluções global do MoE, permitindo melhor exploração e especialização dos especialistas, particularmente em cenários onde os especialistas são compartilhados entre múltiplas tarefas ou idiomas.

Os autores enfatizam que essa abordagem aborda uma lacuna específica na pesquisa atual de PEFT-MoE: a incapacidade de modelar estratégias diversas de exploração para especialistas candidatos quando o número de especialistas é limitado. Os resultados sugerem que o RotMoLE oferece um método robusto e eficiente em parâmetros para aprimorar a capacidade representacional dos LLMs em aplicações complexas e multi-domínio, sem incorrer nos custos computacionais de transformações de rank completo.