Sintesi Tecnica: MoSE – La Auto-Distillazione Gerarchica Migliora gli Embedding degli Strati Iniziali

1. Definizione del Problema

Il deployment di Large Language Models (LLM) in produzione richiede spesso di gestire un compromesso tra accuratezza e prestazioni per soddisfare rigidi vincoli di latenza e risorse. La distillazione tradizionale dei modelli riduce le dimensioni del modello ma comporta costi significativi, richiedendo l'addestramento di modelli "studenti" separati per emulare i modelli "teacher". Inoltre, gli approcci esistenti spesso assumono che l'ultimo strato di un transformer produca le rappresentazioni più utili, mentre intuizioni recenti suggeriscono che gli strati intermedi possano contenere le informazioni semanticamente più ricche per specifici task a valle. È necessario un'architettura unificata che permetta un deployment flessibile (selezionando diverse porzioni del modello) senza sacrificare le prestazioni o richiedere pipeline di distillazione separate.

2. Metodologia

Gli autori introducono MODULARSTARENCODER (MOSE), un encoder multi-exit da 1 miliardo di parametri basato sull'architettura StarCoder-2. L'innovità principale è un meccanismo di Auto-Distillazione (Self-Distillation) che migliora le rappresentazioni degli strati inferiori addestrando simultaneamente molteplici punti di uscita all'interno di un singolo stack di modelli.

Architettura e Strategia di Addestramento

• Design Multi-Exit: Il modello presenta 36 strati nascosti con testine di uscita (exit heads) agli strati 4, 9, 18, 27 e 36. Durante l'inferenza, gli utenti possono selezionare un punto di uscita in base ai vincoli computazionali (ad esempio, fermandosi allo strato 4 per la velocità o allo strato 36 per la massima accuratezza).
• Loss di Auto-Distillazione: L'obiettivo di addestramento aggrega le perdite (loss) da tutti gli strati di uscita selezionati. Gli strati superiori guidano quelli precedenti, incoraggiando gli strati inferiori a imparare rappresentazioni robuste. La loss totale è calcolata come una somma pesata: $L = \sum_{i \in \iota} L_i \cdot \alpha_i$, dove $\alpha_i$ aumenta con la profondità dello strato per dare priorità agli strati più profondi.
• Obiettivi di Pre-training:
  • Masked Language Modeling (MLM): Predizione standard del token.
  • In-Context Classification (ICC): Sostituisce la tradizionale perdita Next Sentence Prediction (NSP). L'ICC classifica se frammenti di codice concatenati originano dallo stesso repository. Ciò aumenta la densità dell'input (riducendo il padding), espande l'utilizzo della finestra di contesto e favorisce una migliore comprensione a livello di chunk, essenziale per il recupero semantico.
• Miglioramenti Contestuali: Il modello utilizza Grouped Query Attention (GQA), Rotary Positional Encoding (RoPE) e FlashAttention V2. La maschera causale di StarCoder-2 è stata rimossa per abilitare l'attenzione bidirezionale, e l'attenzione a finestra scorrevole (sliding window attention) è stata sostituita con l'attenzione completa per evitare limitazioni nel campo recettivo.

Dati e Fine-tuning

• Dataset SYNTHCoNL: Un nuovo dataset di 1,07 milioni di triplette linguaggio naturale-codice-codice. È stato costruito traducendo frammenti di codice dal dataset CodeSearchNet in più linguaggi (Python, Java, Go, PHP, ecc.) utilizzando Qwen2.5-Coder-7B-Instruct. Questo crea coppie cross-language codice-a-codice legate alla stessa descrizione in linguaggio naturale, espandendo i benchmark text-to-code.
• Fine-tuning: MOSE è stato sottoposto a fine-tuning su SYNTHCoNL per il recupero text-to-code e code-to-code utilizzando una loss in stile CLIP con testine di proiezione multi-strato. È stato inoltre sottoposto a fine-tuning su BigCloneBench per il rilevamento dei code clone.

3. Contributi Chiave

1. Metodologia di Auto-Distillazione: Un approccio innovativo che addestra molteplici modelli all'interno di un unico stack di strati, eliminando la necessità di modelli teacher/student separati e riducendo la ridondanza.
2. Architettura MOSE: Un encoder pre-addestrato e sottoposto a fine-tuning da 1 miliardo di parametri con cinque punti di uscita, che consente agli utenti di selezionare dinamicamente la dimensione del modello in base ai vincoli di memoria e computazionali mantenendo prestazioni competitive.
3. Dataset SYNTHCoNL: Una nuova risorsa di benchmark contenente 1.071.367 triplette generate tramite traduzione del codice, che facilita la ricerca sul recupero cross-language codice-a-codice.

4. Risultati

Le valutazioni sono state condotte sulla suite di benchmark CodeXGLUE (CodeSearchNet, Code Translation, POJ104) e BigCloneBench.

• Recupero Text-to-Code (CodeSearchNet): MOSE ha raggiunto un MRR (Mean Recivocal Rank) medio di 81,0, superando tutti i modelli open-source, inclusi CodeT5+ (770M), UniXcoder e ModernBERT-Large. Ha ridotto il divario di prestazioni con il modello chiuso OpenAI Text-Embedding-3-Large.
• Recupero Code-to-Code:
  • Sul dataset POJ104 (ricerca semantica intra-linguaggio), MOSE ha ottenuto un mAP (mean Average Precision) di 75,9 in un setting zero-shot, il più alto tra i modelli open-source.
  • Sul dataset Code Translation (CT), MOSE ha ottenuto prestazioni comparabili al modello di embedding di OpenAI e ha superato altri modelli aperti.
• Rilevamento Code Clone (BigCloneBench): MOSE ha ottenuto un punteggio F1 di 94,2 allo strato 18, eguagliando le prestazioni degli strati più profondi e competendo con CodeT5+ (95,1) e UniXcoder (95,2).
• Compromesso Efficienza-Accuratezza: Il paper dimostra che ridurre il modello da 36 a 4 strati comporta una riduzione del 90% delle operazioni in virgola mobile (GFLOPs) con un calo assoluto di soli 6,4% nell'MRR per il recupero text-to-code. Notevolmente, per i task code-to-code, le migliori prestazioni sono state spesso osservate negli strati iniziali (es. Strato 4), sfidando l'assunto che gli strati più profondi siano sempre superiori.

5. Significato e Rivendicazioni

Il paper sostiene che MOSE offre un approccio fondato per scambiare il costo di inferenza con l'accuratezza. Sfruttando l'Auto-Distillazione, il modello migliora le rappresentazioni intermedie, rendendo gli strati iniziali validi per il deployment senza l'overhead della distillazione tradizionale.

• Flessibilità di Deployment: L'architettura fornisce un meccanismo naturale di selezione del modello dove gli utenti possono scegliere il punto di uscita che meglio si adatta ai loro vincoli di latenza ed energia.
• Divergenza delle Rappresentazioni: I risultati evidenziano che lo strato ottimale per l'allineamento semantico dipende dalla modalità di input (text-to-code vs code-to-code), con gli strati meno profondi che spesso performano meglio per i task code-to-code.
• Efficienza: L'approccio supporta un deployment sostenibile consentendo configurazioni di modelli più piccole ed efficienti dal punto di vista energetico senza una degradazione significativa delle prestazioni.

Gli autori riconoscono i limiti, tra cui i vincoli computazionali che hanno limitato l'estesa ottimizzazione degli iperparametri e l'impatto non chiaro dell'uso di codice sintetico generato tramite traduzione per il fine-tuning. Il lavoro futuro proposto mira ad esplorare ulteriormente le interazioni tra i tipi di task e la profondità dello strato.