pretrain.models#

The pretrained encoders, their maskers, and the pretrainers that fit them.

Each model is one directory under src/pretrain/models/, pairing an encoder with the pretrainer that fits it, and a masker where the objective needs one. This page is grouped the same way, one section per directory.

NDT Stitch#

`NDTStitch`	Multi-session NDT with per-session stitchers [Ye and Pandarinath, 2021].
`NDTStitchMasker`	Temporal masking for spikes (BERT-style).
`NDTStitchPretrain`	Masked spike modelling for NDT Stitch, which fits one stitcher per session.

MtM#

`MtM`	Multi-task masked transformer over binned spikes [Zhang et al., 2024].
`MtMMasker`	Module for masking spikes.
`MtMPretrain`	Masked spike modelling for MtM, whose masks carry a mode token.

NDT2#

`NDT2`	Multi-context masked autoencoder over patched spike tokens [Ye et al., 2023].
`NDT2Masker`	ShuffleInfill-style MAE masker for NDT2.
`NDT2Pretrain`	Masked spike modelling for NDT2, which masks context tokens and accumulates gradients.

NEDS#

`NEDS`	Multimodal masked transformer over spikes and behavior [Zhang et al., 2025].
`NEDSMasker`
`NEDSPretrain`	Multimodal masked modelling for NEDS, over spikes and the behaviors alike.

POYO#

`POYO`	Transformer-based model for neural decoding from spike trains [Azabou et al., 2023].
`POYOSingleTaskPretrain`	Supervised pretraining of POYO on one behavior.

POYO+#

`POYOPlus`	POYO+ over chained spike tokens [Azabou et al., 2025].
`POYOPlusMultitaskPretrain`	Supervised pretraining of POYO+ on every behavior at once, via its readout heads.

POSSM#

Pretrained on one behavior or on several at once, hence two trainers.

`POSSM`	POSSM (POYO + State Space Model) for the IBL benchmark [Ryoo et al., 2025].
`POSSMSingleTaskPretrain`	Supervised pretraining of POSSM on one behavior.
`POSSMMultitaskPretrain`	Multi-task pretrainer for POSSM.

RRR#

`RRRDecoder`	Reduced-rank decoder mapping a window of binned spikes to behavior [Zhang et al., 2026].
`RRRSingleTaskPretrain`	Supervised pretraining of the shared basis `V` on one behavior.

NuCLR#

A unit encoder for TS3, trained against the contrastive loss listed here.

`NuCLR`	Contrastive unit encoder over population context [Arora et al., 2025].
`NuCLRLoss`
`NuCLRPretrain`	Contrastive pretraining of the NuCLR unit encoder, against `NuCLRLoss`.

NEMO#

A unit encoder for TS3, assembled from the submodules listed here.

`NEMO`	NEMO bimodal contrastive learning model for IBL brain region pretraining [Yu et al., 2025].
`NEMOPretrain`	Trainer for NEMO SSL pretraining (CLIP-style contrastive learning).
`WVFEncoder`
`ACGEncoder`
`LinearProjector`	Linear -> LayerNorm projection head (from NEMO LinearProjector, layer_norm=True).