Resumo Técnico: PreLort – Prefix-Nested LoRA para Fine-Tuning Federado sob Heterogeneidade de Rank

1. Definição do Problema

O Aprendizado Federado (FL) permite o treinamento colaborativo de Grandes Modelos de Linguagem (LLMs) sem o compartilhamento de dados brutos, mas enfrenta desafios significativos ao adaptar modelos de fundação via métodos de ajuste fino de parâmetros eficientes (PEFT), como o Low-Rank Adaptation (LoRA). Um obstáculo primário é a heterogeneidade de rank: em configurações realistas de dispositivos cruzados (cross-device), os clientes possuem diferentes restrições computacionais e de memória, levando-os a inicializar adaptadores LoRA com diferentes ranks ($r$).

Soluções existentes para LoRA federado heterogêneo concentram-se primariamente na fase de agregação, frequentemente empregando o preenchimento com zeros (zero-padding) para alinhar adaptadores de diferentes ranks antes da média. No entanto, os autores identificam uma limitação fundamental nessas abordagens: o desalinhamento de representação. Mesmo quando a agregação é matematicamente definida, combinar modelos treinados independentemente com diferentes ranks leva a inconsistências. Especificamente, clientes de rank inferior não se beneficiam das atualizações de clientes de rank superior porque a informação não é distribuída hierarquicamente através das dimensões de rank. O zero-padding também dilui as contribuições de clientes de alto rank ao introduzir segmentos com valores zero, resultando em um uso subótimo das representações de baixo rank compartilhadas.

2. Metodologia: PreLort

Os autores propõem o PreLort (Prefix-Nested LoRA), um framework que aborda a heterogeneidade de rank estruturando conjuntamente a dinâmica de treinamento e as regras de agregação. A filosofia central é organizar as dimensões do adaptador em uma hierarquia de prefixos, garantindo que as dimensões de rank inferior codifiquem informações relevantes para a tarefa compartilhadas entre todos os clientes, enquanto as dimensões de rank superior capturem a capacidade incremental.

2.1 Estrutura de Baixo Rank Aninhada (Nested Low-Rank Structure)

O PreLort impõe uma organização de prefixo estruturada nas dimensões do adaptador. Para um cliente com rank $r_k$, apenas as primeiras $r_k$ dimensões estão ativas. Crucialmente, durante o treinamento local, o adaptador não é otimizado apenas em seu rank total. Em vez disso, ele é otimizado sob múltiplas configurações truncadas (sub-ranks de prefixo).

• Mecanismo: Um cliente com rank $r_k$ é treinado para performar bem não apenas em $r_k$, mas também em todos os níveis de rank inferiores presentes na federação ($r \in \{r \mid r \le r_k\}$).
• Objetivo: Isso força o modelo a concentrar os sinais relevantes da tarefa nas dimensões iniciais (de baixo rank), garantindo que essas dimensões permaneçam informativas e úteis mesmo quando truncadas.

2.2 Estratégia de Treinamento Local

O objetivo local para um cliente $k$ é definido como a média das perdas de entropia cruzada computadas sobre todos os sub-ranks de prefixo válidos $R_k$:
$$\mathcal{L}_k = \frac{1}{|R_k|} \sum_{r \in R_k} \mathcal{L}_{CE}(f(\cdot; \theta, r), \mathcal{D}_k)$$
Na prática, cada etapa de treinamento envolve uma passagem de forward/backward de rank total e passagens adicionais para todos os sub-ranks. Esta estratégia garante que os segmentos de prefixo sejam explicitamente otimizados na mesma granularidade usada durante a agregação, incentivando uma hierarquia consistente onde as dimensões inferiores capturam a informação mais crítica.

2.3 Agregação por Segmento

Para evitar a diluição causada pelo zero-padding, o PreLort introduz uma regra de agregação por segmento que opera no espaço do gradiente (atualização).

• Processo: Em vez de tirar a média de todos os clientes para cada dimensão, o servidor agrega os gradientes para cada segmento de rank $[r_{l-1}, r_l)$ apenas sobre o subconjunto de clientes cujo rank $r_k$ é suficiente para contribuir para esse segmento (ou seja, $r_k \ge r_l$).
• Fórmula: Para o $l$-ésimo segmento, o gradiente agregado é:
  $$\nabla A_{global}^{r_{l-1}:r_l} = \frac{\sum_{k: r_k \ge r_l} n_k \nabla A_k^{r_{l-1}:r_l}}{\sum_{k: r_k \ge r_l} n_k}$$
• Benefício: Isso evita a diluição de sinal de clientes de rank inferior com zero-padding e garante que as dimensões de alto rank sejam atualizadas apenas por clientes que realmente possuem a capacidade de aprendê-las, enquanto as dimensões inferiores se beneficiam de toda a população.

3. Principais Contribuições

O artigo delineia três contribuições primárias:

1. Identificação do Desalinhamento de Representação: Os autores apontam que a falta de um alinhamento estruturado entre as dimensões de rank é uma limitação chave nos métodos atuais de LoRA federado heterogêneo, levando a representações inconsistentes.
2. Formulação Aninhada e Estratégia Unificada: Eles propõem uma formulação de baixo rank aninhada que impõe uma estrutura de prefixo comum. Isso é acoplado a uma estratégia de treinamento e agregação unificada que permite a agregação consistente entre ranks heterogêneos sem o uso de zero-padding.
3. Validação Empírica: O método é demonstrado como capaz de melhorar consistentemente a qualidade e a estabilidade do modelo em relação aos métodos anteriores de LoRA federado heterogêneo em benchmarks de ajuste fino de instrução.

4. Resultados Experimentais

Os autores avaliaram o PreLort em dois modelos base (TinyLlama-1.1B e Qwen2.5-0.5B) através de três conjuntos de dados: Alpaca, Databricks-dolly-15 e 20 Newsgroups.

• Métricas de Desempenho: O PreLort superou consistentemente os baselines (ZeroPad, HetLoRA, FLoRA) em Acurácia (MMLU, 20 Newsgroups) e pontuações ROUGE-L.
  • No TinyLlama-1.1B, o PreLort alcançou 33,86% de acurácia no MMLU (Dolly) e 30,76% no Alpaca, superando os melhores baselines.
  • No Qwen2.5-0.5B, alcançou 39,85% de acurácia no MMLU (Dolly).
• Perplexidade: O PreLort alcançou a menor perplexidade no Dolly e NewsGroup para ambos os modelos base, demonstrando uma verossimilhança superior ao nível de token.
• Robustez de Rank: A análise mostrou que o PreLort mantém um desempenho estável através de diferentes níveis de truncamento de prefixo. Em contraste, baselines como "Nested Aggregation Only" exibiram instabilidade, particularmente em ranks mais altos, confirmando que a agregação sozinha é insuficiente sem o objetivo de treinamento aninhado.
• Importância por Dimensão: A visualização da importância da dimensão revelou que o PreLort concentra o sinal relevante da tarefa nas dimensões iniciais (uma distribuição estruturada), enquanto o treinamento padrão produz uma distribuição uniforme. Essa priorização hierárquica facilita atualizações estáveis entre ranks heterogêneos.

5. Significância e Alegações

O artigo afirma que o PreLort representa uma mudança no enfrentamento da heterogeneidade de rank: em vez de tratá-la como um problema de agregação post-hoc, ela é resolvida estruturando o próprio processo de treinamento.

• Argumento Central: Os autores argumentam que abordar a heterogeneidade de rank requer não apenas uma agregação melhorada, mas também o alinhamento estruturado de representações entre os ranks.
• Eficiência: O método permite que clientes de baixo rank se beneficiem de informações mais ricas contribuídas por clientes de alto rank, pois as dimensões de prefixo são consistentemente aprendidas e agregadas.
• Modéstia sobre Limitações: Os autores reconhecem que a estratégia de treinamento aninhado introduz aproximadamente 40% de overhead computacional adicional devido às múltiplas passagens de forward/backward por lote. Eles também observam que a eficácia depende de passos de otimização local suficientes; com épocas locais muito limitadas, o modelo pode não aprender plenamente as representações de prefixo de baixo rank bem alinhadas.

Em conclusão, o PreLort demonstra que a heterogeneidade de rank no LoRA federado é melhor tratada ao estruturar conjuntamente como os adaptadores são treinados (via otimização de prefixo aninhada) e como suas atualizações são agregadas (via regras de segmentação), resultando em representações de prefixo alinhadas e robustas que superam os métodos de agregação heterogênea existentes.