Resumo Técnico: MoSE – Auto-destilação Hierárquica Melhora Embeddings de Camadas Iniciais

1. Declaração do Problema

A implantação de Grandes Modelos de Linguagem (LLMs) em produção frequentemente exige navegar em um equilíbrio entre precisão e desempenho para atender a rigorosas restrições de latência e recursos. A destilação de modelos tradicional reduz o tamanho do modelo, mas incorre em custos significativos ao exigir o treinamento de modelos "estudantes" separados para emular modelos "professores". Além disso, abordagens existentes frequentemente assumem que a camada final de um transformer produz as representações mais úteis, enquanto insights recentes sugerem que camadas intermediárias podem conter a informação semanticamente mais rica para tarefas específicas de downstream. Há uma necessidade de uma arquitetura unificada que permita a implantação flexível (selecionando diferentes partes do modelo) sem sacrificar o desempenho ou exigir pipelines de destilação separados.

2. Metodologia

Os autores introduzem o MODULARSTARENCODER (MOSE), um codificador de 1 bilhão de parâmetros com múltiplas saídas construído sobre a arquitetura StarCoder-2. A inovação central é um mecanismo de Auto-destilação que melhora as representações de camadas inferiores através do treinamento de múltiplos pontos de saída simultaneamente dentro de uma única pilha de modelo.

Arquitetura e Estratégia de Treinamento

• Design de Múltiplas Saídas (Multi-Exit): O modelo apresenta 36 camadas ocultas com cabeças de saída nas camadas 4, 9, 18, 27 e 36. Durante a inferência, os usuários podem selecionar um ponto de saída com base em restrições computacionais (ex: interrompendo na camada 4 para velocidade ou na camada 36 para máxima precisão).
• Perda de Auto-destilação (Self-Distillation Loss): O objetivo de treinamento agrega perdas de todas as camadas de saída selecionadas. Camadas superiores guiam as anteriores, incentivando as camadas inferiores a aprender representações robustas. A perda total é computada como uma soma ponderada: $L = \sum_{i \in \iota} L_i \cdot \alpha_i$, onde $\alpha_i$ aumenta com a profundidade da camada para priorizar camadas mais profundas.
• Objetivos de Pré-treinamento:
  • Modelagem de Linguagem Mascarada (MLM): Predição padrão de tokens.
  • Classificação em Contexto (ICC): Substitui a tradicional perda de Predição de Próxima Sentença (NSP). A ICC classifica se fragmentos de código concatenados originam-se do mesmo repositório. Isso aumenta a densidade de entrada (reduzindo o padding), expande a utilização da janela de contexto e promove uma melhor compreensão em nível de bloco (chunk-level), essencial para recuperação semântica.
• Melhorias Contextuais: O modelo utiliza Grouped Query Attention (GQA), Rotary Positional Encoding (RoPE) e FlashAttention V2. O mascaramento causal do StarCoder-2 foi removido para permitir atenção bidirecional, e a atenção de janela deslizante foi substituída por atenção total para evitar limitações no campo receptivo.

Dados e Fine-tuning

• Dataset SYNTHCoNL: Um novo dataset de 1,07 milhão de tripletos de linguagem natural-código-código. Foi construído traduzindo fragmentos de código do dataset CodeSearchNet para múltiplas linguagens (Python, Java, Go, PHP, etc.) usando o Qwen2.5-Coder-7B-Instruct. Isso cria pares de código-para-código cross-language vinculados à mesma descrição em linguagem natural, expandindo os benchmarks de texto-para-código.
• Fine-tuning: O MOSE foi ajustado (fine-tuned) no SYNTHCoNL para recuperação de texto-para-código e código-para-código usando uma perda do tipo CLIP com cabeças de projeção de múltiplas camadas. Também foi ajustado no BigCloneBench para detecção de clones de código.

3. Principais Contribuições

1. Metodologia de Auto-destilação: Uma abordagem inovadora que treina múltiplos modelos dentro de uma única pilha de camadas, eliminando a necessidade de modelos professor/estudante separados e reduzindo a redundância.
2. Arquitetura MOSE: Um codificador pré-treinado e ajustado de 1 bilhão de parâmetros com cinco pontos de saída, permitindo que os usuários selecionem dinamicamente o tamanho do modelo com base em restrições de memória e computação, mantendo um desempenho competitivo.
3. Dataset SYNTHCoNL: Um novo recurso de benchmark contendo 1.071.367 tripletos gerados via tradução de código, facilitando a pesquisa de recuperação de código-para-código cross-language.

4. Resultados

As avaliações foram conduzidas no conjunto de benchmarks CodeXGLUE (CodeSearchNet, Code Translation, POJ104) e BigCloneBench.

• Recuperação de Texto-para-Código (CodeSearchNet): O MOSE alcançou um MRR (Mean Reciprocal Rank) médio de 81,0, superando todos os modelos de código aberto, incluindo CodeT5+ (770M), UniXcoder e ModernBERT-Large. Ele reduziu a lacuna de desempenho com o modelo fechado OpenAI Text-Embedding-3-Large.
• Recuperação de Código-para-Código:
  • No dataset POJ104 (busca semântica intra-lingual), o MOSE alcançou uma precisão média (mAP) de 75,9 em um cenário zero-shot, o mais alto entre os modelos de código aberto.
  • No dataset de Tradução de Código (CT), o MOSE teve um desempenho comparável ao modelo de embedding da OpenAI e superou outros modelos abertos.
• Detecção de Clone de Código (BigCloneBench): O MOSE alcançou um F1 score de 94,2 na Camada 18, igualando o desempenho de camadas mais profundas e competindo com CodeT5+ (95,1) e UniXcoder (95,2).
• Troca Eficiência-Precisão: O artigo demonstra que reduzir o modelo de 36 para 4 camadas resulta em uma redução de 90% nas operações de ponto flutuante (GFLOPs) com uma queda absoluta de apenas 6,4% no MRR para recuperação de texto-para-código. Notavelmente, para tarefas de código-para-código, o melhor desempenho foi frequentemente observado nas camadas iniciais (ex: Camada 4), desafiando a suposição de que camadas mais profundas são sempre superiores.

5. Significância e Alegações

O artigo afirma que o MOSE oferece uma abordagem fundamentada para trocar custo de inferência por precisão. Ao utilizar a Auto-destilação, o modelo melhora as representações intermediárias, tornando as camadas iniciais viáveis para implantação sem o overhead da destilação tradicional.

• Flexibilidade de Implantação: A arquitetura fornece um mecanismo natural de seleção de modelo, onde os usuários podem escolher o ponto de saída que melhor se adapta às suas restrições de latência e energia.
• Divergência de Representação: Os resultados destacam que a camada ideal para alinhamento semântico depende da modalidade de entrada (texto-para-código vs. código-para-código), com camadas rasas frequentemente performando melhor para tarefas de código-para-código.
• Eficiência: A abordagem suporta uma implantação sustentável ao permitir configurações de modelos menores e energeticamente eficientes sem degradação significativa de desempenho.

Os autores reconhecem limitações, incluindo restrições computacionais que limitaram o ajuste extensivo de hiperparâmetros e o impacto incerto do uso de código sintético gerado via tradução para o fine-tuning. Trabalhos futuros são propostos para explorar ainda mais as interações entre tipos de tarefas e profundidade de camada.