Updates over HEIR, het compilerproject voor homomorfe encryptie

Voor wie het corporate blogbericht niet leest: HEIR is een compiler die een inputprogramma omzet in een programma dat direct opereert op versleutelde data. De garantie van homomorfe encryptie is dat, uitgaande van het feit dat de cryptografie niet gekraakt is, de computer die het programma uitvoert op geen enkel moment ook maar één bit aan informatie krijgt over de platte tekst (cleartext) die is gebruikt om de versleutelde inputs te genereren. Er is geen informatie beschikbaar over de inputs, outputs of enige tussenwaarden.[^1]

Het blogbericht richt zich op het vermogen van HEIR om vooraf getrainde ML-modellen te compileren en geeft vier voorbeelden van kleine, maar niet-triviale modellen die gecompileerd kunnen worden. Homomorfe encryptie kan diensten dus in staat stellen om volledig private inferentie te bieden. Later in dit artikel zal ik dieper ingaan op wanneer en waar dit nuttig is.

Eerst wil ik een concreter beeld geven van hoe HEIR werkt in de context van deze voorbeelden, en de roadmap voor de toekomst van het project schetsen. Ik zal niet diep ingaan op de interne werking van HEIR om het artikel behapbaar te houden, maar er is veel documentatie te vinden op heir.dev en er is een recente presentatie beschikbaar die ik dit jaar gaf op ASPLOS.

De repository achter de blogpost

Het blogbericht eindigt met een lijst van voorbeelden die met HEIR zijn gecompileerd. Deze voorbeelden verwijzen naar een GitHub-repository[^2] die je kunt klonen en zelf kunt uitvoeren. De grootste hindernis is het installeren van Bazel; daarna beheert Bazel alles hermetisch.[^3]

Enkele runtime-vergelijkingen

Het eenvoudigste en snelste voorbeeld om te proberen is de detector voor creditcardfraude. Dit is een eenvoudig feed-forward netwerk van drie lagen met sigmoid-activaties, getraind op een Kaggle-dataset. De lineaire lagen hebben dimensies 128, 64 en 2 (waarbij de laatste de logits zijn voor de twee klassen: fraude en geen fraude).

Je kunt het basisvoorbeeld met één regel uitvoeren: bazel run -c opt //demos/ccfraud/lattigo:evaluatefhe

Dit commando compileert het (vooraf getrainde) cc_fraud-model naar de Lattigo-backend en voert het uit op een voorbeeldinput. De output is als volgt:

Loading test row 0 from /home/jeremy/fully-homomorphic-encryption/demos/cc_fraud/data/test_rows.csv...
Took 83.226µs
Expected label (is_fraud): 0
Feature vector size: 82
First 5 features: [-0.31676582 0.85089076 -0.40874073 -0.1833772 -1.7155787]
Configuring Lattigo context...
Took 2.164051998s
Encrypting input features...
Took 19.430069ms
Running preprocessing...
Took 573.537941ms
Running FHE evaluation (preprocessed)...
Took 2.020821739s
Decrypting output...
Took 488.237µs
Decrypted logits: [16.464235 -16.781752]
Predicted class: 0
SUCCESS: Predicted class matches expected label!

Het centrale punt hier is dat de evaluatie van het model op versleutelde inputs ongeveer 2 seconden duurde op een single-threaded CPU. Ter vergelijking: dezelfde uitvoering op platte tekst (waarbij dit de latentie van een enkele inferentie is, zonder voordeel van amortisatie):

$ bazel run -c opt //demos/ccfraud/cleartext:evaluatecleartext

Loading model from: demos/cc_fraud/data/mlp_fraud_model_sigmoid.pt
Evaluating Credit Card Fraud Sample Index: 0
True Label:      0 (LEGITIMATE)
Predicted Label: 0 (LEGITIMATE)
Fraud Probability: 0.000000
Result:          CORRECT
Latency:         0.5233 ms

Wie bekend is met HE, heeft waarschijnlijk gehoord dat het traag is. Een runtime van 2 seconden voor HE-inferentie tegenover 0,5 ms voor platte tekst is een vertraging van 4.000x. De berekening omvat twee matrix-vectorproducten (waarbij de matrix niet privaat is) met twee evaluaties van een sigmoid-functie.

Er zijn enkele belangrijke kanttekeningen bij deze demo:

  • Het model is klein genoeg zodat de private informatie in één enkele (CKKS) cijfertekst past. Input-tensoren met meer dan bijvoorbeeld 32k elementen vereisen meerdere cijferteksten en brengen meer overhead met zich mee.
  • Het model is klein genoeg zodat bootstrapping (het traagste deel van HE) niet nodig is.
  • Er is geen gesimuleerde netwerkoverhead meegerekend, noch de eenmalige setup per gebruiker voor het genereren en uploaden van het relevante sleutelmateriaal.
  • De server voert een aanzienlijke hoeveelheid (eenmalige, modelspecifieke) pre-computatie uit in de stap "Running preprocessing".

De andere voorbeelden in de repository zijn complexer en hebben daardoor langere latenties en een grotere overhead ten opzichte van platte tekst (met een geheugenvereiste van 60-90 GiB). In het bijzonder:

  • network_anomaly, een ensemble van auto-encoders: 30 seconden voor inferentie.
  • criteo, een recommender-model specifiek aangepast voor HE: 5 minuten voor inferentie.
  • hotword, een convolutioneel netwerk van 10 lagen: 20 minuten voor inferentie.[^4]

Dit klinkt slecht, maar onthoud dat dit single-threaded CPU-uitvoering is. Collega's die HEIR integreren met GPU's rapporteren dat de criteo-workload in ongeveer 500ms draait op een enkele GPU (vergelijkbaar met een H100). Vergeleken met de criteo/cleartext-demo die in 10ms draait, is de vertraging teruggebracht naar 50x.

Mijn punt is dat de uitvoeringstijden blijven verbeteren en voor sommige kleine problemen redelijk kunnen worden genoemd. Bovendien is dit nog zonder versnelling via FPGA's en ASIC's, die qua prestaties nog veelbelovender zijn. In plaats van ruwe prestaties wilde de corporate blogpost vooral de expressiviteit van HEIR tonen: het kan veel modellen compileren en de prestaties zijn op sommige daarvan acceptabel.

Functionaliteiten van HEIR

De repository laat ook zien hoe men HEIR kan gebruiken om:

  • Meerdere HE-backends te targeten en te vergelijken: De repo bevat Lattigo en OpenFHE als voorbeelden.
  • Timing per laag te analyseren: HEIR kan debugging-callbacks invoegen in het gecompileerde programma. In het fraude-voorbeeld zie je bijvoorbeeld dat de eerste lineaire laag het grootste deel van de runtime in beslag neemt (~1,2 seconden van de 2 seconden totaal).
  • Incorrecte inferentie te onderzoeken: De debug-callback kan cijferteksten ontsleutelen om ze te controleren op correctheid of precisieverlies. Bij de fraude-demo is te zien dat het gebruik van HE aan het einde van de inferentie ongeveer 2 bits aan precisieverlies heeft veroorzaakt ten opzichte van het model met platte tekst. Dit verlies kan worden afgestemd met compiler-flags, waarbij een afweging wordt gemaakt met de prestaties.

Een model in HEIR krijgen

Het proces om een vooraf gecompileerd model om te zetten naar iets wat HEIR kan verwerken, is nog niet geautomatiseerd. De twee belangrijkste beperkingen zijn:

  1. Je moet het model kunnen converteren naar MLIR, de intermediate representation die HEIR gebruikt. Veel ML-frameworks zoals PyTorch en JAX hebben tools voor conversie naar MLIR.
  2. Je moet het programma handmatig annoteren met HEIR-specifieke annotaties die aangeven (a) welke inputs van de inferentie geheim zijn en (b) wat de grenzen zijn van de input-ranges voor elke activatiefunctie.

De demo's tonen hoe dit voor PyTorch werkt. Er wordt gewerkt met de maintainers van torch-mlir aan een functie om dit te automatiseren (met gebruik van een validatieset om ranges te schatten). Omdat veel vroege fasen van de compiler-pipeline (zoals het herkennen van activaties) gebaseerd zijn op hoe torch-mlir modellen exporteert, zal ondersteuning voor ONNX of JAX waarschijnlijk niet direct "out of the box" werken. Daarnaast is er nog geen volledige dekking van torch-operators; de ondersteuning van lineair algebraïsche operators in HE is een complex en actief onderzoeksonderwerp.

De compiler aanroepen

De interface van HEIR lijkt meer op LLVM dan op Clang. Er zijn twee binaries, heir-opt en heir-translate, die respectievelijk compiler-passes en codegen afhandelen. Deze binaries zijn verpakt in Bazel-rules (rulesheir), die vervolgens weer zijn verpakt in macros die de optimizer en codegen-binaries koppelen en de resultaten omzetten in een cclibrary, golibrary of rustlibrary.

Het doel is om uiteindelijk één of meerdere Clang-stijl producten te maken met een nauwer bereik (bijv. ondersteuning voor één frontend) en conceptueel eenvoudigere flags (zoals -O2). Voor frontends zoals JAX/PyTorch zal dit waarschijnlijk worden verpakt in een Python-bibliotheek.

Python-bibliotheken

Er bestaat een heir_py frontend-bibliotheek. Deze is momenteel niet gerelateerd aan ML, maar compileert een zeer beperkte subset van Python-bytecode (via de numba-frontend) naar homomorfe encryptie. Het genereert OpenFHE C++ code, compileert en linkt dit met Clang, en laadt de machinecode vervolgens terug in een Python-module. Deze frontend heeft nog veel functies nodig om echt nuttig te worden, zoals het behouden van numpy-ops als MLIR linalg-ops.

Wie heeft er behoefte aan HE?

Er is veel discussie in de HE-community over de "killer app". Veel academici wijzen naar "private LLM's". Hoewel veel onderzoek zich richt op het versnellen van versleutelde LLM-inferentie, is de huidige stand van zaken (SOTA) nog ver verwijderd van praktisch gebruik. Latenties worden gemeten in seconden per token, en vrijwel alle oplossingen vereisen een round-trip naar de client om de output te ontsleutelen voor de volgende stap.

Mijn focus bij Google ligt op de "laaghangende vruchten": welke applicaties kunnen gebruikmaken van de huidige HE-mogelijkheden? Hierbij hanteer ik de volgende heuristieken voor private inferentie:

Modellen uit 2020

Hoewel de latentie is verbeterd, zijn modellen met miljarden parameters nog onbereikbaar. Dit sluit grote transformers uit. Echter, in veel domeinen is de kerncomplexiteit van modelarchitecturen sinds 2020 niet significant veranderd; convoluties zijn nog steeds dominant. Er zijn dus nuttige taken die geen enorme modellen vereisen.

Kritieke privacy voor beide partijen

Voor HE is het nuttig wanneer de privacy van beide partijen kritiek is.

  • De gebruiker: Wil voorkomen dat gevoelige data lekt (bijv. locatiegeschiedenis). Dit kan vaak lokaal op het apparaat, maar dat is niet altijd mogelijk.
  • De provider: Het getrainde model is vaak intellectueel eigendom. Het versturen van een model naar een apparaat (zelfs met een secure enclave) riskeert blootstelling van de modelweights aan concurrenten. Bovendien kunnen aanvallers in een onbeperkte omgeving aanvallen engineeren tegen het productiemodel.

Niet te veel gelijktijdige gebruikers

Een groot bottleneck is sleutelbeheer. HE vereist "evaluatie-sleutels" die uniek zijn per gebruiker en schalen met de programmagrootte. De grootste voorbeelden in de HEIR-demo vereisen ~40 GiB aan evaluatiesleutels. Bij een miljoen gebruikers zou dit enorme opslagkosten en enorme RAM-vereisten met zich meebrengen. Bovendien zorgt het laden van deze 40 GiB van disk naar GPU bij een verzoek voor een latentie die in dezelfde orde van grootte ligt als de gehele HE-berekening.

De server hoeft de output niet te zien

Vaak wil de server één bit informatie weten over de input (bijv. "is dit vingerafdruk correct?"). Omdat HE dit voorkomt, zou de client het resultaat moeten ontsleutelen en terugsturen. Om te voorkomen dat de client liegt, is een zero-knowledge proof nodig, wat de latentie verder verhoogt. Toepassingen waarbij de server de output niet hoeft te zien, komen sneller in productie.

Toepassingen die aan deze eisen voldoen

Enkele ideeën die levensvatbaar lijken:

  • Remote diagnostics: Een fabrikant diagnoseert een machine van een klant. De output (diagnoserapport) is een batch-job die de server niet direct hoeft in te zien, er zijn relatief weinig fabrieken, en de modellen zijn niet extreem complex.
  • B2B-analyse: Twee bedrijven analyseren gezamenlijk data. Er zijn slechts twee gebruikers (sleutelmateriaal is geen probleem), er is wederzijds vertrouwen om de output te delen, en business analytics is meestal tolerant voor latentie.
  • Gespecialiseerde biometrie: Gezichts- of vingerafdrukherkenning bij zeldzame gebeurtenissen (bijv. onboarding van een nieuwe werknemer of accountherstel), waarbij de latentie acceptabel is en de werkgever geen database met biometrie wil bijhouden.

Notities over de roadmap van HEIR

Cryptografie in de compiler

Momenteel zijn HEIR-backends externe library-API's. We bouwen nu aan een echte "backend" waarbij de HE-cryptografie direct als compiler-passes wordt geïmplementeerd. Hiermee kunnen we compileren naar:

  1. Een library-API (hoog niveau).
  2. Modulaire polynoom-aritmetiek (gebruikt door sommige accelerators).
  3. Gevectoriseerde modulaire aritmetiek (bijv. GPU/TPU voor 64-bit integer math modulo custom primes).
  4. LLVM, en vervolgens x86 of WASM.

Dit is cruciaal voor hardwarepartners zoals Niobium, Cornami en Optalysys.

GPU- en TPU-ondersteuning

Er lopen parallelle inspanningen voor de integratie van TPU- en GPU-backends (vooralsnog als library-API's), waaronder de CHEDDAR-bibliotheek, de GPU-bibliotheek van Belfort en CROSS/jaxite.

Nieuwe schema's

HEIR is bedoeld om alle nieuwe schema's te ondersteunen, waaronder:

  • Gentry-Lee: Geoptimaliseerd voor matrixvermenigvuldiging.
  • Poulpy: Vereenvoudigt de onderliggende polynoom-wiskunde van CKKS.
  • Gao-Zheng: Maakt zowel rekenkundige als logische operaties mogelijk in één schema.

Benchmarking

Via fhe-benchmarking.org wordt gewerkt aan een industriestandaard voor benchmarking. HEIR moet code kunnen genereren die direct aan deze repository kan worden ingediend, zodat optimalisaties makkelijk getest en getoond kunnen worden.

Onderzoek

Onderzoekers kunnen hun ideeën direct in HEIR implementeren. Hoewel HEIR zelf nog geen formele publicatie heeft, is het al geciteerd in vier gepubliceerde papers die gebaseerd zijn op onderzoek gedaan in HEIR.

Sluit je bij ons aan!

HEIR is volledig open source. We hebben wekelijkse office hours en een maandelijkse meeting. Iedereen is welkom, maar we vragen om eerst deel te nemen aan een meeting (of een afspraak te maken met een maintainer) voordat er een PR wordt ingediend, om AI-spam te verminderen.

***

[^1]: Zie het introductieartikel voor meer details. [^2]: Deze repo is hergebruikt van het oorspronkelijke "Google transpiler"-project. [^3]: Bazel wordt gebruikt omdat Google dit overal gebruikt en het uitstekend is voor projecten met meerdere talen (C++, Go, Rust, Python). [^4]: Dit voorbeeld heeft prestatieproblemen die nog opgelost moeten worden. [^5]: De anomaly detection demo is bijzonder omdat het model getraind moet worden op referentieverkeer, wat betekent dat men ook HE voor training zou moeten gebruiken.