Technische Samenvatting: PowerSoftmax: Naar Veilige LLM-Inferentie over Versleutelde Data

Probleemstelling

Moderne privacybehoudende Large Language Models (LLM's) vertrouwen vaak op Homomorf Versleuteling (HE), specifiek het CKKS-schema, dat berekeningen op versleutelde data mogelijk maakt zonder decryptie. CKKS ondersteunt echter uitsluitend polynoomoperaties. Standaard Transformer-architecturen zijn sterk afhankelijk van niet-polynomiale componenten, met name de Softmax-functie in self-attention, die exponentiële functies, deling en logaritmische bewerkingen omvat (bijvoorbeeld voor numerieke stabiliteit via de log-sum-exp-truc).

Bestaande benaderingen om deze kloof te overbruggen ondervinden aanzienlijke beperkingen:

1. Post-Training Benadering (PTA): Directe benadering van niet-polynomiale functies in vooraf getrainde modellen met behulp van polynomen van hoge graad. Dit resulteert vaak in inefficiënte inferentie vanwege een hoge multiplicatieve diepte en opgebouwde ruis.
2. Benaderingsbewuste Training (AAT): Vervanging van niet-polynomiale primitieven door alternatieven die makkelijker te benaderen zijn (bijvoorbeeld puntsgewijze attention of geschaalde ReLU) voordat er getraind wordt. Hoewel deze methoden effectief zijn voor kleinere modellen, hebben ze te kampen met stabiliteits- en schaalbaarheidsproblemen, waardoor ze geen LLM's met miljarden parameters of diepe architecturen (bijvoorbeeld 32 lagen) kunnen ondersteunen.

De kernuitdaging bestaat uit het ontwerpen van een Transformer-architectuur die van nature compatibel is met HE-beperkingen (polynomen van lage graad, stabiele numerieke bereiken), terwijl tegelijkertijd de expressieve kracht, redeneervermogen en contextueel leren (ICL) van standaard Softmax-gebaseerde Transformers op de schaal van miljarden parameters behouden blijven.

Methodologie

De auteurs stellen PowerSoftmax voor, een nieuwe, HE-vriendelijke variant van self-attention die ontworpen is om de kernprincipes van Softmax na te bootsen, terwijl deze vatbaar blijft voor polynoombenadering. De methodologie omvat drie distincte fasen:

1. HE-Vriendelijk Attention-ontwerp

De auteurs distilleren de essentiële eigenschappen van Softmax: normalisatie naar $[0, 1]$, versterking van scoreverschillen en behoud van volgorde. Ze vervangen de exponentiële functie $\exp(x)$ in Softmax door een machtsfunctie $x^p$ (waarbij $p$ een positief even geheel getal is).
Het voorgestelde attention-mechanisme wordt gedefinieerd als:
$$\text{PowerSoftmax}(x)_j = \frac{x_j^p}{\sum_i x_i^p}$$
Deze formulering vermijdt exponentiële functies en deling door nul-problemen die inherent zijn aan directe polynoombenaderingen van Softmax.

2. Trainingsstabiliteit en Numerieke Robuustheid

Om trainingsinstabiliteit veroorzaakt door de machtsfunctie aan te pakken (waarbij invoer $>1$ explodeert en invoer $<1$ verdwijnt), introduceren de auteurs twee belangrijke modificaties:

• Stabiele Variant: Geïnspireerd door de log-sum-exp-truc normaliseren ze invoer met een constante $c = \|x\|_\infty + \delta$ voordat ze de machtsfunctie toepassen. Dit zorgt ervoor dat invoer binnen een beheersbaar bereik $(0, 1)$ blijft, waardoor overloop/onderloop wordt voorkomen en zwevendekommaprecisie behouden blijft.
• $\epsilon$-Gebonden Deling: Om de delingsoperatie te hanteren, die niet-polynomiaal is en gevoelig bij nul, modificeren ze de noemer naar $\epsilon + \sum x_i^p$. Dit creëert een $1/\epsilon^2$-Lipschitz-continue functie, waardoor de daaropvolgende polynoombenadering van de deling aanzienlijk hanteerbaarder wordt en de vereiste multiplicatieve diepte afneemt.

3. Lengte-agnostische Benadering

Standaard delingsbenadering (bijvoorbeeld het Goldschmidt-algoritme) heeft moeite met de onbegrensde som van attention-scores, die lineair schaalt met de sequentielengte $L$. De auteurs stellen een Lengte-agnostische PowerSoftmax voor:
$$\text{Lengte-agnostisch}(x)_j = \frac{x_j^p}{L \cdot \text{Mean}(x^p)}$$
Omdat $L$ (sequentielengte) publiek is, kan de term $1/L$ direct worden berekend zonder benadering. Het model benadert de deling van het gemiddelde in plaats van de som, waardoor het benaderingsdomein onafhankelijk van de sequentielengte begrensd blijft.

4. Constructiepiplijn (Algoritme 1)

De volledige pijplijn voor het creëren van een polynoom-LLM omvat:

1. Architecturale Wijziging: Vervanging van Softmax door de Stabiele PowerSoftmax-variant en training vanaf nul (of via continue voortraining) met doelstellingen voor bereikminimalisatie.
2. Bereikminimalisatie: Toepassing van een gespecialiseerde verliesfunctie om de invoerbereiken van niet-polynomiale lagen (PowerSoftmax, GELU, LayerNorm) te beperken, zodat polynoombenaderingen van lage graad accuraat zijn.
3. Polynoomvervanging: Na training worden de resterende niet-polynomiale operaties (deling, inverse vierkantswortel) vervangen door polynoombenaderingen van lage graad (bijvoorbeeld de Goldschmidt-methode) om een volledig polynoommodel te verkrijgen.

Belangrijkste Bijdragen

1. PowerSoftmax-mechanisme: Een nieuwe variant van self-attention die exponentiële schaling vervangt door polynoomschaling, specifiek op maat gemaakt voor HE-omgevingen terwijl de kernontwerpprincipes van attention behouden blijven.
2. Schaalbare Polynoom-LLM's: De ontwikkeling van de eerste polynoom-LLM's met meer dan één miljard parameters (specifiek 1,4 miljard parameters met 32 lagen). Deze modellen vertonen redeneer- en contextleercapaciteiten (ICL) die vergelijkbaar zijn met standaard Transformers van dezelfde grootte.
3. Stabiele Trainings- en Inferentiestrategieën: Introductie van numeriek stabiele trainingsvarianten en lengte-agnostische inferentiestrategieën die de instabiliteits- en schaalbaarheidsproblemen van eerdere polynoom-Transformer-benaderingen overwinnen.
4. Uitgebreide Evaluatie: Een gedetailleerde latency-benadering voor versleutelde berekeningen en ablatiestudies die de ontwerpkeuzes valideren, inclusief de impact van de machtsparameter $p$ en de $\epsilon$-gebonden deling.

Experimentele Resultaten

De auteurs evalueerden hun methode op causale (GPT/Pythia) en bidirectionele (RoBERTa) architecturen:

• Prestaties: Het polynoommodel met 1,4 miljard parameters behaalde prestaties die vergelijkbaar waren met de originele Softmax-gebaseerde Pythia-modellen op zero-shot en few-shot benchmarks (bijvoorbeeld Lambada, PIQA, WinoGrande, ARC). Opmerkelijk is dat het model concurrerende prestaties liet zien op redeneerbenchmarks zoals ARC-Challenge, een capaciteit die eerder niet was waargenomen bij polynoom-LLM's.
• Bidirectionele Modellen: Een polynoom RoBERTa-Base-variant vertoonde slechts een degradatie van ongeveer 1% ten opzichte van de originele Softmax-gebaseerde RoBERTa op GLUE-taken (SST-2, QNLI, MNLI), en presteerde beter dan eerdere polynoombenaderingen zoals Nexus.
• HE-inferentie: Experimenten met het Pythia 70M-model met behulp van het HElayers-framework bevestigden dat versleutelde inferentie resultaten oplevert die bijna identiek zijn aan plaintext-inferentie (maximale MSE-verlies van 0,005).
• Efficiëntie: De voorgestelde methode vermindert de multiplicatieve diepte aanzienlijk ten opzichte van state-of-the-art alternatieven (Nexus, Cho et al.). Voor een sequentie van 128 tokens verlaagde PowerSoftmax de latency met een factor van ongeveer 9,7 ten opzichte van Taylor/Goldschmidt-benaderingen en met ongeveer 20 ten opzichte van square-strategy-benaderingen.

Betekenis en Claims

Het artikel claimt een doorbraak op het gebied van privacybehoudend machine learning door aan te tonen dat volledig polynoom-LLM's kunnen schalen naar het regime van miljarden parameters zonder in te leveren op redeneer- of contextleercapaciteiten.

• Schaalbaarheid: In tegenstelling tot eerdere werken die beperkt waren tot kleine modellen of eenvoudige classificatietaken, traint deze benadering succesvol diepe (32-laags) Transformers met 1,4 miljard parameters.
• Niet-interactieve Beveiliging: Door gebruik te maken van volledig homomorf versleuteling (FHE) met een volledig polynoommodel, maakt de oplossing niet-interactieve inferentie mogelijk, waarbij de communicatie-overhead en potentiële kwetsbaarheden die geassocieerd zijn met multi-party computation (MPC) of interactieve protocollen worden vermeden.
• Praktijktoepasbaarheid: De gedetailleerde latency-analyse en de succesvolle implementatie van een model met 1,4 miljard parameters suggereren dat veilige, private LLM-inferentie praktisch haalbaar wordt, wat de weg effent voor toepassingen in datagevoelige sectoren zoals gezondheidszorg en financiën.

De auteurs erkennen dat, hoewel de modellen goed presteren, een volledige evaluatie van autoregressieve generatieve vaardigheden in sequentiële decoding over versleutelde data een richting blijft voor toekomstig werk.