Shubhronyu/LFPformer-large-MaskPred-35k

LFPformer Large — MaskPred

Masked latent prediction without a spatial auxiliary objective.

This is the 35,000-step Large checkpoint from the NeurIPS 2026 LFPformer matched decoder-capacity comparison. It has 24 Transformer layers, hidden width 1,024, 16 attention heads, and about 317.5 million backbone parameters including the pretraining components. It was initialized from facebook/wav2vec2-large and trained on IBL full-recording LFP windows. With effective global batch 512, this checkpoint represents 17,920,000 window presentations, not that many unique recordings. It is a fixed-budget comparison checkpoint, not the latest training endpoint or a claim of universal best performance.

Extract LFP embeddings

Supply already preprocessed single-channel LFP windows of 3 seconds at 1,250 Hz as windows.npz, with a finite signals array of shape (N, 3750). Preserve the training signal units and dataset-specific preprocessing. The helper resamples to 16 kHz without per-window normalization, then averages the final hidden states to give 1,024-dimensional embeddings.

pip install -r requirements.txt
python extract_embeddings.py --model . --input windows.npz --output embeddings.npz --device cpu

Install matching PyTorch and torchaudio builds for your CPU/CUDA environment. Use --device cuda for GPU extraction. The weights load through standard Wav2Vec2ForPreTraining.from_pretrained(...); no custom remote code is required. This is an LFP representation model, not an automatic speech recognizer. Do not apply the default audio processor's waveform normalization to LFP.

Files and compatibility

• model.safetensors, config.json: full standard wav2vec2 pretraining model.
• auxiliary.safetensors: the original wrapper's auxiliary head and buffers, retained losslessly. The auxiliary XYZ head had xyz_lambda=0 and is not a trained localization decoder.
• training_config.yaml: embedded source settings with machine paths removed.
• lfp_preprocessing.json, provenance.json, SHA256SUMS: preprocessing, checkpoint identity, validation and file integrity.
• restore_lfpformer_checkpoint.py: reconstructs the exact original wrapper tensor dictionary in a local .pt file if required by research code.

The existing LFPformer benchmark embedding loader hardcodes 256-dimensional codebook projections; these Large checkpoints use 768. Use the supplied native Transformers helper, or the compatible LFPformer pretraining backbone with the supplied configuration and strict loading.

These files contain no fitted downstream localization/classification decoder, decoder whitening, optimizer, scheduler, or training data. Localization requires a compatible trained readout and its training normalization. Paper results depend on the decoder, cohort, aggregation and postprocessing; encoder weights alone do not reproduce a complete-system metric.

Training and evaluation scope

The spatial auxiliary objectives are implemented in the LFPformer research code, not standard Transformers' pretraining loss. The embedded max_global_step may retain an earlier gate limit, and its batch size may be per-rank; consult provenance.json for the checkpoint's actual step and effective batch. This snapshot is not a complete optimizer-resume package.

Downstream anatomical labels are used when training supervised localization readouts; they are not needed to extract embeddings. IBL refinement uses inverse-variance channel pooling followed by unweighted Kabsch fitting of known electrode geometry to predictions. Allen PCA-line projection is a separate evaluation protocol.

Validation and source

All 417 original state tensors are preserved exactly. Strict loading and a finite (1, 1024) embedding smoke check passed; native and original embeddings were exactly equal on the synthetic test input. This checks export mechanics, not scientific performance. See provenance.json for versions and source SHA-256.

Research code: https://github.com/neuroinfolab/LFPformer (access may be restricted). Code provenance is pinned to commit 139b8d90054aa99608d77ba89832050b6eb3cf3b. The upstream initialization retains its applicable terms. A separate license for these fine-tuned weights has not been specified in this package.