Sintesi Tecnica: SecRL-Prune

Problematica

I grandi Modelli Linguistici per il Codice (CodeLLM) possiedono la capacità di generare e riscrivere programmi, una caratteristica che può essere sfruttata per creare diverse varianti di malware capaci di eludere il rilevamento basato su firme. Una questione di sicurezza critica sorge: questi modelli possono essere compressi (tramite pruning, distillazione o quantizzazione) per adattarsi a budget hardware limitati (ad esempio, incorporati in payload di malware o in esecuzione su hardware di consumo) mantenendo al contempo la loro capacità di eseguire mutazioni di codice che preservano la semantica?

La ricerca attuale sulla compressione dei modelli si concentra ampiamente sui compiti del linguaggio naturale, lasciando inesplorato l'impatto della compressione strutturata sulla correttezza dell'esecuzione e sulla diversità della mutazione dei CodeLLM. Inoltre, gli esistenti metodi di pruning strutturato richiedono spesso risorse computazionali significative, come il mantenimento simultaneo del modello teacher e dello studente nella memoria GPU durante l'addestramento, il che ne limita la fattibilità in scenari avversari con risorse limitate.

Metodologia: SecRL-Prune

Gli autori propongono SecRL-Prune, un framework di apprendimento per rinforzo (RL) strutturato progettato per eseguire il pruning dei CodeLLM preservandone la fedeltà nella mutazione del codice. L'approccio opera sui canali delle reti Feed-Forward (FFN/MLP) all'interno dei blocchi transformer, poiché questi strati contengono la maggior parte dei parametri del modello.

Componenti Core

1. Apprendimento della Policy basato su RL:

   • Il problema è formulato come una ricerca della maschera di pruning ottimale (una decisione binaria di mantenere o rimuovere ogni canale MLP) attraverso tutti i livelli.
   • Una rete di policy apprende ad assegnare punteggi di importanza a ciascun canale. Durante l'addestramento, la policy campiona le maschere per creare modelli "studenti".
   • L'obiettivo è identificare quali canali sono essenziali per mantenere il comportamento del modello teacher e quali possono essere rimossi in sicurezza.

2. Segnale di Ricompensa basato sulla Divergenza KL:

   • La funzione di ricompensa si basa sulla Distillazione della Conoscenza (Knowledge Distillation). Il modello studente con pruning viene valutato rispetto al modello "teacher" originale.
   • La ricompensa è definita come la divergenza negativa di Kullback-Leibler (KL) tra la distribuzione di probabilità del prossimo token del teacher e la distribuzione dello studente.
   • Per garantire che lo studente preservi le capacità del teacher, la policy viene aggiornata (usando l'algoritmo REINFORCE) per massimizzare questa ricompensa, minimizzando efficacemente la divergenza.

3. Meccanismo di Caching Top-$P$:

   • Un'innovazione chiave per migliorare l'efficienza è il caching degli output del teacher. Invece di eseguire il modello teacher e lo studente in parallelo (il che richiederebbe il doppio della memoria GPU), il modello teacher viene eseguito una sola volta su un set di calibrazione.
     im Solo le predizioni dei token top-$P$ (logits e indici) e un bucket aggregato "Other" vengono memorizzati in cache per ogni time step.
   • Durante l'addestramento RL, lo studente viene valutato solo rispetto a questo target compatto memorizzato in cache. Ciò riduce significativamente l'uso di picco della memoria GPU, poiché solo lo studente risiede in memoria durante l'ottimizzazione.

Contributi Principali

Il documento delinea quattro contributi primari:

1. Framework SecRL-Prune: Un nuovo framework di pruning strutturato basato su RL per i CodeLLM che utilizza un obiettivo guidato dalla KL per preservare il comportamento del teacher sotto compressione.
2. Addestramento Efficiente della Memoria: L'introduzione di un meccanismo di caching che pre-calcola gli output del teacher, consentendo alla policy di pruning di addestrare con solo lo studente in memoria. Ciò produce un uso di picco della memoria GPU inferiore rispetto alle precedenti tecniche di pruning strutturato.
3. Studio Empirico sulla Mutazione: Una valutazione sistematica di come la compressione strutturata influenzi la capacità dei CodeLLM di eseguire mutazioni di codice che preservano la funzionalità, utilizzando metriche basate sull'esecuzione.
4. Caso di Studio sull'Evasione di Malware: Una valutazione di come le mutazioni che preservano la semantica, generate da modelli compressi, influenzino i tassi di rilevamento attraverso i motori per malware standard del settore (VirusTotal).

Risultati Sperimentali

Gli autori hanno valutato SecRL-Prune su tre CodeLLM da 7B (CodeLlama-7B-Instruct, CodeLlama-7B-Python e Qwen2.5-Coder-7B-Instruct) a rapporti di compressione dal 10% al 30%, confrontandoli con il teacher non compresso e con il baseline allo stato dell'arte PruneNet.

Metriche di Performance

• Correttezza dell'Esecuzione (pass@k): SecRL-Prune ha superato costantemente PruneNet. Al 10% di compressione, ha mantenuto l'80–85% della performance pass@k del teacher, mentre PruneNet ha mantenuto solo il 45–72%. Con una compressione aggressiva del 30%, SecRL-Prune ha mantenuto una correttezza significativamente più alta (ad esempio, 63,7% contro 35,7% su Qwen2.5-Coder).
• Diversità della Mutazione (var@k): Il framework ha preservato meglio la diversità delle soluzioni corrette rispetto al baseline. Notevolmente, al 10% di compressione, SecRL-Prune su Qwen2.5-Coder ha leggermente migliorato la diversità delle soluzioni rispetto al modello denso, mentre PruneNet ha mostrato un netto declino.
• Efficienza delle Risorse: In termini di uso di picco della memoria GPU durante l'addestramento al 20% di compressione, SecRL-Prune ha utilizzato meno della metà della memoria di PruneNet (riducendo l'uso da circa 33–37 GiB a circa 14–16 GiB, una riduzione del ~55–60%).

Caso di Studio sull'Evasione di Malware

Utilizzando campioni di malware reali (Python e JavaScript), gli autori hanno testato se le mutazioni generate da un CodeLLM con il 20% di pruning potessero eludere il rilevamento.

• Risultati: I modelli compressi hanno generato con successo riscritture che preservano la semantica che hanno ridotto significativamente i tassi di rilevamento su VirusTotal.
  • Un campione Python ha visto il rilevamento scendere da 28/63 a 9/63 (riduzione del 68%).
  • Un campione JavaScript ha visto il rilevamento crollare da 9/62 a 1/62 (riduzione dell'88%).
• Implicazione: Anche con una compressione significativa, i modelli hanno mantenuto la capacità di generare varianti che aggirano l'analisi statica e il rilevamento basato su firme.

Significato e Rivendicazioni

Il documento afferma che la capacità di mutazione del codice può sopravvivere a un significativo pruning strutturato. Questa scoperta evidenzia un rischio critico per la sicurezza: la fattibilità di distribuire CodeLLM altamente compressi e specializzati direttamente all'interno di payload di malware o su hardware consumer per automatizzare la generazione di varianti di malware evasive.

Gli autori sottolineano che, sebbene i modelli compressi siano ancora capaci, non sono sostituti perfetti del teacher denso (persiste un divario di prestazioni). Tuttavia, la capacità di comprimere questi modelli mantenendo la loro utilità avversaria abbassa la barriera all'ingresso per la creazione di "motori di mutazione miniaturizzati". Il lavoro funge da avvertimento per i difensori, suggerendo che il malware futuro potrebbe sfruttare modelli compressi e sottoposti a RL-pruning per automatizzare il polimorfismo, rendendo necessari nuovi meccanismi di difesa che vadano oltre le firme statiche.

Il documento conclude notando i limiti, come la dipendenza da un supporto Top-$P$ troncato per la funzione di ricompensa e l'esclusione del pruning delle teste di attenzione, suggerendo che i lavori futuri dovrebbero esplorare dati di calibrazione più ampi e strategie di pruning congiunte.