Beyond Perplexity: A Geometric and Spectral Study of Low-Rank Pre-Training

Resumo Técnico: Além da Perplexidade: Um Estudo Geométrico e Espectral de Pré-treinamento de Baixo Rango

Declaração do Problema
O pré-treinamento de Modelos de Linguagem de Grande Escala (LLMs) é limitado pelos custos de memória associados ao armazenamento de pesos, gradientes e estados do otimizador de posto completo. O pré-treinamento de baixo rango emergiu como uma solução para reduzir esses custos, com métodos evoluindo para três categorias: baseados em otimizador (por exemplo, GaLore, Fira), baseados em arquitetura (por exemplo, CoLA, SLTrain) e baseados em adaptadores (por exemplo, ReLoRA). No entanto, a avaliação desses métodos permanece limitada. As práticas atuais dependem quase exclusivamente da perplexidade de validação de execuções com única semente. Os autores argumentam que a perplexidade é um mau proxy para a qualidade da solução; modelos com perplexidade idêntica podem convergir para regiões fundamentalmente diferentes da paisagem de perda, com representações internas distintas. Além disso, comparações existentes frequentemente falham em distinguir quão bem um método ajusta finamente um modelo pré-treinado versus quão bem um modelo pré-treinado com esse método se generaliza a jusante.

Metodologia
Os autores introduzem um framework de diagnóstico para caracterizar as soluções encontradas por cinco métodos de pré-treinamento de baixo rango (GaLore, Fira, CoLA, SLTrain, ReLoRA) em comparação com o treinamento de posto completo em três escalas de modelo (60M, 130M, 350M). O estudo avalia esses métodos usando 16 métricas em quatro dimensões:

1. Paisagem de Perda 1-D (Direções Aleatórias): Os autores perturbam os parâmetros do modelo ao longo de 100 direções aleatórias para medir a média e a variância da perda de validação. Eles derivam métricas escalares para Nitidez (a taxa de crescimento da perda afastando-se do checkpoint) e Variância Direcional (anisotropia da superfície de perda).
2. Paisagem de Perda 1-D (Direções Top-K PCA): Perturbações são aplicadas ao longo das direções singulares principais das matrizes de pesos para analisar a geometria ao longo dos eixos mais significativos.
3. Interpolação 1-D: Os autores medem a perda ao longo de caminhos lineares entre checkpoints. Eles introduzem duas métricas específicas:
   • Altura da Barreira entre Checkpoints Consecutivos (CCBH): Mede a barreira entre checkpoints sequenciais do mesmo método.
   • Altura da Barreira Inter-Método (IMBH): Uma métrica nova que quantifica a separação geométrica entre diferentes métodos de treinamento no mesmo passo de treinamento. Isso revela se diferentes métodos convergem para a mesma bacia ou para soluções geometricamente distintas.
4. Métricas Espectrais e de Rango: O estudo rastreia a estrutura espectral de pesos e atualizações, incluindo Rango Efetivo, Rango Estável, Gap Espectral e Rango de Limiar.
5. Similaridade de Ativação: Os autores comparam estados ocultos de modelos de baixo rango com baselines de posto completo usando Distância L2 de Ativação, Semelhança Cosseno e CKA Linear (Alinhamento de Kernel Centralizado).

Finalmente, os autores testam se essas características geométricas e espectrais podem prever o desempenho zero-shot a jusante melhor do que apenas a perplexidade, usando regressão linear com validação cruzada leave-one-size-out (LOSO) e leave-one-method-out (LOMO).

Principais Contribuições

1. Framework de Diagnóstico: O artigo propõe uma suíte de avaliação abrangente incluindo a métrica nova IMBH, que quantifica a separação geométrica entre algoritmos de treinamento em vez de apenas entre sementes ou épocas.
2. Descobertas Empíricas Trans-Família: O estudo estabelece que métodos de baixo rango não são equivalentes ao treinamento de posto completo nem entre si. Descobertas-chave incluem:
   • Separação Geométrica: Métodos convergem para bacias geometricamente distintas, com IMBH crescendo monotonicamente conforme o treinamento progride.
   • Assinaturas Espectrais: GaLore corresponde estreitamente à estrutura espectral das atualizações de posto completo, enquanto CoLA não.
   • Divergência de Representação: As ativações de baixo rango divergem crescentemente das de posto completo nas camadas posteriores. CoLA exibe ativações ortogonais, enquanto Fira recupera o alinhamento apenas na camada final.
3. Melhor Predição a Jusante: A perplexidade de validação sozinha é um preditor incompleto do desempenho a jusante. Aumentar a perplexidade com características geométricas e espectrais melhora significativamente a precisão da predição (a correlação de Pearson LOSO aumenta de 0,873 para 0,913).
4. Base de Código Unificada: Os autores lançam um pipeline de código aberto integrando pré-treinamento, computação de métricas e avaliação a jusante para facilitar pesquisas futuras.

Resultados

• Geometria da Paisagem de Perda: O treinamento de posto completo estabelece-se em bacias mais nítidas ao longo de direções aleatórias, mas em bacias mais planas ao longo da direção PCA top-1. Por outro lado, métodos de baixo rango como Fira e SLTrain convergem para bacias mais planas ao longo de direções aleatórias, enquanto CoLA e ReLoRA atingem as bacias mais nítidas. GaLore exibe um perfil misto: moderadamente plano ao longo de direções aleatórias, mas nítido ao longo de sua direção PCA top-1.
• Barreiras Inter-Método: Todos os métodos ocupam bacias geometricamente separadas. Na escala de 60M, a barreira entre posto completo e qualquer método de baixo rango é maior do que as barreiras entre métodos de baixo rango. Fira converge para a bacia geometricamente mais próxima do posto completo, enquanto SLTrain permanece o mais isolado.
• Análise Espectral: O posto completo domina todas as métricas de rango. GaLore rastreia o posto completo estreitamente em rango efetivo e estável, mas exibe um gap espectral menor. Fira e GaLore reproduzem mais fielmente a distribuição espectral das atualizações de pesos de posto completo.
• Alinhamento de Ativação: ReLoRA rastreia as representações de posto completo mais estreitamente em todas as escalas. CoLA mostra o maior desvio, com ativações tornando-se ortogonais à baseline nas camadas posteriores. Fira mostra um desvio médio por camada alto, mas surpreendentemente baixo desvio na camada final.
• Desempenho a Jusante: A perplexidade de validação não prevê consistentemente a classificação a jusante. Por exemplo, Fira alcança a menor perplexidade em todas as escalas, ainda assim CoLA lidera em múltiplas categorias a jusante em 130M e 350M, apesar de não estar classificado entre os dois primeiros em perplexidade.

Significado
O artigo conclui que os métodos de pré-treinamento de baixo rango não são aproximações intercambiáveis do treinamento de posto completo; eles alteram fundamentalmente as soluções alcançadas. O estudo demonstra que a perplexidade de validação fornece uma imagem incompleta da qualidade do modelo. Ao incorporar métricas geométricas e espectrais, os pesquisadores podem prever melhor o desempenho a jusante e entender os vieses indutivos distintos introduzidos por diferentes estratégias de baixo rango. O trabalho destaca que, enquanto alguns métodos (como Fira) convergem para bacias geometricamente próximas ao posto completo, outros (como CoLA) ocupam vales distintos, sugerindo que a escolha do método de baixo rango dita a região específica da paisagem de perda que o modelo habita.