Technische Samenvatting: Efficiënte Adaptieve Data-acquisitie via Pretrained Belief Representaties (POLAR)

Probleemstelling

Adaptieve data-acquisitie omvat sequentiële leerproblemen zoals Bayesian Experimental Design (BED), Bayesian Optimisation (BO) en Active Learning (AL). In deze settings moet een leerder sequentieel queries (designs) selecteren om een downstream nut (bijv. informatiewinst, identificatie van het globale optimum of voorspellende nauwkeurigheid) te maximaliseren onder een eindig budget.

Bestaande benaderingen voor het leren van acquisitiebeleid vallen over het algemeen in twee problematische categorieën:

1. Posterior-gebaseerde methoden: Deze vertrouwen op het fitten van een probabilistisch model aan geobserveerde data om een posterior te berekenen, die vervolgens wordt gebruikt om een acquisitiefunctie te evalueren (bijv. Expected Improvement, Expected Information Gain). Deze methoden worden beperkt door de kwaliteit van de posterior-approximatie. Model-misspecificatie of approximatiefouten kunnen de acquisitieregel verstoren, en deze fouten stapelen zich op over opeenvolgende rondes.
2. Directe beleidsleer-methoden (Direct policy-learning methods): Deze mappen ruwe observatiegeschiedenissen direct naar beslissingen (bijv. via Reinforcement Learning). Hoewel ze expliciete posterior-inferentie vermijden, moeten ze een complexe mapping van ruwe data naar beslissingen vanaf nul leren. Dit vereist het gelijktijdig leren hoe men de dataset moet representeren én hoe men op die representatie moet handelen, wat vaak leidt tot sample-inefficiëntie en trainingsproblemen.

Methodologie: POLAR

De auteurs introduceren POLAR (Policy Learning with Belief Representations), een framework dat representatie-leren ontkoppelt van beleids-leren. De kerninzicht is dat de optimale sequentiële beslissing alleen afhangt van de observatiegeschiedenis via een belief state. In plaats van deze staat vanaf nul te leren of te vertrouwen op expliciete posterior-approximaties, maakt POLAR gebruik van pre-trained predictive foundation models als belief-state encoders.

Architectuur

1. Belief Encoder (Backbone): POLAR maakt gebruik van een Tabular Foundation Model (TFM), zoals TabPFN of TabICL, die is voorgetraind op diverse synthetische taken. Gegeven een geschiedenis van observaties $D_{t-1}$ en een kandidaat-design $x$, produceert de TFM een kandidaat-geconditioneerde verborgen representatie $r_\psi(x, D_{t-1})$.

   • In tegen tegenstelling tot de uiteindelijke voorspellende distributie-output van de TFM, wordt deze verborgen representatie behandeld als de belief representatie. Het codeert hoe voorspellingen veranderen met data (voorspellende onzekerheid, lokale structuur, context-afhankelijk gedrag) zonder zich vast te leggen op een specifieke posterior-approximatie.
   • Het framework werkt in een pool-gebaseerde setting, waarbij het beleid selecteert uit een eindige set kandidaten $C_t$. De backbone verwerkt de geschiedenis één keer om representaties voor alle kandidaten in de pool te genereren.

2. Policy Head: Een lichtgewicht, leerbare policy head $g_\eta$ mapt de belief-representatie $r_\psi$ naar een scalaire logit voor elke kandidaat. Een softmax over deze logits definieert het stochastische beleid $\pi_\eta$. Bij implementatie wordt de kandidaat met de hoogste score geselecteerd.

3. Trainingsdoel (Training Objective):

   • Policy Loss: De policy wordt getraind via policy gradients op gesimuleerde rollouts. Het doel is om een taakspecifiek nut te maximaliseren (bijv. Expected Information Gain voor BED, terminal best value voor BO). De auteurs gebruiken een lokale credit assignment strategie (één-stap incrementen) om variantie te verminderen, in plaats van een volledige long-horizon value function.
   • Backbone Adaptatie (Optioneel): Om de pretrained backbone af te stemmen op de specifieke downstream taakdistributie, wordt een auxiliary supervised prediction loss ($L_{pred}$) toegepast op de parameters van de backbone.
   • Ontkoppeling (Decoupling): Een cruciale ontwerpkeuze is dat policy gradients niet worden gepropageerd naar de backbone. De backbone wordt alleen bijgewerkt via de supervised prediction loss, terwijl de policy head wordt bijgewerkt via de policy gradient. Dit voorkomt dat hoog-variante reinforcement learning signalen de voorgetrainde representaties destabiliseren.

Theoretische Fundering

Het artikel biedt een verenigd decision-theoretic perspectief, waarbij wordt aangetoond dat het Bayes-optimale beleid $\pi^*$ een functie is van de taak-relevante posterior $p(w|D_{t-1})$. Omdat de posterior een functie is van de geschiedenis, is elke representatie $e_t$ die de posterior mogelijk kan herstellen (via een decoder), een geldige belief-representatie. POLAR stelt dat voorgetrainde TFMs dergelijke representaties bieden die rijk genoeg zijn om optimale acties te bepalen, maar gestructureerd genoeg zijn om de mapping naar de policy head eenvoudig te maken, waardoor de extremen van ruwe geschiedenis (te complex) of expliciete posterior (gevoelig voor approximatiefouten) worden vermeden.

Belangrijkste Bijdragen

1. Verenigd Perspectief: De auteurs formaliseren adaptieve data-acquisitie als het leren van policies op basis van belief-representaties, waarmee ze de kloof overbruggen tussen ruwe geschiedenis-gebaseerde en posterior-gebaseerde benaderingen.
2. Ontkoppelde Architectuur: Ze stellen een eenvoudige architectuur voor die voorgetrainde foundation models gebruikt als fixed (of gefinetuned) belief encoders en een aparte policy head traint, waardoor representatie-leren en besluitvorming worden gescheiden.
3. Sample Efficiëntie: Empirische resultaten tonen aan dat POLAR state-of-the-art prestaties bereikt over diverse taken (BED, BO, AL) terwijl het aanzienlijk minder training samples vereist (tot wel $100\times$ minder) vergeleken met bestaande geamortiseerde methoden.

Experimentele Resultaten

De auteurs evalueren POLAR op drie belangrijke benchmarks:

1. Bayesian Experimental Design (BED):

   • Taken: Location Finding (2D en 5D) en Constant Elasticity of Substitution (CES).
   • Resultaten: POLAR presteert aanzienlijk beter dan state-of-the-art geamortiseerde methoden (DAD, RL-BOED, ALINE). In de 5D Location Finding taak overtreft POLAR de piekprestaties van RL-BOED en DAD met 10.000$\times$ minder training samples en overtreft het ALINE met 100$\times$ minder samples.
   • Ablaties:
     • Frozen vs. Finetuned Backbone: Het finetunen van de backbone met een auxiliary prediction loss levert betere prestaties op dan het bevriezen van de backbone, hoewel zelfs de bevroren versie competitief is.
     • Decoupled vs. Coupled Gradients: Het toestaan dat policy gradients de backbone updaten (coupled) degradeert de prestaties ernstig, wat de noodzaak van ontkoppeling bevestigt.
     • Pretrained vs. Scratch: Het trainen van de architectuur vanaf nul convergeert veel langzamer en presteert slechter, wat de waarde van het voorgetrainde foundation model benadrukt.

2. Bayesian Optimisation (Hyperparameter Optimisation):

   • Benchmark: HPO-B (Hyperparameter Optimisation Benchmark).
   • Resultaten: POLAR bereikt de laagste gemiddelde regret en de beste gemiddelde rang over zes zoekruimtes, waarbij het klassieke GPs, Meta-GPs en end-to-end geamortiseerde methoden zoals NAP overtreft. Het demonstreert dat het gebruik van een voorgetrainde TFM backbone sterkere prestaties oplevert dan het leren van representaties vanaf nul.

3. Moleculaire Optimisatie:

   • Benchmark: DOCKSTRING (high-dimensional moleculaire docking).
   • Resultaten: In een 512-dimensionale discrete ruimte bereikt POLAR de laagste regret vergeleken met GPs met Tanimoto kernels, Random Forests en een TabICL baseline met een handontworpen acquisitiefunctie. Dit bevestigt dat het direct leren van de acquisitieregel vanuit de belief-representatie beter presteert dan vaste acquisitiestrategieën.

Betekenis en Claims

Het artikel claimt dat POLAR een belangrijke stap is in de schaalbaarheid en efficiëntie van geamortiseerde data-acquisitie. Door het paradigma te verschuiven van het gezamenlijk leren van representaties en beslissingen naar het handelen op voorgetrainde predictive belief-representaties, mitigeert POLAR de sample-inefficiëntie en trainingsinstabiliteit die geassocieerd worden met directe policy-learning en de approximatiefouten van posterior-gebaseerde methoden.

De auteurs benadrukken dat hun aanpak niet gaat over het ontwerpen van een nieuwe estimator of trainingsalgoritme, maar over het hergebruiken van een voorgetraind predictief model als belief encoder. Dit maakt het mogelijk om een kleine policy head te trainen met standaard policy gradients, wat resulteert in superieure prestaties met drastisch lagere computationele kosten voor training. Het framework wordt gepresenteerd als een verenigde oplossing toepasbaar op BED, BO en AL, waarbij enkel de taakspecifieke utility wordt gebruikt voor training.