OmniMouse: Scaling properties of multi-modal, multi-task Brain Models on 150B Neural Tokens

Titel: OmniMouse: Schaalbaarheidseigenschappen van multi-modale, multi-task hersenmodellen op 150 miljard neurale tokens

1. Het Probleem

De kunstmatige intelligentie (AI) heeft enorme vooruitgang geboekt door het schalen van datasets en modelgroottes (bijv. Large Language Models), wat heeft geleid tot fundamentele modellen die diverse taken kunnen uitvoeren. Het is echter onduidelijk of dezelfde schaalwetten van toepassing zijn op het modelleren van hersenactiviteit.

• Data-uitdaging: In tegenstelling tot internet-grootte corpora voor taal en beeld, zijn neurale datasets vaak klein, gefragmenteerd en minder divers.
• Bestaande beperkingen: Eerdere modellen voor neurale voorspelling waren meestal beperkt tot één modale (bijv. alleen neurale activiteit), één taak (bijv. alleen stimulus-encoding), of misten schaalbaarheid over verschillende datasets. Ze konden vaak geen combinatie van neurale voorspelling, gedragsdecoding en stimulus-encoding in één architectuur hanteren.
• De vraag: Kunnen hersenmodellen profiteren van schaling, en wat is de huidige bottleneck: modelgrootte of datahoeveelheid?

2. Methodologie: OmniMouse

De auteurs introduceren OmniMouse, een unificerend, multi-modale en multi-task architectuur voor het modelleren van neurale activiteit in de visuele cortex van muizen.

• Dataset:

  • Grootte: Meer dan 3 miljoen enkele neurale eenheden (single-unit recordings) van 73 muizen over 323 sessies.
  • Totale data: Meer dan 150 miljard neurale tokens.
  • Inhoud: Calcium-imaging data (GCaMP6s) uit de visuele cortex, opgenomen terwijl muizen wakker waren, vastgehouden werden en zich gedroegen op een loopband.
  • Stimuli: Natuurlijke films, beelden (ImageNet), en parametrische stimuli (zoals Gabor-patches).
  • Gedragsvariabelen: Loop-snelheid, pupilverwijding, pupilocatie en afgeleiden.

• Architectuur en Tokenisatie:

  • Tokenisatie: Het model gebruikt een flexibele tokenisatiestrategie. Neurale responsen worden omgezet in tokens via een gestructureerde 1D-convolutie. Elke token krijgt leerbare identiteits-embeddings voor de neuron, sessie en dier.
  • Multi-modale input: Het model verwerkt drie modaliteiten: neurale responsen, visuele stimuli (via een Hiera Vision Transformer) en gedragsdata.
  • Flexibel Maskeren: Een kerninnovatie is het gebruik van dynamische masking-configuraties. Het model kan willekeurige combinaties van taken uitvoeren door bepaalde tokens te maskeren (verbergen) en andere als doel te gebruiken. Dit ondersteunt:
    • Neurale voorspelling: Voorspellen van toekomstige activiteit op basis van verleden activiteit.
    • Stimulus-encoding: Voorspellen van activiteit op basis van visuele input.
    • Gedragsdecoding: Voorspellen van gedrag op basis van neurale activiteit.
    • Sub-populatie voorspelling: Voorspellen van een subset van neuronen op basis van een andere subset.
  • Encoder-Decoder: Het model gebruikt een cross-attention encoder om variabele lengte input te comprimeren naar een vaste latent sequence, gevolgd door een multi-modal fusion stack (Transformer) en een decoder die de gemaskerde tokens reconstrueert.

• Training:

  • Het model wordt getraind met een combinatie van Poisson-verlies (voor neurale activiteit) en Mean Squared Error (voor gedragsdata).
  • Er zijn schalingsexperimenten uitgevoerd met modelgroottes variërend van 1M tot 300M parameters, getraind op geneste datasets (van 8 tot 323 sessies).

3. Belangrijkste Bijdragen

1. Systematische Schalingsanalyse: De auteurs tonen aan dat prestaties systematisch verbeteren met meer data, maar verzadigen bij het vergroten van het model (boven ~80M parameters). Dit inverteert het standaard AI-verhaal: in hersenmodellen is data de bottleneck, niet de rekenkracht of modelgrootte.
2. Unificerend Model: OmniMouse is het eerste model dat neurale voorspelling, stimulus-encoding en gedragsdecoding in één architectuur combineert, met flexibele masking voor willekeurige taakcombinaties.
3. State-of-the-Art Prestaties: Het model overtreft gespecialiseerde baselines (zoals MtM, CEBRA, POYO+) in bijna alle evaluatieregimes, zelfs wanneer het op dezelfde hoeveelheid data wordt getraind als de baselines.

4. Resultaten

• Schaalwetten:
  • Data: Het vergroten van het dataset (van 8 naar 323 sessies) leidt tot consistente en voorspelbare verbeteringen in prestaties voor alle modelgroottes.
  • Modelgrootte: Het vergroten van het model (van 1M naar 300M parameters) leidt tot verbeteringen tot ongeveer 80M parameters. Daarna vertonen de leercurven verzadiging of overfitting. Dit suggereert dat de huidige datasets niet complex of divers genoeg zijn om grotere modellen volledig te benutten.
• Prestatieverbetering:
  • OmniMouse behaalt nieuwe state-of-the-art resultaten op de SENSORIUM 2022 en 2023 benchmarks.
  • Bij gedragsdecoding (bijv. pupilgrootte, loop-snelheid) toont het model de meest veelbelovende schalingsdynamiek, waarbij grotere modellen consistent betere correlaties behalen zonder duidelijke verzadiging.
  • Voor neurale voorspelling met visuele input blijven grotere modellen (80M+) verbeteren bij het toevoegen van meer data, wat aangeeft dat ze nog steeds data-gelimiteerd zijn.
• Generalisatie: Het model generaliseert goed naar masking-condities die niet tijdens het trainingstijdperk zijn gezien (bijv. andere populatiegroottes of tijdsvensters), wat aantoont dat het robuuste, generaliseerbare representaties heeft geleerd.

5. Betekenis en Conclusie

OmniMouse markeert een mijlpaal in de computationele neurowetenschap door te bewijzen dat schalingseffecten ook van toepassing zijn op hersenmodellen, maar met een cruciaal verschil ten opzichte van taalmodellen:

• Data-tekort: Terwijl taalmodellen profiteren van enorme datasets om grotere modellen te rechtvaardigen, zijn neurale modellen momenteel beperkt door de beschikbaarheid van hoogwaardige, diverse neurale data. Zelfs in het relatief simpele visuele systeem van de muis zijn de huidige datasets nog niet groot genoeg om de voordelen van zeer grote modellen volledig te realiseren.
• Toekomstperspectief: De observatie van systematische schaling suggereert dat het verzamelen van nog rijkere en grotere datasets (bijv. met meer stimuli, dieren en gedragscontexten) mogelijk "fase-overgangen" kan veroorzaken, waarbij modellen kwalitatief nieuwe vaardigheden ontwikkelen, vergelijkbaar met emergente eigenschappen in Large Language Models.
• Open Science: De code is openbaar beschikbaar, en de auteurs zijn van plan de volledige dataset binnen zes maanden te publiceren, wat een belangrijke stap is voor de gemeenschap om foundation models voor neurowetenschap verder te ontwikkelen.

Samenvattend biedt OmniMouse een krachtig bewijs dat de toekomst van hersenmodellen ligt in het verzamelen van schaalbare, multi-modale datasets, en niet primair in het vergroten van modelarchitecturen.