Riepilogo Tecnico: PowerSoftmax: Verso l'Inferenza Sicura di LLM su Dati Criptati

Enunciato del Problema

I moderni Modelli Linguistici su Grande Scala (LLM) che preservano la privacy spesso fanno affidamento sulla Crittografia Omomorfica (HE), in particolare sullo schema CKKS, che consente calcoli su dati criptati senza decrittazione. Tuttavia, CKKS supporta solo operazioni polinomiali. Le architetture Transformer standard dipendono fortemente da componenti non polinomiali, in particolare la funzione Softmax nell'attenzione self-attention, che coinvolge esponenziali, divisioni e operazioni logaritmiche (ad esempio, per la stabilità numerica tramite il trucco log-sum-exp).

Gli approcci esistenti per colmare questo divario presentano limitazioni significative:

1. Approssimazione Post-Addestramento (PTA): Approssimazione diretta delle funzioni non polinomiali nei modelli pre-addestrati utilizzando polinomi di alto grado. Ciò spesso comporta un'inferenza inefficiente a causa della profondità moltiplicativa elevata e del rumore accumulato.
2. Addestramento Consapevole dell'Approssimazione (AAT): Sostituzione delle primitive non polinomiali con alternative più facili da approssimare (ad esempio, attenzione punto per punto o ReLU scalata) prima dell'addestramento. Sebbene efficaci per modelli più piccoli, questi metodi hanno faticato con la stabilità e la scalabilità, non riuscendo a supportare LLM con miliardi di parametri o architetture profonde (ad esempio, 32 livelli).

La sfida fondamentale consiste nel progettare un'architettura Transformer intrinsecamente compatibile con i vincoli HE (polinomi di basso grado, intervalli numerici stabili) mantenendo al contempo il potere espressivo, le capacità di ragionamento e l'apprendimento in contesto (ICL) dei Transformer standard basati su Softmax alla scala dei miliardi di parametri.

Metodologia

Gli autori propongono PowerSoftmax, una variante innovativa e compatibile con HE dell'attenzione self-attention, progettata per replicare i principi fondamentali di Softmax rimanendo suscettibile all'approssimazione polinomiale. La metodologia prevede tre fasi distinte:

1. Progettazione dell'Attenzione Compatibile con HE

Gli autori distillano le proprietà essenziali di Softmax: normalizzazione in $[0, 1]$, amplificazione delle differenze di punteggio e conservazione dell'ordine. Sostituiscono la funzione esponenziale $\exp(x)$ in Softmax con una funzione potenza $x^p$ (dove $p$ è un intero positivo pari).
Il meccanismo di attenzione proposto è definito come:
$$\text{PowerSoftmax}(x)_j = \frac{x_j^p}{\sum_i x_i^p}$$
Questa formulazione evita le esponenziali e i problemi di divisione per zero intrinseci nelle approssimazioni polinomiali dirette di Softmax.

2. Stabilità dell'Addestramento e Robustezza Numerica

Per affrontare l'instabilità dell'addestramento causata dalla funzione potenza (dove gli input $>1$ esplodono e gli input $<1$ svaniscono), gli autori introducono due modifiche chiave:

• Variante Stabile: Ispirati al trucco log-sum-exp, normalizzano gli input mediante una costante $c = \|x\|_\infty + \delta$ prima di applicare la funzione potenza. Ciò garantisce che gli input rimangano entro un intervallo gestibile $(0, 1)$, prevenendo overflow/underflow e preservando la precisione in virgola mobile.
• Divisione Limitata da $\epsilon$: Per gestire l'operazione di divisione, che è non polinomiale e sensibile vicino allo zero, modificano il denominatore in $\epsilon + \sum x_i^p$. Ciò crea una funzione continua $1/\epsilon^2$-Lipschitz, rendendo l'approssimazione polinomiale successiva della divisione significativamente più trattabile e riducendo la profondità moltiplicativa richiesta.

3. Approssimazione Indipendente dalla Lunghezza

L'approssimazione standard della divisione (ad esempio, algoritmo di Goldschmidt) fatica con la somma illimitata dei punteggi di attenzione, che scala linearmente con la lunghezza della sequenza $L$. Gli autori propongono un PowerSoftmax Indipendente dalla Lunghezza:
$$\text{Indipendente dalla Lunghezza}(x)_j = \frac{x_j^p}{L \cdot \text{Media}(x^p)}$$
Poiché $L$ (lunghezza della sequenza) è pubblica, il termine $1/L$ può essere calcolato direttamente senza approssimazione. Il modello approssima la divisione della media anziché della somma, garantendo che il dominio di approssimazione rimanga limitato indipendentemente dalla lunghezza della sequenza.

4. Pipeline di Costruzione (Algoritmo 1)

La pipeline completa per creare un LLM polinomiale prevede:

1. Modifica Architettonica: Sostituzione di Softmax con la variante PowerSoftmax Stabile e addestramento da zero (o tramite pre-addestramento continuo) con obiettivi di minimizzazione dell'intervallo.
2. Minimizzazione dell'Intervallo: Applicazione di una funzione di perdita specializzata per vincolare gli intervalli di input dei livelli non polinomiali (PowerSoftmax, GELU, LayerNorm) per garantire che le approssimazioni polinomiali di basso grado siano accurate.
3. Sostituzione Polinomiale: Dopo l'addestramento, sostituzione delle rimanenti operazioni non polinomiali (divisione, radice quadrata inversa) con approssimazioni polinomiali di basso grado (ad esempio, metodo di Goldschmidt) per ottenere un modello completamente polinomiale.

Contributi Chiave

1. Meccanismo PowerSoftmax: Una nuova variante di attenzione self-attention che sostituisce la scalatura esponenziale con una scalatura polinomiale, specificamente adattata per ambienti HE mantenendo i principi fondamentali di progettazione dell'attenzione.
2. LLM Polinomiali Scalabili: Lo sviluppo dei primi LLM polinomiali che superano il miliardo di parametri (specificamente 1,4 miliardi di parametri con 32 livelli). Questi modelli dimostrano capacità di ragionamento e apprendimento in contesto (ICL) comparabili a quelle dei transformer standard delle stesse dimensioni.
3. Strategie di Addestramento e Inferenza Stabili: Introduzione di varianti di addestramento numericamente stabili e strategie di inferenza indipendenti dalla lunghezza che superano i problemi di instabilità e scalabilità dei precedenti approcci Transformer polinomiali.
4. Valutazione Completa: Un'analisi dettagliata della latenza per i calcoli criptati e studi di ablazione che validano le scelte progettuali, incluso l'impatto del parametro potenza $p$ e della divisione limitata da $\epsilon$.

Risultati Sperimentali

Gli autori hanno valutato il loro metodo su architetture causali (GPT/Pythia) e bidirezionali (RoBERTa):

• Prestazioni: Il modello polinomiale da 1,4 miliardi di parametri ha raggiunto prestazioni comparabili ai modelli Pythia originali basati su Softmax su benchmark zero-shot e few-shot (ad esempio, Lambada, PIQA, WinoGrande, ARC). In particolare, il modello ha dimostrato prestazioni competitive su benchmark di ragionamento come ARC-Challenge, una capacità precedentemente non osservata negli LLM polinomiali.
• Modelli Bidirezionali: Una variante polinomiale di RoBERTa-Base ha mostrato solo un degrado di circa l'1% rispetto al RoBERTa originale basato su Softmax sui compiti GLUE (SST-2, QNLI, MNLI), superando approcci polinomiali precedenti come Nexus.
• Inferenza HE: Esperimenti sul modello Pythia da 70M utilizzando il framework HElayers hanno confermato che l'inferenza criptata produce risultati quasi identici all'inferenza in testo chiaro (perdita MSE massima di 0,005).
• Efficienza: Il metodo proposto riduce significativamente la profondità moltiplicativa rispetto alle alternative all'avanguardia (Nexus, Cho et al.). Per una sequenza di 128 token, PowerSoftmax ha ridotto la latenza di un fattore di circa 9,7 rispetto agli approcci Taylor/Goldschmidt e di circa 20 rispetto agli approcci basati su strategie quadratiche.

Significato e Affermazioni

Il documento afferma una svolta nel campo dell'apprendimento automatico che preserva la privacy, dimostrando che gli LLM completamente polinomiali possono scalare alla regime dei miliardi di parametri senza sacrificare le capacità di ragionamento o di apprendimento in contesto.

• Scalabilità: A differenza di lavori precedenti limitati a modelli piccoli o compiti di classificazione semplici, questo approccio addestra con successo transformer profondi (32 livelli) con 1,4 miliardi di parametri.
• Sicurezza Non Interattiva: Utilizzando la crittografia omomorfica completamente (FHE) con un modello completamente polinomiale, la soluzione abilita un'inferenza non interattiva, evitando il sovraccarico di comunicazione e le potenziali vulnerabilità associate alla computazione multi-partita (MPC) o ai protocolli interattivi.
• Praticità: L'analisi dettagliata della latenza e il dispiegamento riuscito di un modello da 1,4 miliardi di parametri suggeriscono che l'inferenza sicura e privata di LLM sta diventando praticamente fattibile, aprendo la strada ad applicazioni in settori sensibili ai dati come sanità e finanza.

Gli autori riconoscono che, sebbene i modelli funzionino bene, una valutazione completa delle capacità generative auto-regressive nel decoding sequenziale su dati criptati rimane una direzione per lavori futuri.