Technische Samenvatting: SLIM (Sparse Latent Steering voor Interpreteerbare en Eigenschap-gerichte LLM-gebaseerde Moleculaire Redactie)

1. Probleemstelling

Grote Taalmodellen (LLM's) hebben sterke capaciteiten getoond in chemisch redeneren en moleculaire redactie, vaak presterend beter dan traditionele versterkingsleer (RL) en diffusiemodellen door gebruik te maken van voorgeprogrammeerde chemische priors. Een kritieke beperking blijft echter bestaan: eigenschapsrelevante informatie binnen LLM's is impliciet verweven over dichte verborgen toestanden. Deze verwevenheid biedt geen expliciet mechanisme voor eigenschapscontrole.

Bijgevolg vertrouwen bestaande LLM-gebaseerde redactiemethoden op Supervised Fine-Tuning (SFT) voor grove controle. Deze aanpak lijdt aan twee hoofdproblemen:

1. Inefficiëntie: Een aanzienlijk deel van de bewerkingen faalt om de doeleigenschap te verbeteren of verslechtert deze zelfs.
2. Stijfheid: Het verbeteren van een specifieke eigenschap vereist het verzamelen van nieuwe gekoppelde data en het hertrainen van het model, een kostbaar proces dat niet goed schaalbaar is over verschillende eigenschappen of modelarchitecturen.

Er is behoefte aan een methode die LLM's kan sturen naar specifieke moleculaire eigenschappen op het moment van inferentie zonder de parameters van het basismodel te wijzigen, terwijl het tegelijkertijd interpreteerbaarheid biedt over hoe deze eigenschappen worden gecodeerd.

2. Methodologie: Het SLIM-kader

De auteurs stellen SLIM (Sparse Latent Interpretable Molecular editing) voor, een plug-and-play kader dat dichte verborgen toestanden decomposeert in sparse, eigenschap-gealigneerde kenmerken. De pijplijn bestaat uit drie fasen:

Fase 1: Layer Scan

Het kader identificeert eerst het optimale interventiepunt ($l^*$) binnen het model. Met behulp van ridge-regressieprobes op activaties die zijn geëxtraheerd uit 5.000 ZINC-moleculen, scant de methode alle lagen om die te vinden waar de verborgen toestanden het meest voorspellend zijn voor moleculaire eigenschappen (hoogste gemiddelde $R^2$). Interventies op deze laag maximaliseren de invloed omdat de representatie het rijkst is aan de signalen die moeten worden versterkt.

Fase 2: Taakgerichte Sparse Autoencoder (SAE) Training

Een Gated Sparse Autoencoder wordt getraind op de verborgen toestanden op laag $l^*$ met behulp van een multi-objectieve verliesfunctie. In tegenstelling tot standaard SAE's die uitsluitend voor reconstructie worden getraind, is deze SAE taakgericht met vier specifieke doelen:

1. Sparse Reconstructie ($L_{recon}$): Zorgt ervoor dat de SAE activaties trouw reconstrueert om eigenschapsgerelateerde informatie te behouden.
2. Sparsiteit ($L_{sparse}$): Moedigt activatiesparsiteit aan zodat elk molecuul wordt weergegeven door een klein aantal actieve kenmerken.
3. Supervisie Predictie ($L_{sup}$): Per-eigenschap Importance Gates ($w_p$) leren het subset van kenmerken te selecteren dat relevant is voor een specifieke eigenschap. Een MLP-head voorspelt de eigenschapswaarde vanuit deze gegateerde codes, zodat wordt gegarandeerd dat de geselecteerde kenmerken voldoende informatie bevatten om precieze numerieke waarden te regresseren.
4. Contrastieve Alignering ($L_{contrast}$): Een groep-contrastieve loss trekt moleculen met vergelijkbare eigenschapswaarden naar elkaar toe en duwt die met verschillende waarden uit elkaar binnen de gegateerde sparse code-ruimte.
5. Gradient Alignering ($L_{grad}$): Een belangrijke innovatie waarbij de SAE-basis wordt gealigneerd met causale stuurrichtingen. Het model berekent vooraf een op gradiënten gebaseerde richting ($d_{grad}$) vanuit SFT-trainingparen (bron versus verbeterde target). De SAE wordt getraind om ervoor te zorgen dat zijn encoder-decode-pijplijn deze richting behoudt, waardoor wordt gegarandeerd dat de sparse basis causaal is gealigneerd met eigenschapsverbetering in plaats van slechts gecorreleerd.

Fase 3: Activatiesturing op het Moment van Inferentie

Tijdens inferentie genereert het kader een sparse stuurvector voor een doeleigenschap $p$:

1. De op gradiënten gebaseerde richting $d_{grad}^{(p)}$ wordt geprojecteerd door de getrainde SAE.
2. De top-$k$ kenmerken met de grootste magnitudes worden behouden om een sparse stuurvector $d_{steer}^{(p)}$ te vormen.
3. Deze vector wordt toegevoegd aan de residustroom op laag $l^*$ voor alle tokenposities: $h'_{l^*,t} = h_{l^*,t} + \alpha \cdot d_{steer}^{(p)}$.

Dit proces vereist geen hertraining van het basismodel; alleen de basisgewichten blijven bevroren terwijl de stuurrichting wordt toegepast. Voor multi-eigenschapsoptimalisatie worden vectoren gecombineerd via optelling.

3. Belangrijkste Bijdragen

• Plug-and-Play Kader: SLIM decomposeert dichte LLM-verborgen toestanden in sparse, eigenschap-gealigneerde kenmerken, waardoor training-vrije eigenschapscontrole op het moment van inferentie mogelijk wordt.
• Taakgerichte SAE: Introduceert een SAE getraind met contrastieve, voorspellende en gradient-aligneringsdoelen, gekoppeld aan per-eigenschap Importance Gates. Dit zorgt ervoor dat de geleerde kenmerken causaal zijn gealigneerd met doeleigenschappen.
• Interpreteerbaarheid: De sparse basis maakt analyse van redactiegedrag mogelijk, waarbij individuele SAE-kenmerken herkenbare chemische semantics coderen (bijvoorbeeld specifieke functionele groepen of structurele motieven).

4. Experimentele Resultaten

Experimenten werden uitgevoerd op de MolEditRL-benchmark over vier modelarchitecturen (DrugAssist/LLaMA-2, GeLLM3O/LLaMA-3, GeLLM3O/Mistral, MolGen/BART) en acht moleculaire eigenschappen (QED, DRD2, logP, MW, RotBond, SA, HBA, HBD).

• Prestatiewinsten: SLIM verbeterde consistent de succespercentages van bewerkingen (Acc@$\tau$) over alle modellen en eigenschappen.
  • Verbeteringen reikten tot 42,4 procentpunten (bijvoorbeeld MolGen op RotBond, stijgend van 43,2% naar 85,6%).
  • Zelfs de sterkste baseline (GeLLM3O-LLaMA3) zag een gemiddelde verbetering van +3,9 punten.
  • Bij de strikte similariteitsdrempel ($\tau=0.65$) behielden alle modellen positieve gemiddelde winsten, wat aangeeft dat SLIM eigenschappen verbetert zonder structurele similariteit op te offeren.
• Vergelijking met Baselines: SLIM presteerde beter dan externe baselines zoals REINVENT4 (RL-gebaseerd) en MolEditRL (Diffusion+RL) op 4 van de 6 overlappende eigenschappen, met hogere gemiddelde scores ondanks dat het een training-vrije module is.
• Ablatiestudies:
  • Stuurrichting: Op gradiënten gebaseerde richtingen (SLIM) presteerden significant beter dan eenvoudige middelpunt-verschilrichtingen (Contrastive Activation Addition) en willekeurige vectoren. Middelpunt-verschilmethoden degradeerden vaak de prestaties op tel-eigenschappen (HBA, HBD) of slaagden er niet in om continue eigenschappen effectief te verbeteren.
  • SAE Training: De volledige taakgerichte SAE presteerde beter dan een vanilla reconstructie-only SAE. De gradient-aligneringsloss ($L_{grad}$) bleek essentieel, met name voor discrete eigenschapscontrole (tel-eigenschappen), en voorkwam de degradatie die werd waargenomen wanneer deze werd verwijderd.
• Multi-eigenschapssturing: De sparse richtingen bleken compositabel. Het toevoegen van stuurvectoren voor meerdere eigenschappen (bijvoorbeeld HBA + HBD) verbeterde succesvol de gezamenlijke succespercentages, wat minimale interferentie tussen eigenschapsspecifieke kenmerkenubsets aantoont.

5. Betekenis en Claims

Het artikel claimt dat SLIM de "black box"-karakteristiek van eigenschapscontrole in LLM's aanpakt. Door verborgen toestanden te decomponeren in sparse, interpreteerbare kenmerken, biedt het kader:

1. Precieze Controle: Het biedt een expliciete greep voor eigenschapscontrole zonder de kosten van hertraining, en corrigeert specifieke faalmodi van SFT-redactiemethoden (bijvoorbeeld eigenschap-omkerende bewerkingen of inefficiënte modificaties).
2. Interpreteerbaarheid: Casestudies onthullen dat individuele SAE-kenmerken specifieke chemische semantics coderen (bijvoorbeeld Kenmerk F31553 dat moleculair gewicht codeert met een hoge correlatie, of F5418 dat moleculaire complexiteit vastlegt). Dit stelt onderzoekers in staat om te begrijpen waarom een stuurrichting werkt.
3. Generaliseerbaarheid: De methode is architectuur-agnostisch, en werkt effectief over alleen-decoder causale LM's en encoder-decoder modellen heen.

De auteurs merken beperkingen op, waaronder de noodzaak van per-eigenschap afstemming van stuursterkte en interventielagen, de computatiekosten van SAE-training, en de huidige afhankelijkheid van berekenbare eigenschapsorakels in plaats van experimentele assay-eindpunten. Het werk vestigt echter een nieuw paradigma voor controleerbare moleculaire redactie dat de kloof overbrugt tussen dichte neurale representaties en interpreteerbaar chemisch redeneren.