Technische Samenvatting: It Just Takes Two: Schalen van Geamortiseerde Inferentie naar Grote Mengen

Probleemstelling

Neurale Posteriorschatting (NPE) is uitgegroeid tot een standaardinstrument voor geamortiseerde probabilistische inferentie, met name in wetenschappelijke domeinen waar men een doelparameter $\theta$ moet afleiden uit een verzameling waarnemingen $X = \{x_i\}_{i=1}^N$. Een kritieke uitdaging doet zich voor wanneer de waarnemingen in de verzameling gekoppeld zijn door gedeelde, onbekende storende factoren $\psi$ (bijvoorbeeld gemeenschappelijke detector-calibratie, scèneverlichting of moleculaire identiteit) die niet direct van belang zijn, maar die de hele verzameling collectief beïnvloeden.

In dergelijke scenario's vereist optimale inferentie de gezamenlijke verwerking van de verzameling om rekening te houden met deze gedeelde factoren. Bestaande benaderingen staan echter voor een schaalbaarheidstrade-off:

1. Strategieën per waarneming: Trainen op individuele metingen en scores hercombineren op het moment van testen. Hoewel dit schaalbaar is in $N$, faalt dit wanneer gedeelde storende factoren de waarnemingen koppelen, wat leidt tot incorrecte inferentie.
2. Strategieën voor gezamenlijke verzamelingen: Gebruik van permutatie-invariante architecturen (bijvoorbeeld Deep Sets) die end-to-end worden getraind op verzamelingen van grootte $N$. Hoewel dit principieel correct is, schalen de geheugen- en rekentkosten voor training lineair met de doel-implementatiecardinaliteit $N$. Dit maakt training op grote verzamelingen (bijvoorbeeld $N$ in de duizenden) computergewijs onhaalbaar.

Beoefenaars worden dus gedwongen te kiezen tussen een schaalbare maar suboptimale schatter of een correcte maar niet-schaalbare schatter.

Methodologie: PAIRS

De auteurs introduceren PAIRS (Pretraining Aggregators for Inference at aRbitrary Set-sizes), een drie-trapsprocedure die representatieleren ontkoppelt van posterieormodellering om schaalbaarheid te bereiken zonder correctheid op te offeren.

Theoretisch Inzicht

De kerntheoretische bijdrage is het bewijs dat het leren van een permutatie-invariante datarepresentatie (de encoder) geen training op grote verzamelingen vereist. Onder de aanname dat de data toelaat tot een eindig-dimensionale voldoende statistiek via gemiddelde pooling, bewijzen de auteurs dat training op verzamelingen van grootte $N \le 2$ voldoende is om dezelfde encoder te herstellen (tot op een affiene herparameterisatie) als training op elke grotere cardinaliteit. Dit berust op het feit dat paren de kleinste verzamelingsgrootte zijn die de koppeling blootlegt die wordt veroorzaakt door gedeelde storende factoren, terwijl grotere verzamelingen geen nieuwe informatie bieden over het onderliggende generatieve proces boven wat paren al onthullen.

De PAIRS-procedure

De methode bestaat uit drie fasen:

1. Pretraining (Kleine Cardinaliteit): Een Deep Set-encoder (gemiddelde pooling) en een inferentiehoofd (conditioneel generatief model) worden gezamenlijk getraind op verzamelingen van waarnemingen van grootte $N \in \{1, 2\}$. Deze fase leert de representatie $t_\omega(x)$ en de initiële posterieure mapping.
2. Embedding Caching: De getrainde encoder wordt bevroren. Voor de finetuning-dataset (die verzamelingen van willekeurige grote grootte $N$ kan bevatten), worden individuele waarnemingen door de encoder gevoerd, en worden de resulterende embeddings vooraf samengevoegd via gemiddelde pooling.
3. Finetuning (Grote Cardinaliteit): Het inferentiehoofd wordt gefinetuned op de vooraf samengevoegde embeddings. Omdat de encoder bevroren is en embeddings vooraf zijn berekend, zijn de geheugen- en rekentkosten van deze fase onafhankelijk van de implementatieverzamelingsgrootte $N$.

Belangrijkste Bijdragen

1. Theoretische Garantie (Stelling 1): De auteurs bewijzen dat onder milde regulariteitsvoorwaarden een Deep Set-encoder met gemiddelde pooling, getraind op verzamelingen van maximaal twee elementen, een voldoende statistiek voor $\theta$ herstelt bij elke cardinaliteit $N \ge 1$. Dit impliceert dat de kosten voor representatieleren ontkoppeld zijn van de implementatieverzamelingsgrootte.
2. Het PAIRS-algoritme: Een praktische, drie-trapspijplijn die training op kleine verzamelingen mogelijk maakt, terwijl inferentie op verzamelingen van duizenden elementen wordt mogelijk gemaakt. De trainingskosten zijn in wezen onafhankelijk van de doelwaarde $N$.
3. Empirische Validatie: Uitgebreide experimenten over scalaire, beeld-, multi-view 3D-, moleculaire en hoog-dimensionale conditionele generatie-benchmarks met $N$ variërend tot duizenden.

Experimentele Resultaten

Het artikel evalueert PAIRS tegen verschillende baselines, waaronder marginalisatie per waarneming ($N=1$), uitgebreide pretraining ($N=1\text{--}10$), Gaussische regressie en end-to-end training bij grote $N$ (waar haalbaar).

• Prestaties: PAIRS komt overeen met of overtreft standaard baselines in alle taken. Opmerkelijk is dat het de analytische posterior van een bivariate Gaussisch model volgt tot $N=10^5$ en overeenkomt met MCMC-schattingen bij een "bump hunt"-taak in de deeltjesfysica met gedeelde storende factoren.
• Schaalbaarheid: In taken waar end-to-end training bij $N=1000$ computergewijs haalbaar is, bereikt PAIRS vergelijkbare prestaties voor een fractie van de rekentkosten (ongeveer 100x goedkoper). Bij hoog-dimensionale generatieve taken (synthese van nieuwe weergaven) is end-to-end training bij grote $N$ onhaalbaar, terwijl PAIRS succesvol schaalt.
• Embedding-dimensie: Experimenten tonen aan dat de prestaties verzadigen zodra de embedding-dimensie $l$ groot genoeg is om de complexiteit van het probleem te vangen, wat bevestigt dat $l$ niet hoeft te groeien met $N$.
• Kosten-prestatie trade-off: PAIRS ligt op de Pareto-grens en biedt de beste balans tussen posterior-kwaliteit en rekentkosten. Strategieën met pretraining op $N=1\text{--}10$ bleken aanzienlijk duurder en presteerden vaak slechter dan PAIRS.

Betekenis en Beweringen

Het artikel claimt de trade-off tussen schaalbaarheid en correctheid in geamortiseerde inferentie voor verzameling-gebaseerde data op te lossen. Door aan te tonen dat "klein trainen, groot implementeren" theoretisch onderbouwd is voor architecturen met gemiddelde pooling, stelt PAIRS de toepassing van neurale posteriorschatting op problemen met grote waarnemingsverzamelingen mogelijk die eerder onhandelbaar waren.

De auteurs benadrukken dat hun aanpak niet louter een heuristiek is, maar wordt ondersteund door een bewijs dat aantoont dat paren voldoende zijn om de noodzakelijke additieve structuur van de voldoende statistiek te leren. Zij merken op dat hoewel de theorie uitgaat van een marginaal likelihood van het exponentiële familie-type (een beperking die inherent is aan Deep Sets met gemiddelde pooling), de empirische resultaten suggereren dat deze idealisatie mild genoeg is om zich goed te generaliseren naar realistische, niet-EF-problemen. Het werk opent wegen voor het toepassen van geamortiseerde inferentie op grote wetenschappelijke datasets zonder de prohibitieve rekentkosten van end-to-end training bij implementatie-cardinaliteiten.