Technische Samenvatting: Discrete Langevin-geïnspireerde Posterior Sampling (∆LPS)

Probleemstelling

Het artikel behandelt de uitdaging om inverse problemen op te lossen in discrete toestandsruimten. Inverse problemen betreffen het herstellen van een onbekend signaal $x$ uit indirecte, schaarse en ruisbeïnvloede metingen $y$ (gemodelleerd als $y = A(x) + n$). Hoewel continue diffusiemodellen de standaard zijn geworden voor het oplossen van dergelijke problemen via posterior sampling, blijven hun discrete tegenhangers onderbelicht.

Bestaande methoden voor discrete posterior sampling lijden onder drie primaire beperkingen:

1. Continue Relaxaties: Methoden zoals G2D2 relaxeren categorische variabelen naar continue ruimten om gradiëntgeleiding te gebruiken, waardoor de inferentiegeometrie wordt gewijzigd en het risico ontstaat dat geprojecteerde discrete tokens niet consistent decoderen.
2. Sequentiële Updates: Methoden zoals SGDD vertrouwen op Gibbs-stijl updates, die sequentieel zijn en traag mengen in hoogdimensionale ruimten.
3. Specialisatie: Methoden zoals APS zijn sterk gekoppeld aan specifieke corruptieprocessen (bijvoorbeeld gemaskerde diffusie) en trainingsparadigma's, wat hun generaliteit beperkt.

Het doel is het ontwikkelen van een posterior sampler die volledig discreet blijft, alle token-dimensies parallel updatet en agnostisch is ten opzichte van het specifieke trainingsparadigma van de discrete diffusie-prior.

Methodologie: ∆LPS

De auteurs stellen ∆LPS (Discrete Langevin-geïnspireerde Posterior Sampler) voor, een trainingsvrij algoritme dat gradiëntinformatie gebruikt om discrete overgangen te sturen zonder de discrete toestandsruimte te verlaten.

Kernformulering

Het probleem wordt geformuleerd in de discrete latente ruimte $z_0 \in \{1, \dots, K\}^L$ van een encoder-decoder-model. Het doel is om te bemonsteren uit de posterior $p(z_0|y) \propto p(y|z_0)p_\theta(z_0)$, waarbij $p_\theta$ een vooraf getrainde discrete diffusie-prior is. De sampler benadert de posterior-overgang $p(z_s|z_t, y)$ door eerst een schone latent $z_0$ af te leiden die consistent is met de meting en de prior, en deze vervolgens opnieuw te verstoren naar de volgende tijdstap.

Discrete Langevin-voorstel

De kerninnovatie is het aanpassen van de Langevin-dynamica aan discrete ruimten. In continue omgevingen verplaatst een Langevin-update een steekproef $z_0$ in de richting van de gradiënt van de log-posterior-potentiaal $U(z_0) = \log p(y|z_0) + \log p_\theta(z_0; z_t)$. Aangezien $z_0$ discreet is, voegen de auteurs geen gradiëntvector direct toe. In plaats daarvan gebruiken ze de Gaussische voorstelverdeling uit continue Langevin-dynamica om overgangskansen te definiëren over geldige discrete kandidaten $z'_0$:

$$q(z'_0|z_0; z_t, y) \propto \exp\left( -\frac{1}{4\eta} \|z'_0 - z_0 - \eta \nabla_{z_0} U(z_0; z_t, y)\|_2^2 \right)$$

Door de exponent uit te werken en te normaliseren over de eindige set geldige tokens, factoriseert het voorstel over token-coördinaten. Dit stelt de sampler in staat om voor elke tokenpositie $\ell$ onafhankelijk een categorische verdeling te berekenen:
$$q_\ell(z'_0[\ell]) = \text{Cat}\left( \sigma\left( \frac{1}{2} \nabla_{z_0[\ell]} U \cdot (k - z_0[\ell]) - \frac{(k - z_0[\ell])^2}{4\eta} \right) \right)$$

Deze mechanisme maakt volledig parallelle updates mogelijk over alle $L$ token-dimensies, waardoor de sequentiële bottleneck van Gibbs-samplers wordt vermeden.

Generalisatie en Preconditie

• Prior-agnosticisme: De methode vereist alleen voorspellende kansen over categorische toestanden, waardoor deze compatibel is met gemaskerde diffusie (MDLM) en uniform-toestand diffusie (Duo) priors.
• Blinde inverse problemen: Voor niet-differentieerbare of onbekende voorwaartse operatoren wordt de likelihood-gradiënt vervangen door een geleerde contrastieve surrogate (met behulp van encoders voor data en metingen) om een glad signaal voor leiding te bieden.
• Adam-preconditie: Om slecht geschaalde gradiënten in hoogdimensionale latente ruimten te hanteren, wordt de ruwe gradiënt vervangen door een Adam-gepreconditioneerde richting. Dit past de updates adaptief aan, dempt ruisige grote gradiënten terwijl consistent nuttige worden behouden, en verbetert zo de menging en convergentie.

Belangrijkste Bijdragen

1. Volledig Discrete Posterior Sampling: ∆LPS is een trainingsvrije sampler die elke iteratie in de discrete toestandsruimte houdt, waardoor soft-token optimalisatie en projectiefouten worden vermeden.
2. Efficiënte Parallelle Updates: Door het voorstel te factoriseren over coördinaten, update de methode alle token-dimensies parallel, wat de efficiëntie aanzienlijk verbetert ten opzichte van sequentiële Gibbs-stijl samplers.
3. Prior-agnostisch Ontwerp: De sampler is niet gebonden aan een specifiek corruptieproces (gemaskerd versus uniform) en werkt met elke vooraf getrainde discrete diffusie-prior die voorspellende kansen biedt.

Experimentele Resultaten

De auteurs evalueren ∆LPS op een diverse reeks inverse problemen over MNIST, CIFAR-10 en FFHQ, waarbij lineaire, niet-lineaire en blinde situaties worden bestreken.

• Beeldherstel (FFHQ): Op taken zoals inpainting, Gaussian deblurring, super-resolutie, HDR en motion deblurring, bereikt ∆LPS prestaties die vergelijkbaar zijn met of beter zijn dan continue baselines (bijvoorbeeld DPS, PSLD) en presteert beter dan bestaande discrete baselines (SGDD, G2D2, APS) op het gebied van perceptuele kwaliteit (LPIPS).
• Discrete Taken (MNIST/CIFAR): De methode toont robuustheid over verschillende moeilijkheidsgraden (Makkelijk, Gemiddeld, Moeilijk). Het komt consistent overeen met of overtreft SGDD en G2D2 in token-nauwkeurigheid en PSNR, zelfs met aanzienlijk lagere rekenbudgetten (minder functiewaarderingen) dan de "volledige" SGDD-variant.
• Blinde Inverse Problemen: In een ruimtelijke mapping-taak (het afleiden van plattegronden uit wandeltrajecten waarbij de voorwaartse operator onbekend is), presteert ∆LPS beter dan continue baselines (CoGuide, DPS-varianten) op de F1-score en Intersection over Union (IoU).

Betekenis en Claims

Het artikel stelt dat volledig discrete, gradiënt-informeerde posterior samplers een schaalbaar en algemeen pad bieden voor het oplossen van inverse problemen over discrete representaties.

• Schaalbaarheid: Het parallelle update-mechanisme adresseert de beperkingen in mengsnelheid van discrete samplers in hoogdimensionale ruimten.
• Generaliteit: De methode ontkoppelt het inferentie-algoritme van het specifieke trainingsparadigma van de prior, waardoor het kan profiteren van het groeiende ecosysteem van discrete diffusiemodellen (tekst, code, afbeeldingen, audio).
• Praktijktoepasbaarheid: Door continue relaxaties te vermijden, zorgt ∆LPS ervoor dat de uiteindelijke steekproeven geldige discrete tokens zijn, waardoor de "projectiegap" wordt voorkomen waarbij een continue benadering faalt om te decoderen naar een consistent discreet signaal.

De auteurs concluderen dat naarmate discrete generatieve modellen groeien, principiële inferentiemethoden zoals ∆LPS steeds kritischer zullen worden voor toepassingen in biomoleculair ontwerp, symbolische muziek, code-aanvulling en andere domeinen waar data inherent discreet is.