NDT

Transformer encoder trained with a masked Poisson spike objective.

Source

  • Registry name: ndt
  • Model class: ladys.models.ndt.NDT
  • Config class: ladys.models.ndt.NDTConfig
  • Source file: src/ladys/models/ndt.py

When to use

Use NeuralDataTransformer (NDT) as a self-supervised sequence baseline for binned spike counts. This adapter follows the Lorenz NDT configuration from snel-repo/neural-data-transformers: per-neuron spike-count embeddings, optional local temporal attention, learnable positions, pre-norm transformer layers, and a Poisson decoder trained on randomly masked observations.

Assumptions

NeuralDataTransformer (NDT) expects raw nonnegative spike counts. Dataset-level smoothing should be disabled for this model. The model returns natural-space rates for metrics; internally, the default decoder predicts log rates for stable Poisson loss.

Outputs

forward returns nonnegative rate predictions, transformer factor trajectories in latents, and masking diagnostics in extras. During training, loss uses only masked entries. During evaluation, it uses all entries so validation reports a full reconstruction objective.

Configuration

Config for the masked-count NeuralDataTransformer (NDT) adapter.

Field Type Default
name Literal['ndt'] 'ndt'
objective str 'masked_poisson_nll'
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 int 2
learnable_position bool True
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 True
mask_ratio float 0.25
mask_mode Literal['full', 'timestep', 'neuron', 'timestep_only'] 'timestep'
mask_token_ratio float 1.0
mask_random_ratio float 0.5
mask_max_span int 1
mask_span_expand_prob float 0.0
mask_span_ramp_start int 0
mask_span_ramp_end int 0
use_zero_mask bool True
topk_loss_fraction float 1.0
nlb_decoder Literal['direct', 'latents'] 'direct'
optimization OptimizationConfig OptimizationConfig(name='gradient', optimizer='Adam', lr=0.001, weight_decay=0.0, gradient_clip=200.0)

Contracts