LangevinFlow

LangevinFlow sequential VAE for binned neural spike counts.

Source

  • Registry name: langevin_flow
  • Model class: ladys.models.langevin_flow.LangevinFlow
  • Config class: ladys.models.langevin_flow.LangevinFlowConfig
  • Source file: src/ladys/models/langevin_flow.py

When to use

Use LangevinFlow as a nonlinear latent dynamics model for raw spike-count sequences. A GRU encoder updates short-range hidden state, latent position and velocity variables evolve through an underdamped Langevin step with a locally coupled oscillator potential, and a one-layer Transformer decoder reads the whole latent sequence into Poisson firing rates.

Training Budget

LangevinFlow needs longer training runs than the quick smoke-test settings used for development. For reproduction-style runs, prefer the YAML experiment budgets: the NLB configs use the released-scale epoch counts, and the Lorenz config is set to a longer default. Short runs such as 20 epochs are useful only for checking that the loss and co-bps move in the right direction.

Assumptions

LangevinFlow expects nonnegative spike counts. On synthetic datasets the readout reconstructs the observed neurons. When built by Experiment on an NLB dataset, output_mode: auto sizes the readout to reconstruct held-in plus held-out training neurons and evaluates the held-out output slice.

Outputs

forward returns natural-space firing rates, concatenated [position, velocity, hidden] latent trajectories, and ELBO diagnostics in extras. loss computes Poisson reconstruction with a scheduled Langevin KL penalty and coordinated-dropout gradient masking.

Configuration

Config for the LangevinFlow sequential VAE.

Field Type Default
name Literal['langevin_flow'] 'langevin_flow'
objective str 'langevin_flow_elbo'
hidden_size int 64
output_neurons Optional[int] None
output_mode Literal['auto', 'heldin', 'heldin_heldout'] 'auto'
fwd_steps int 0
dropout float 0.05
gamma float 0.55
langevin_step float 0.01
potential_groups int 4
potential_kernel_size int 3
transformer_heads int 2
transformer_feedforward int 512
coordinated_dropout_rate float 0.5
kl_weight float 0.1
kl_warmup_epochs int 500
weight_decay_warmup_epochs int 500
velocity_prior_var float 0.1
log_rate_min float -8.0
log_rate_max float 8.0
posterior_logvar_min float math.log(0.0001)
posterior_logvar_max float 5.0
sample_train bool True
sample_eval bool False
prediction_samples int 1
optimization OptimizationConfig OptimizationConfig(name='gradient', optimizer='Adam', lr=0.003, weight_decay=2e-05, gradient_clip=200.0)

Contracts