Resumo Técnico: Um Estudo sobre Destilação de Camadas Ocultas para Pré-treinamento de Modelos de Linguagem de Grande Escala

Declaração do Problema

Embora a Destilação de Conhecimento (KD) seja um paradigma estabelecido para comprimir Modelos de Linguagem de Grande Escala (LLMs) transferindo conhecimento de um grande "professor" para um menor "aluno", a maioria das pesquisas atuais baseia-se exclusivamente em logits de saída. Essa abordagem negligencia a informação semântica contida nas representações intermediárias ocultas do professor. Embora a Destilação de Camadas Ocultas (HLD) tenha mostrado promessa em arquiteturas de codificador e modelos de menor escala, sua aplicação ao pré-treinamento exclusivo de decodificadores em escala permanece amplamente inexplorada. Além disso, comparações existentes frequentemente falham em contabilizar o custo computacional introduzido pela HLD (por exemplo, armazenamento de ativações completas, custos de regressores) ou não controlam rigorosamente o ajuste de hiperparâmetros, tornando difícil determinar se a HLD oferece vantagens genuínas sobre a KD baseada em logits bem ajustada.

Metodologia

Os autores realizam uma avaliação controlada por computação da HLD aplicada ao pré-treinamento de LLMs causais usando a família de modelos Gemma.

• Modelos:
  • Professor: Um modelo Gemma de 3,4B (Transformador de decodificador único com 34 camadas, dimensão de incorporação de 2560) pré-treinado em 106B tokens únicos do conjunto de dados C4.
  • Alunos: Dois modelos baseados em Gemma com inicialização aleatória padrão: um modelo de 123M parâmetros (18 camadas) e um modelo de 735M parâmetros (26 camadas).
• Métodos de Destilação Avaliados:
  1. KD Padrão (Baseada em Logits): O aluno minimiza uma função de perda composta de Verossimilhança Negativa (NLL) e divergência de Kullback-Leibler (KL) contra os logits de saída do professor.
  2. Otimização Sequencial (HLDF): Adaptado do FitNets, envolve um protocolo de duas fases. A Fase 1 alinha o estado intermediário oculto do aluno (na camada mediana) com o do professor usando um regressor aprendido (perceptron denso de uma camada) via uma perda de Erro Quadrático Médio (MSE) normalizada. A Fase 2 realiza a KD padrão.
  3. Otimização Conjunta (HLDC): Uma abordagem de treinamento de estágio único que minimiza uma função objetivo agregada combinando NLL, destilação de logits e uma perda de alinhamento de camada oculta (perda de incorporação) simultaneamente.
• Rigor Experimental:
  • Correspondência de Computação: As comparações são estritamente baseadas em orçamentos computacionais equivalentes (FLOPs), não apenas em contagens de tokens. Os autores refinam a aproximação padrão de computação $6ND$ para incluir o custo significativo das camadas de desincorporação e do mapeamento do regressor, que é não desprezível para modelos com menos de um bilhão de parâmetros e vocabulários grandes.
  • Controle de Hiperparâmetros: O estudo evita "ajustar para vencer" avaliando os métodos em uma grade compartilhada de hiperparâmetros (taxa de aprendizado, temperatura, peso da NLL) e comparando resultados ponto a ponto. Também identifica a melhor configuração para cada método para garantir uma comparação justa "teto a teto".
  • Dados: O treinamento ocorre no subconjunto em inglês do conjunto de dados C4 (até 168B tokens para o aluno de 735M).

Contribuições Principais

1. Avaliação Sistemática: O artigo fornece o primeiro estudo de código aberto de HLD especificamente durante a fase de pré-treinamento de LLMs exclusivos de decodificadores, escalando para modelos com centenas de milhões de parâmetros.
2. Linhas de Base Rigorosas: Estabelece um quadro de referência que avalia métodos com base em equivalência computacional estrita (FLOPs), contabilizando explicitamente os custos não desprezíveis de desincorporação e cálculo de perda em modelos pequenos, que estudos anteriores frequentemente ignoravam.
3. Benchmarking Abrangente: A avaliação abrange múltiplas tarefas downstream (Wikitext 103, HellaSwag, WinoGrande, LAMBADA, PIQA, ARC-E) e analisa o desempenho em uma ampla grade de hiperparâmetros, em vez de depender de uma única configuração "melhor".

Resultados

Os experimentos produzem achados matizados sobre a eficácia da HLD em comparação com a KD padrão:

• Desempenho Downstream: A HLD não supera consistentemente a KD baseada em logits padrão em tarefas de avaliação downstream. Na comparação de "melhor hiperparâmetro", nenhuma variante de HLD domina a KD em todos os benchmarks; as diferenças são marginais e divididas entre os métodos.
• Ganhos de Perplexidade: Apesar da falta de superioridade downstream, a HLDF (Otimização Sequencial) produz uma melhoria sistemática, embora pequena, na perplexidade do C4 em todas as configurações de hiperparâmetros compartilhados. A melhoria média é de aproximadamente 0,26% para ambos os tamanhos de aluno.
• Contexto dos Ganhos: Os ganhos de perplexidade da HLDF são mais pronunciados em regimes onde a KD padrão performa mal. Em regimes onde a KD já é eficiente, a vantagem da HLD desaparece.
• Downstream vs. Perplexidade: O ganho sistemático de perplexidade não se traduz em melhorias consistentes nas pontuações downstream. Os autores sugerem que isso indica que, embora exista um sinal latente nas camadas ocultas, ele pode não alterar fundamentalmente as capacidades de granulação grossa capturadas pelos benchmarks downstream, ou que o sinal é útil apenas para ajuste distribucional de granulação fina.

Significado e Alegações

O artigo conclui que, embora a HLD seja um conceito atraente para capturar conhecimento interno profundo, ela não fornece uma melhoria consistente e confiável sobre uma linha de base de KD baseada em logits bem ajustada quando os orçamentos de treinamento são estritamente equalizados.

• Resiliência da KD de Logits: Os achados sublinham a resiliência da destilação baseada em logits padrão, sugerindo que, para o pré-treinamento exclusivo de decodificadores, a complexidade adicional e a sensibilidade a hiperparâmetros da HLD podem superar seus benefícios na implantação prática.
• Discrepâncias Metodológicas: Os autores atribuem a contradição com trabalhos anteriores (que frequentemente relatam ganhos de HLD) a quatro fatores:
  1. Contabilidade de Computação: Estudos anteriores frequentemente correspondiam contagens de tokens em vez de FLOPs, ignorando o custo extra da HLD.
  2. Arquitetura/Fase: Trabalhos anteriores focaram em codificadores ou pós-treinamento, onde as representações já são moldadas por objetivos de pré-treinamento. Decodificadores causais treinados do zero apresentam um cenário de otimização diferente.
  3. Ajuste da Linha de Base: Estudos anteriores frequentemente comparavam HLD contra linhas de base de KD mal ajustadas.
  4. Escalabilidade: Adaptar a HLD de CNNs pequenas para transformadores de bilhões de parâmetros requer escolhas de design mais profundas (perdas, adaptadores) do que simples regressores lineares.
• Perspectiva Futura: O artigo postula que o ganho sistemático de perplexidade observado na HLDF sugere que um "sinal latente" pode ser extraído, mas realizar um avanço provavelmente exigiria inovações significativas de design ou evidências estatísticas mais fortes em escalas maiores e diferentes famílias de modelos. Os autores enfatizam a necessidade de contabilidade exata de FLOPs ao avaliar técnicas de supervisão inovadoras.