ZorBA: Zeroth-order Federated Fine-tuning of LLMs with Heterogeneous Block Activation

O artigo propõe o ZorBA, um framework de ajuste fino federado para grandes modelos de linguagem que utiliza otimização de ordem zero e ativação heterogênea de blocos para reduzir significativamente o uso de VRAM e a sobrecarga de comunicação sem comprometer a convergência.

O essencial

Imagine que você e seus amigos querem treinar um gênio superinteligente (um Modelo de Linguagem Grande, ou LLM) para escrever poemas, resolver problemas ou traduzir textos. O problema é que esse "gênio" é enorme, pesando bilhões de "células" (parâmetros), e ninguém de vocês tem uma mente grande o suficiente (memória de vídeo/VRAM) para segurá-lo todo sozinho.

Além disso, vocês não podem simplesmente enviar o cérebro do gênio de um lado para o outro pela internet, porque a conexão seria lenta demais e o arquivo pesaria uma tonelada.

Aqui entra o ZorBA, a solução proposta neste artigo. Pense no ZorBA como um sistema de treinamento em equipe inteligente e econômico. Vamos desmontar como ele funciona usando analogias do dia a dia:

1. O Problema: A "Mochila" Muito Pesada

Normalmente, para treinar uma IA, você precisa calcular como cada parte do cérebro do modelo deve mudar. Isso exige guardar muitos "rascunhos" (gradientes) na memória.

• A analogia: É como tentar montar um quebra-cabeça de 1 bilhão de peças em uma mesa de cozinha pequena. A mesa (sua memória de vídeo/VRAM) não tem espaço para todas as peças espalhadas ao mesmo tempo. Se você tentar, a mesa quebra (o computador trava).

2. A Solução Mágica: "Adivinhar" sem Ver Tudo (Otimização de Ordem Zero)

O ZorBA usa uma técnica chamada Otimização de Ordem Zero.

• A analogia: Em vez de calcular matematicamente exatamente para onde cada peça do quebra-cabeça deve ir (o que exige muita memória), o ZorBA faz um "teste e erro" inteligente. Ele dá um leve "empurrão" aleatório no modelo, vê se o resultado ficou melhor ou pior, e deduz a direção correta.
• O benefício: Você não precisa guardar os "rascunhos" complexos da matemática. Você só precisa guardar o modelo e o resultado do teste. Isso economiza muito espaço na sua mesa (VRAM).

3. O Grande Truque: "Ativação Heterogênea de Blocos"

Aqui está a parte mais criativa do ZorBA. O modelo é dividido em vários "blocos" (como capítulos de um livro). Nem todo mundo precisa ler e anotar todos os capítulos ao mesmo tempo.

• A analogia: Imagine que o modelo é um livro com 24 capítulos.
  • O jeito antigo: Todos os 50 alunos da turma têm que ler e anotar os 24 capítulos. Ninguém consegue, a mesa é pequena.
  • O jeito ZorBA: O professor (servidor central) diz: "Aluno A, você foca nos capítulos 1 a 5. Aluno B, você foca nos 6 a 10. Aluno C, nos 11 a 15...".
  • Cada aluno só precisa de espaço na mesa para os seus capítulos. No final, o professor junta as anotações de todos para ter o conhecimento completo do livro.
  • Por que isso é bom? Isso permite que alunos com mesas pequenas (celulares ou computadores fracos) participem, e o livro todo é aprendido mais rápido porque o trabalho é dividido de forma inteligente.

4. Economizando a Internet: "Sementes Compartilhadas"

Normalmente, para coordenar esse teste e erro, os alunos teriam que enviar seus "rascunhos" para o professor, o que lotaria a internet.

• A analogia: O ZorBA usa Sementes Aleatórias Compartilhadas.
  • O professor entrega a todos um "livro de receitas" (uma semente aleatória) que diz exatamente quais "empurrões" fazer.
  • Como todos têm o mesmo livro de receitas, todos geram os mesmos testes aleatórios independentemente.
  • Em vez de enviar o resultado do teste (que é grande), o aluno só envia um pequeno bilhete dizendo: "O resultado ficou X% melhor". O professor, tendo a mesma receita, consegue reconstruir a lógica inteira apenas com esse bilhete.
  • Resultado: A internet não fica congestionada. É como enviar um SMS em vez de um vídeo pesado.

5. O Algoritmo "Lexicográfico": O Maestro da Orquestra

O papel mais difícil é decidir quem lê quais capítulos para que o livro todo seja aprendido o mais rápido possível sem sobrecarregar ninguém.

• A analogia: O ZorBA usa um algoritmo inteligente (chamado de $\epsilon$-constraint lexicographic) que age como um maestro. Ele olha para a "popularidade" de cada capítulo.
  • Se um capítulo é lido por pouquíssimas pessoas, ele se torna um "gargalo" e o aprendizado fica lento.
  • O algoritmo ajusta a lista de tarefas para garantir que os capítulos mais importantes sejam lidos por várias pessoas, equilibrando a carga de trabalho (memória) com a velocidade de aprendizado.

Resumo dos Resultados

Os testes mostraram que o ZorBA é incrível:

1. Economia de Memória: Reduziu o uso de memória de vídeo em até 62% comparado a métodos antigos. É como se você pudesse treinar o modelo em computadores que antes nem conseguiam abrir o arquivo.
2. Velocidade: Convergência mais rápida (aprendizado mais rápido) do que outros métodos de otimização sem gradiente.
3. Internet Leve: O tráfego de dados entre os alunos e o professor é minúsculo, quase imperceptível.

Em suma: O ZorBA é a maneira de treinar IAs gigantes em uma rede de computadores pequenos e fracos, dividindo o trabalho de forma inteligente, economizando memória e sem lotar a internet, tudo isso sem precisar de supercomputadores caros.

Resumo técnico

Resumo Técnico: ZorBA

1. Problema Abordado

O ajuste fino (fine-tuning) de Modelos de Linguagem de Grande Escala (LLMs) em um cenário de Aprendizado Federado (FL) enfrenta dois desafios críticos:

1. Uso Excessivo de VRAM: LLMs possuem bilhões de parâmetros. O uso tradicional de retropropagação (backpropagation) exige o armazenamento de gradientes e ativações de todas as camadas, o que frequentemente excede a capacidade de memória de vídeo (VRAM) de dispositivos de borda (clientes) com recursos limitados.
2. Custo de Comunicação: A troca frequente de gradientes ou atualizações de modelos de alta dimensão entre clientes e o servidor central gera uma sobrecarga de comunicação proibitiva.
3. Limitações de Métodos Existentes: Abordagens de otimização de ordem zero (zeroth-order), que evitam o cálculo de gradientes explícitos, ainda exigem o armazenamento de ativações de todos os blocos do modelo e sofrem com taxas de convergência lentas em espaços de parâmetros de alta dimensão devido à variância introduzida pelos vetores de perturbação.

2. Metodologia Proposta: ZorBA

O ZorBA é um framework de ajuste fino federado baseado em otimização de ordem zero, projetado para mitigar os problemas acima através de três mecanismos principais:

• Otimização de Ordem Zero (Zeroth-Order Optimization):

  • Substitui a retropropagação por uma abordagem baseada apenas em passagens forward.
  • Estima os gradientes utilizando diferenças finitas de valores de função de perda gerados por vetores de perturbação aleatórios.
  • Elimina a necessidade de armazenar gradientes nos clientes, reduzindo drasticamente o uso de VRAM.

• Ativação Heterogênea de Blocos (Heterogeneous Block Activation):

  • Em vez de atualizar todos os blocos de transformadores em todos os clientes, o servidor central aloca subconjuntos diferentes de blocos para cada cliente.
  • Cada cliente atualiza apenas os blocos ativados, congelando os demais.
  • Isso permite que clientes com menos VRAM participem do treinamento, reduzindo o uso de memória proporcionalmente ao número de blocos ativados.

• Sementes Aleatórias Compartilhadas (Shared Random Seeds):

  • Para reduzir a sobrecarga de comunicação, o servidor e os clientes compartilham sementes aleatórias em vez de trocar vetores de perturbação ou gradientes completos.
  • Ambos geram independentemente os mesmos vetores de perturbação.
  • Os clientes transmitem apenas as diferenças finitas (escalares) das estimativas de gradiente, e não os vetores de gradiente de alta dimensão, permitindo que o servidor reconstrua a atualização global.

3. Análise Teórica e Formulação de Otimização

Os autores realizam uma análise teórica rigorosa para entender como as decisões de ativação de blocos afetam a convergência:

• Limites de Convergência: Derivam limites de convergência para cenários não convexos, mostrando que a taxa de convergência depende de uma métrica chamada $\Lambda(A)$, que é uma função da matriz de ativação de blocos $A$.
• Insight Chave: Minimizar $\Lambda(A)$ (o que acelera a convergência) é equivalente a maximizar a "popularidade mínima" dos blocos entre os clientes. Ou seja, é crucial garantir que os blocos menos populares (aqueles ativados por menos clientes) tenham a maior popularidade possível, equilibrando a carga.
• Trade-off: Existe um compromisso fundamental entre a taxa de convergência (que favorece mais blocos ativados) e o uso de VRAM (que favorece menos blocos).
• Algoritmo de Solução: Para resolver o problema de otimização multi-objetivo (minimizar $\Lambda(A)$ e o uso de VRAM), os autores propõem um algoritmo lexicográfico com restrição $\epsilon$:
  1. Problema PI: Maximiza a popularidade mínima de todos os clientes (determinando o limite inferior de blocos que cada cliente deve ter).
  2. Problema PII: Um algoritmo guloso que ativa blocos adicionais para minimizar o número de clientes que ainda estão no limite mínimo de popularidade, respeitando as restrições de VRAM.

4. Resultados Experimentais

O ZorBA foi avaliado em conjuntos de dados de classificação de texto (AG-News, SST-2, SNLI) utilizando os modelos OPT-125M e OPT-1.3B, comparado com baselines como FedIT, FedZO e DeComFL.

• Redução de VRAM: O ZorBA reduziu o uso total de VRAM em até 62,41% em comparação com as abordagens de base (FedIT, FedZO, DeComFL), permitindo o ajuste fino em dispositivos com recursos limitados.
• Convergência: O ZorBA convergiu mais rápido que as abordagens de ordem zero existentes (FedZO e DeComFL), demonstrando que a ativação heterogênea otimizada supera a ativação de todos os blocos em termos de eficiência de convergência.
• Custo de Comunicação: O framework gerou uma sobrecarga de comunicação insignificante em comparação com métodos que transmitem gradientes completos, sendo comparável a métodos de compressão avançados, mas com menor complexidade.

5. Significado e Contribuições

O trabalho ZorBA é significativo por:

1. Viabilizar o FL para LLMs em Dispositivos de Borda: Ao eliminar a necessidade de gradientes e permitir a ativação parcial de blocos, torna possível ajustar modelos grandes em hardware com VRAM limitada.
2. Superar a Ineficiência da Otimização de Ordem Zero: Demonstra que a otimização de ordem zero não precisa ser lenta se combinada com uma estratégia inteligente de seleção de blocos, superando a limitação de convergência lenta tradicional.
3. Redução de Comunicação: A técnica de sementes compartilhadas e transmissão de diferenças finitas oferece uma solução escalável para a comunicação em redes federadas com modelos massivos.
4. Fundação Teórica: Fornece uma análise teórica que conecta a estrutura de ativação de blocos à taxa de convergência, oferecendo um guia prático para o design de sistemas federados futuros.

Em resumo, o ZorBA representa um avanço crucial na democratização do ajuste fino de LLMs, permitindo colaboração federada eficiente, privada e viável em termos de recursos computacionais.