Resumo Técnico: ZO-Act

Declaração do Problema

O ajuste fino (fine-tuning) de Grandes Modelos de Linguagem (LLMs) normalmente depende de otimização de primeira ordem (backpropagation), que se torna proibitiva em termos de memória à medida que os modelos escalam para bilhões de parâmetros devido aos requisitos de armazenamento de ativações, estados do otimizador e computações de gradiente. A otimização de ordem zero (ZO) oferece uma alternativa baseada apenas no forward pass, estimando direções de atualização usando apenas avaliações de perda, evitando assim o backpropagation. No entanto, os métodos ZO existentes enfrentam desafios significativos:

1. Alta Variância: Perturbar o espaço total de parâmetros (como no MeZO) resulta em estimativas de gradiente de alta variância, levando a uma convergência lenta e instabilidade.
2. Subespaços Subótimos: Embora alguns métodos restrinjam as perturbações a subespaços de baixa dimensão (ex: LOZO, SubZero), eles frequentemente dependem de subespaços construídos aleatoriamente ou perturbações aleatórias que podem não se alinhar com o fluxo real do gradiente.
3. Compatibilidade com Otimizadores: Muitos métodos ZO têm dificuldade em utilizar efetivamente otimizadores padrão baseados em momentum (como Adam), porque a estrutura de perturbação não expõe variáveis treináveis explícitas e leves.
4. Restrições de Quantização: O ajuste fino de LLMs quantizados (ex: INT4) é difícil com métodos que exigem a materialização de perturbações de peso total.

Metodologia: ZO-Act

Os autores propõem o ZO-Act, um método de ajuste fino ZO de um único passo (one-shot) informado pela ativação. A ideia central é que os componentes dominantes dos gradientes de peso em camadas lineares residem no subespaço gerado pelos vetores singulares à direita da matriz de ativação de entrada.

Mecanismo Central

1. Inicialização de Passo Único (One-Shot): Para cada camada linear, o ZO-Act realiza um único forward pass em um lote de calibração para computar a matriz de ativação de entrada $X$. Em seguida, computa os $r$ principais vetores singulares à direita ($V_r$) de $X$ via uma SVD (Decomposição de Valor Singular) fina.
2. Parametrização de Subespaço Fixo: A atualização de peso é restrita a um subespaço de baixo posto definido por esses vetores de ativação congelados. A atualização de peso efetiva é parametrizada como:
   $$\Delta W = V_r B$$
   onde $V_r \in \mathbb{R}^{m \times r}$ é a base de ativação congelada, e $B \in \mathbb{R}^{r \times n}$ é uma matriz de coeficientes leve e treinável. Os pesos pré-treinados originais $W$ e a base $V_r$ permanecem congelados durante todo o treinamento.
3. Otimização no Espaço de Coeficientes: Em vez de perturbar a matriz de peso total $W$, o ZO-Act amostra perturbações aleatórias $Z$ no espaço de coeficientes de baixa dimensão ($Z \sim \mathcal{N}(0, I)$). O peso efetivo perturbado torna-se $W_{eff} = W + V_r(B + \mu Z)$.
4. Estimativa Apenas de Forward: O gradiente para $B$ é estimado usando avaliações de perda de diferença finita:
   $$\hat{g}_B = \frac{L(B + \mu Z) - L(B)}{\mu} Z$$
   Este gradiente estimado é então usado para atualizar $B$ usando otimizadores padrão de primeira ordem, como Adam.

Análise Teórica

O artigo analisa o ZO-Act como uma otimização de ordem zero sobre um espaço de coeficientes restrito $\beta$ (dimensão $k$) em vez do espaço de pesos total (dimensão $d$).

• Redução de Variância: A variância do estimador ZO escala linearmente com a dimensão da perturbação. Ao reduzir a dimensão de $d$ (pesos totais) para $k = \sum r n_\ell$ (coeficientes, onde $r \ll m$), o ZO-Act reduz significamente o termo de convergência dependente da variância e o erro de diferença finita.
• Compromisso Viés-Variância (Bias-Variance Trade-off): Restringir as atualizações ao subespaço informado pela ativação introduz um viés de aproximação de subespaço. No entanto, os autores argumentam que as ativações e gradientes de LLMs exibem estruturas de baixo posto, o que significa que as direções de atualização dominantes são bem capturadas pelos principais vetores singulares das ativações de entrada. Assim, o viés é controlado e mitigado pela natureza informada pelos dados do subespaço.

Principais Contribuições

1. Subespaço Informado pela Ativação: Diferente de métodos de subespaço anteriores que utilizam matrizes de baixo posto aleatórias ou estáticas, o ZO-Act deriva o subespaço de perturbação diretamente das ativações de entrada, garantindo que o subespaço se alinhe com as direções de ativação dominantes da camada.
2. Otimização de Coeficientes Explícita: Ao congelar a base e otimizar apenas a matriz de coeficientes $B$, o ZO-Act transforma o ajuste fino ZO em um problema de otimização explícito sobre um espaço de baixa dimensão. Isso permite a aplicação direta de otimizadores baseados em momentum (ex: Adam) e evita a materialização de perturbações de peso total.
3. Compatibilidade com Quantização: O método suporta naturalmente LLMs quantizados (ex: INT4) porque os pesos de baixa precisão permanecem congelados, e apenas as matrizes de coeficientes de baixa precisão são atualizadas.
4. Garantias Teóricas: O artigo fornece uma análise de convergência mostrando que o ZO-Act reduz tanto a variância do estimador de gradiente quanto o erro de diferença finita em comparação com o ZO de peso total, ao custo de um viés de subespaço gerenciável.

Resultados Experimentais

Os autores avaliaram o ZO-Act no Llama-3-8B, OPT-13B e um Llama-3-8B quantizado em INT4 através de tarefas de compreensão de linguagem, resposta a perguntas e raciocínio de senso comum.

• Desempenho: O ZO-Act superou consistentemente fortes baselines de ZO (incluindo MeZO, LOZO, SubZero, AGZO e ZO-Muon).
  • No Llama-3-8B, alcançou os melhores resultados de ZO de precisão total em SST-2, RTE, BoolQ e WiC. Notavelmente, melhorou o desempenho de RTE de 81.2 (ZO-Muon) para 87.0 e CB de 69.6 para 89.3.
  • No OPT-13B, obteve resultados semelhantes de topo em SST-2, RTE, CB e BoolQ.
  • Em Modelos Quantizados em INT4, o ZO-Act permaneceu competitivo com os baselines de ZO de precisão total, demonstrando sua eficácia para ajuste fino com restrição de memória.
• Eficiência: O ZO-Act igualou o menor uso de memória entre os métodos comparados (aprox. 16.1 GB para Llama-3-8B, caindo para 5.8 GB para INT4) e alcançou o tempo de execução mais rápido (42.0 minutos vs. 44.5+ minutos para outros) sob um orçamento fixo de consultas de forward.
• Sensibilidade ao Posto (Rank): Experimentos mostraram que um posto de $r=1$ muitas vezes era suficiente para modelos de precisão total, produzindo a melhor acurácia enquanto minimizava a dimensão da perturbação. Postos mais altos ($r=32$) foram benéficos para modelos quantizados para compensar a redução de capacidade.
• Estabilidade: Estudos diagnósticos confirmaram que o subespaço informado pela ativação permanece estável durante o ajuste fino e captura uma parte significativa da energia do gradiente total, validando o design de inicialização de "passo único".

Significância e Alegações

O artigo afirma que o ZO-Act representa um passo significativo para tornar o ajuste fino de ordem zero uma alternativa prática e de alto desempenho para LLMs grandes e quantizados. Sua significância reside em:

• Reduzir a Lacuna: Ele estreita a lacuna de desempenho entre o ajuste fino de ordem zero e o de primeira ordem ao reduzir a variância do estimador através de subespaços informados pelos dados.
• • Habilitar Otimizadores: Resolve a dificuldade de aplicar otimizadores padrão como Adam a métodos ZO ao expor variáveis treináveis explícitas e de baixa dimensão.
• Escalabilidade: Fornece uma solução eficiente em termos de memória para o ajuste fino de modelos quantizados, o que é crítico para implantação em hardware com restrição de recursos.

Os autores reconhecem limitações, observando que o ZO-Act ainda não fecha totalmente a lacuna de desempenho com métodos de primeira ordem devido ao ruído inerente da estimativa de ordem zero e ao viés de subespaço. Adicionalmente, o ZO-Act continua sendo mais lento em tempo de execução (wall-clock time) do que métodos baseados em backpropagation devido às múltiplas passagens de forward necessárias por atualização. No entanto, o método posiciona-se como uma escolha superior para cenários onde o backpropagation é proibitivo ou indisponível em termos de memória.