STNDT

Spatiotemporal Neural Data Transformer for binned spike counts.

Source

  • Registry name: stndt
  • Model class: ladys.models.stndt.STNDT
  • Config class: ladys.models.stndt.STNDTConfig
  • Source file: src/ladys/models/stndt.py

When to use

Use STNDT when NDT-style masked count modeling should also learn spatial dependencies between neurons. The adapter ports the PyTorch STNDT design: temporal self-attention over time, spatial self-attention over neurons, a spatially mixed temporal stream, and optional SimCLR-style contrastive consistency between two independently masked views.

Assumptions

STNDT expects raw nonnegative spike counts. On synthetic datasets it reconstructs the observed neurons. When built by Experiment on an NLB dataset with output_mode: auto, it expands the readout to held-in plus held-out training neurons, feeds zeros for unavailable held-out inputs, and scores the held-out output slice with co-smoothing bits/spike.

Outputs

forward returns natural-space rates, spatiotemporal transformer factors, and masking/contrast diagnostics in extras. Training uses masked held-in entries plus any available NLB held-out/forward targets; evaluation uses all available target entries.

Configuration

Config for the Spatiotemporal Neural Data Transformer adapter.

Field Type Default
name Literal['stndt'] 'stndt'
objective str 'stndt_masked_poisson_nll'
ensemble bool False
ensemble_size int 2
output_neurons Optional[int] None
output_mode Literal['auto', 'heldin', 'heldin_heldout'] 'auto'
fwd_steps int 0
context_forward int 4
context_backward int 8
context_wrap_initial bool False
full_context bool False
hidden_size int 128
dropout float 0.1
dropout_rates float 0.2
dropout_embedding float 0.2
num_heads int 2
num_layers int 6
activation Literal['relu', 'gelu'] 'relu'
linear_embedder bool False
embed_dim Literal[0, 1] 1
learnable_position bool False
max_spike_count int 20
lograte bool True
log_rate_min float -8.0
log_rate_max float 8.0
spike_log_init bool False
fixup_init bool True
pre_norm bool True
scale_norm bool False
decoder_layers int 1
position_offset bool False
mask_ratio float 0.25
mask_mode Literal['full', 'timestep', 'neuron', 'timestep_only'] 'full'
mask_token_ratio float 1.0
mask_random_ratio float 0.5
mask_max_span int 1
mask_span_expand_prob float 0.0
use_zero_mask bool True
topk_loss_fraction float 1.0
do_contrast bool True
contrast_mask_ratio float 0.05
contrast_mask_mode Literal['full', 'timestep', 'neuron', 'timestep_only'] 'full'
contrast_mask_token_ratio float 0.5
contrast_mask_random_ratio float 0.5
contrast_mask_max_span int 1
contrast_mask_span_expand_prob float 0.0
temperature float 0.07
contrast_lambda float 0.1
use_contrast_projector bool False
linear_projector bool True
contrast_layer Literal['embedder', 'decoder'] 'embedder'
nlb_decoder Literal['direct', 'latents'] 'direct'
optimization OptimizationConfig OptimizationConfig(name='gradient', optimizer='Adam', lr=0.001, weight_decay=5e-05, gradient_clip=200.0)

Contracts