Scaling an Autoregressive Transformer for Single-Cell Generation

Technische Samenvatting: Schalen van een Autoregressieve Transformer voor Single-Cell Generatie

Probleemstelling
Het artikel adresseert de uitdaging van het genereren van realistische single-cell genexpressievectoren voor specifieke celtypen. Hoewel bestaande methoden cellijnen vertegenwoordigen via categorische namen, enkele voorbeeldvectoren of sequenties van vectoren, is er behoefte aan een schaalbaar, zelfgesuperviseerd raamwerk dat synthetische genexpressieprofielen kan genereren om schaarse metingen aan te vullen en in-silico experimentatie mogelijk te maken. Bovendien ontbert in het vakgebied een duidelijk begrip van hoe de pretraining-loss schaalt met de modelgrootte ($N$) en de datagrootte ($D$) voor single-cell foundation models, waarbij tegenstrijdige rapporten bestaan over de vraag of deze modellen een voorspelbare power-law scaling volgen of afvlakken door beperkingen in de datadiversiteit.

Methodologie
De auteurs stellen een autoregressieve transformer-architectuur voor, gekoppeld aan een geleerde gekwantiseerde VAE-tokenizer. De aanpak omvat twee hoofdcomponenten die sequentieel worden getraind:

1. Tokenizer (RQ-VAE): Een Residual Quantized VAE brengt continue single-cell genexpressievectoren (18.080 genen) in kaart naar een sequentie van 8 discrete tokens. De encoder comprimeert de input naar een 256-dimensionale latente vector, die vervolgens grof-naar-fijn wordt gekwantiseerd over acht residuele codeboeken. Deze tokenizer wordt eenmaal getraind op de scBaseCount atlas en vervolgens bevroren om te dienen als een discrete representatielaag voor de transformer.
2. Transformer Model: Een causale transformer (gebaseerd op de LLaMA-architectuur) modelleert sequenties van deze discrete tokens. Een "zin" wordt geconstrueerd door tokenblokken van 32 cellen van hetzelfde celtype aan elkaar te koppelen, gescheiden door geleerde separator-tokens. Het model wordt getraind via next-token prediction met behulp van cross-entropy loss.
3. Scaling Grid: Om scaling laws te onderzoeken, hebben de auteurs een gecontroleerd grid geconstrueerd waarin de modelgrootte ($N$) en de datagrootte ($D$) onafhankelijk van elkaar variëren.
   • Modelgroottes: 5 configuraties variërend van 1,31M tot 83,9M parameters, waarbij een vaste breedte-diepteverhouding wordt aangehouden.
   • Datagroottes: Trainingsbudgetten variërend van 38,1M tot 362M cellen, gesampled uit een vaste subset van ongeveer 200M cellen uit de scBaseCount atlas met proportionele sampling om de natuurlijke celtype-compositie te behouden.
   • Totaal: 22 verschillende $(N, D)$-punten werden geëvalueerd.

Belangrijkste Bijdragen

1. Taakformulering: Definieerde een zelfgesuperviseerde leeropdracht voor het genereren van synthetische genexpressievectoren gegeven een cellijn-context, gebruikmakend van discrete tokenisatie.
2. Evaluatiemetrieken: Stelde downstream metrieken vast voor biologische getrouwheid, inclus\nief Pearson-correlatie van gemiddelde expressie, Mean Absolute Error (MAE) en een celtype-discriminatiescore.
3. Loss Correlatie: Demonstreerde dat de pretraining cross-entropy loss sterk correleert met downstream generatiekwaliteitsmetrieken.
4. Joint Scaling Law: Paste een gezamenlijke, twee-exponentiële scaling law $L(N, D) = E + A N^{-\alpha} + B D^{-\beta}$ toe op de pretraining loss. Naar weten van de auteurs is dit de eerste dergelijke fit voor een single-cell transformer op een onafhankelijk gevarieerd $N \times D$ grid.
5. Compute-Optimale Frontier: Leidde een compute-optimale allocatiestrategie af op basis van de passende exponenten, wat suggereert dat computationele middelen op het huidige schaalniveau beter besteed kunnen worden aan het vergroten van het datavolume dan aan modelparameters.

Resultaten

• Generatiekwaliteit: Het model genereert celtype-specifieke vectoren die nauw aansluiten bij empirische replica's. Bij het laagste validatie-loss checkpoint bereikte de gegenereerde populatie een Pearson-correlatie van 0,960 en een MAE van 0,041 ten opzichte van de uit de dataset gehouden empirische data, waarmee het bijna de baseline evenaart van het vergelijken van twee disjuncte empirische subsets (0,964 en 0,038, respectievelijk). De celtype-discriminatie was 0,990, wat aangeeft dat de generator niet instort naar een enkele modus.
• Scaling Law Fit: De passende scaling law leverde exponenten $\alpha \approx 0,81$ (modelgrootte) en $\beta \approx 0,57$ (datagrootte) op, met een onherleidbare loss-vloer $E \approx 2,90$ nats. De fit bereikte een $R^2$ van 0,977.
• Voorspellende Kracht: De passende wet voorspelde helden-runs (waarbij zowel $N$ als $D$ varieerden) met fouten binnen 0,5% van de gemeten waarden, wat de bekwaamheid van het model om te extrapoleren valideert.
• Compute-Optimale Frontier: Omdat $\alpha > \beta$, helt de optimale allocatie naar data ($D^* \propto C^{0,59}$ vs. $N^* \propto C^{0,41}$).

Betekenis en Claims
Het artikel claimt het debat over single-cell scaling laws te hebben opgelost. Waar eerdere studies (de "diversity-over-scale" cluster) rapporteerden dat single-cell modellen afvlakken bij grote schalen, wat zij toeschreven aan een gebrek aan contextdiversiteit, stellen de auteurs dat dit een misinterpretatie is van capaciteitsbottlenecks. Hun resultaten tonen aan dat wanneer $N$ en $D$ gezamenlijk worden gevarieerd, een duidelijke power-law scaling verschijnt. De geobserveerde plateaus in eerder werk worden toegeschreven aan het vergroten van de data bij een vaste of ongecontroleerde modelgrootte, waarbij het model de capaciteit mist om extra data te exploiteren zodra de dataterm onder de capaciteitsterm valt.

De auteurs positioneren dit werk als een fundamentele stap voor single-cell foundation models, waarbij ze een compute-optimale frontier bieden voor training. Ze bespreken ook een potentiële "Stage 2" toepassing: het fine-tunen van het pretrainde model voor het voorspellen van perturbatie-respons. Ze merken echter op dat een observatie tijdens de pretraining waarbij een helden-loss op een perturbatie-benchmark een "vinkje-vormig" traject vertoonde (dalend en dan stijgend), verschillend van de hoofd-scaling loss, wat wijst op een compositie-mismatch tussen de algemene pretraining corpus en specifieke perturbatie-contexten. Dit blijft een open vraag voor toekomstig onderzoek in plaats van een gerapporteerd resultaat.

Het artikel concludeert dat de voorgestelde architectuur succesvol genexpressievectoren synthetiseert met een hoge biologische getrouwheid en dat het scaling gedrag van single-cell transformers een voorspelbare, gezamenlijke power law volgt, vergelijkbaar met die in natuurlijke taalverwerking, mits het trainingsgrid voldoende gecontroleerd is.