S-CEReBrO Integration - #18
Open
g12bftd wants to merge 16 commits into
Open
Conversation
Adds the shared contracts a separated encoder/head family needs, without altering anything the five bundled families observe. Every change is additive: new fields default to values that leave existing callers unchanged. biofoundation/core/batch.py SignalBatch gains channel_coords, num_padded_channels and num_padded_timesteps. channel_coords (batch, channels, 2, 3) keeps both electrodes of a channel separate; channel_locations (batch, channels, 3) stays as-is for the families that consume electrode midpoints. The two are documented as peers: a model requires one or the other, and neither is derived from the other. BatchRequirements gains matching flags, all defaulting to False, so require_batch_fields validates the new metadata only for families that ask for it. biofoundation/core/protocols.py New SignalEncoder and PredictionHead typing.Protocol definitions describing the token-embedding boundary. Conformance is structural, so bundled families remain valid unchanged while split families gain a checkable contract. Tensor annotations are guarded by TYPE_CHECKING to keep the package importable without PyTorch. biofoundation/core/checkpoints.py New SafetensorsCheckpointMixin maps the load_pretrained_checkpoint and load_safetensors_checkpoint entry points that run_train.py calls onto a task's own load_from_checkpoint, plus split_state_dict_by_prefix for separating encoder and head weights. biofoundation/model_registry.py ModelSpec gains venue, head_targets and size_variants, all defaulted. The five existing entries record their venue and, where applicable, their size variants. tests Venue coverage is now driven from the registry instead of a hard-coded map, and TASK_FILES is derived by globbing tasks/ so a new task file cannot opt out of the shared batch-adapter contract. New cases cover the paired-electrode and padding requirements, confirm the two geometry representations do not satisfy one another, and pin that the added defaults leave existing BatchRequirements equality unchanged. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01TDAmftg4nxMQim2SZFTAha
Imports the S-CEReBrO implementation with no edits of any kind, so that the
adaptation that follows can be reviewed as a diff against known-good upstream
code. Every file in this commit is byte-identical to its TimeFM counterpart and
can be verified with cmp against the source tree.
Imported:
models/s_cerebro.py windowed alternating-attention encoder
models/modules/{attention,patching,pos_embed}.py
models/model_heads/*.py reconstruction, classification,
regression and sequence heads
criterion/{ce,focal,mse,masked_reconstruction}*.py
tasks/{mae_pretraining,classification_task,regression_task}.py
datasets/{lmdb_dataset,tueg_dataset}.py
data_module/{pretraining,finetuning}_data_module.py
make_datasets/*.py preprocessing for the pre-training
corpus and 11 downstream datasets
tests/model_tests/{test_attention,test_pipeline}.py
Deliberately not imported:
schedulers/cosine.py BioFoundation's CosineLRSchedulerWrapper already has a
compatible constructor signature; the existing one is reused.
callbacks/ run_train.py builds ModelCheckpoint from cfg.model_checkpoint,
so NamedModelCheckpoint has no role here.
config/, run_train.py replaced by BioFoundation-shaped equivalents in the
adaptation commit.
empty __init__.py copying these would turn existing BioFoundation namespace
packages into regular packages.
The upstream unit tests are placed under tests/model_tests/ rather than tests/.
They are pytest-style and import torch, while the fast CI job installs neither.
unittest discovery does not traverse a directory without __init__.py, so they stay
invisible to `unittest discover -s tests` and runnable under pytest. Their contents
are unmodified.
This commit is intentionally inert: nothing references the new files yet. Two
mechanical contract checks fail here by design and are resolved in the next commit,
which is the first to edit any of this code:
- Apache headers are absent from all 41 imported files.
- The imported tasks do not yet call as_signal_batch in their step methods.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01TDAmftg4nxMQim2SZFTAha
First commit to modify any imported file. Changes are confined to the repository's
own conventions; no model, head, criterion or dataset logic is altered.
Licensing
Apache 2.0 headers added to all 41 imported files, each recording provenance from
the S-CEReBrO reference implementation. New models/model_heads/__init__.py carries
the header, matching models/modules/.
Batch contract
The three imported tasks now normalise input through as_signal_batch in every
training, validation and test step, satisfying the shared adapter contract.
Each task also validates its family's BatchRequirements once per step via
require_batch_fields. The requirements are looked up from MODEL_REGISTRY using a
model_family key set by the experiment, so the registry becomes load-bearing at
runtime rather than metadata that only the tests read. Tasks fall back to empty
requirements when no family is named, so they stay usable by future families.
Checkpoints
The tasks mix in SafetensorsCheckpointMixin, exposing the load_pretrained_checkpoint
and load_safetensors_checkpoint entry points that run_train.py calls while keeping
their own load_from_checkpoint as the single implementation. Its partial loader
skips shape-mismatched tensors, so an encoder pre-trained at 64 channels can seed a
22-channel fine-tune.
Hydra configuration
config/model/SCEReBrO_{tiny,small,base}.yaml size variants, LUNA-style
config/model_head/*.yaml new group: reconstruction,
classification, regression, sequence
config/task/{pretrain,finetune,finetune_regression}_task_SCEReBrO.yaml
config/criterion/{masked_reconstruction_loss,ce,focal,mse}_criterion.yaml
config/data_module/{pretrain,finetune}_data_module_SCEReBrO.yaml
config/experiment/SCEReBrO_{pretrain,finetune}.yaml
All follow the house convention: @Package directive on line 1, licence header
after it, contents under a wrapper key. Paths use ${env:DATA_PATH} and
${env:CHECKPOINT_DIR} rather than the upstream data_root defaults.
model_head is deliberately absent from config/defaults.yaml and is instead added by
the S-CEReBrO experiments through their own defaults list, so the composition graph
of the ten existing experiments is untouched.
Experiments are flat, matching the other five families. The fine-tuning corpus is a
dataset_root override rather than a file per dataset, so all eleven prepared
datasets are reachable from one experiment.
The data module receives seed: ${seed} through config because run_train.py
instantiates it without arguments, keeping the train/validation split identical
across ranks.
Dependencies
lmdb and asrpy added to requirements.txt, the only imports in the vendored code not
already satisfied.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01TDAmftg4nxMQim2SZFTAha
Makes S-CEReBrO discoverable through the same metadata and documentation surface as the other five families. biofoundation/model_registry.py New "s-cerebro" entry: EEG, windowed alternating-attention Transformer, encoder target models.s_cerebro.SCerebroEncoder, four head targets, three size variants, and BatchRequirements(channel_coords=True). The key is the case-folded display name, matching every existing entry, and a new test pins that invariant so the convention cannot drift. docs/adr/0001-two-electrode-geometry-representations.md Records why channel_locations and channel_coords both exist as first-class fields, why neither is derived from the other, and what that costs. channel_coords is strictly richer, since a midpoint is recoverable from an electrode pair but not the reverse, so making it canonical was the alternative considered. It was rejected because it would have changed how the five existing families obtain geometry. The consequence is recorded plainly: a dataset prepared for one representation cannot feed a family expecting the other without an explicit conversion. Implicit conversion inside require_batch_fields is ruled out, because a model that silently received midpoints where it expected electrode pairs would train without error on quietly wrong geometry. docs/model/SCEReBrO.md Input assumptions, preprocessing, architecture, the alternating attention schedule, the SSL objective, the three downstream layouts, size variants, and runnable commands for pre-training, fine-tuning, dataset switching, sleep staging, regression and linear probing. README.md, models/README.md, docs/README.md, docs/CITATIONS.md Model-zoo row, family table row, documentation index entry, publication-status row and BibTeX. models/README.md now explains that the repository carries two model shapes and which protocols describe the second. CONTRIBUTING.md Adds the additive-change rule for biofoundation/, the requirement to record expensive-to-reverse decisions as an ADR, guidance on the encoder/head option when adding a model, and how to run the PyTorch-dependent tests under tests/model_tests. Across the branch, eight pre-existing files are modified: the four contract modules, three test modules and requirements.txt. No existing model, task, dataset, criterion, experiment or run_train.py is touched, and nothing is deleted. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01TDAmftg4nxMQim2SZFTAha
Verification with hydra-core and PyTorch actually installed, which was not possible when the previous commits were written. Three real defects surfaced and are fixed. Hydra defaults ordering Both S-CEReBrO experiments failed to compose. Hydra requires additions to the defaults list before overrides, so "- /model_head: ..." must precede the "- override /..." entries rather than follow them. Reordered in both experiments. Head-specific keys in the fine-tuning experiment SCEReBrO_finetune set num_classes and num_patches under model_head, which are MlpClassificationHead parameters. Selecting any other head then failed, because the classification keys were merged onto a head that does not accept them. The block is removed: each head config in config/model_head/ carries its own defaults, and per dataset values are given as model_head.num_classes=... on the command line. Documentation The head, task and criterion swap examples used a leading slash. That form belongs in a defaults list inside a config file; a command-line override takes the group name alone. Corrected, with a note explaining the distinction. arXiv identifier Updated to 2607.27913 in the registry, README and CITATIONS. The Hugging Face URL was already correct. tests/test_hydra_composition.py New test composing every file in config/experiment/, not only those a ModelSpec names. The defaults-ordering defect above produced a valid-looking YAML file that Hydra rejected only at run time, and nothing would have caught it. Experiments that do not compose are recorded in KNOWN_UNCOMPOSABLE_EXPERIMENTS with the reason, and the test fails if one starts composing, so an entry cannot rot. FEMBA_quantized is listed: it predates this branch, fails identically on main, is referenced by no registry entry, and has two independent faults - its defaults list requires scheduler/constant_lr, which does not exist, and its target misspells test_24_femba_full_expand2 as test_24_femba_full_expland2. Left unfixed here to keep this branch free of changes to the existing families. Verified All 10 pre-existing experiments resolve byte-identically to main. LUNA, TinyMyo and PanLUNA: identical parameter counts, state_dict signatures and forward-pass output hashes on main and on this branch, from deterministically filled weights. FEMBA and LuMamba cannot be imported without CUDA-built mamba_ssm, on main or here; their sources and resolved configs are unchanged, so their behaviour is unchanged. S-CEReBrO: 16 upstream unit tests pass, and the pre-training, classification, regression and sequence-classification tasks each complete a training step with a finite loss. require_batch_fields rejects a batch without channel_coords, which confirms the registry lookup is live at runtime. All six commands in docs/model/SCEReBrO.md compose. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01TDAmftg4nxMQim2SZFTAha
Running pre-training and fine-tuning end to end for the first time surfaced one real integration defect, fixed here, and produced the tooling to reproduce the run. make_datasets/make_dummy_scerebro_dataset.py Generates synthetic corpora in the exact on-disk formats the readers expect: pickled dictionaries plus a meta.json key index for LMDBDataset, and fixed-size raw blobs plus a key list for TUEGDataset, with padded channels written as all-zero coordinates so num_padded_channels is recovered as it is from real data. Covers the pre-training union, TUAB, ISRUC and SEED-VIG. A label is written only for corpora that have one. Pre-training samples must carry no label at all, because they are concatenated with TUEG, which has none, and a batch mixing samples with and without the key cannot be collated. datasets/lmdb_dataset.py Open the LMDB environment lazily on first read rather than in the constructor. run_train.py deletes and re-instantiates the data module before the rank-zero test pass, and LMDB refuses to open the same path twice in one process, so every fine-tuning run with final_test enabled aborted after training with "The environment ... is already open in this process". The upstream runner reuses one data module object and never hit this, so the defect only appears in BioFoundation. Lazy opening also stops an open handle being inherited by forked dataloader workers. This mirrors TUEGDataset, which already manages its environment this way. Where no meta.json index exists the key scan now uses a temporary environment that is closed again, so construction still leaves no handle open. This is the only vendored file that differs from upstream beyond its licence header. docs/model/SCEReBrO.md New smoke-test section with the generator invocation and the CPU overrides, and a note that on PyTorch 2.6 and newer the final validation and test passes fail when reloading a checkpoint, because torch.load now defaults to weights_only=True and the tasks store their Hydra configuration in the checkpoint. That affects run_train.py for every family; the workaround is TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD=1. Verified on CPU with generated data Pre-training: 9 corpora concatenated to 216 windows, one epoch, train and val loss logged, checkpoints written. Fine-tuning: encoder loaded from the pre-training checkpoint with 97 of 97 tensors matched and none skipped, two epochs, then validation and test with the full metric set reported. All 27 contract tests pass with and without hydra installed, the 16 vendored unit tests pass, and the 10 pre-existing experiments still resolve identically to main. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01TDAmftg4nxMQim2SZFTAha
…normalisation Testing the Hugging Face checkpoint path for the first time found it was broken. tasks/mae_pretraining.py The task aliased the encoder's mask and pad tokens onto itself. Assigning an nn.Parameter to a second module registers it a second time, so every checkpoint carried both "model.mask_token" and "mask_token" backed by one storage, 101 tensors where the model has 99. safetensors refuses shared storage, so util/ckpt_to_safetensor.py failed with "Some tensors share memory" and no S-CEReBrO checkpoint could be published in the format the releases use. The aliases are gone and both tokens are read through self.model. Upstream never hit this because it distributes .ckpt rather than safetensors. tests/model_tests/test_pipeline.py reads the token through the encoder to match. Normalisation S-CEReBrO uses per-channel min-max scaling to [-1, 1], applied by LMDBDataset via apply_minmax and written into TUEG offline by make_tueg, which uses the same formula. The quantile normalisation configured in config/defaults.yaml belongs to the tasks the bundled families use, and the S-CEReBrO tasks never read input_normalization, so it was already inert. Both experiments now set input_normalization.normalize to False explicitly, so a resolved config cannot be misread as enabling it. No behaviour changes. make_datasets/make_dummy_scerebro_dataset.py Synthetic TUEG windows are now min-max normalised at write time, as make_tueg does. TUEGDataset returns the stored waveform unchanged, so without this the TUEG windows sat on a different scale from the LMDB corpora they are concatenated with. docs/model/SCEReBrO.md Pretrained Weights rewritten to match the other families: snapshot_download from PulpBio/S-CEReBrO, then pretrained_safetensors_path with the local path. Documents pretrained_checkpoint_path for a local .ckpt, the conversion command, how to read the loader's counts, and that a size mismatch is silently skipped rather than raised so the printed counts are what to check. Verified on CPU Checkpoint now has 99 tensors with no shared storage; conversion to safetensors succeeds and yields 97 encoder plus 2 head tensors. Fine-tuning from that safetensors file reports "[load:model] loaded=97 shape_mismatch=0 unexpected=0 total_target=97", trains, validates and tests, with the classification head left freshly initialised. Resolved configs confirm input_normalization.normalize=False and apply_minmax=True for both experiments. 27 contract tests, 16 vendored unit tests, and the 10 pre-existing experiments still resolving identically to main. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01TDAmftg4nxMQim2SZFTAha
Answers three questions about using the published checkpoints, and removes a way the
size label could disagree with the encoder it names.
Channel counts
The variable-montage mechanism transferred intact and is now documented and
demonstrated. Nothing in the state dict depends on num_channels: the temporal
position table is sized by max_timesteps and shared across channels, the channel
embedding is an MLP over 3D electrode coordinates with no per-channel parameters,
and the attention blocks take the channel count only as a reshape argument. Loading
the 64-channel checkpoint into encoders built for 1, 4, 6, 16, 22, 32 and 64
channels gives loaded=97 shape_mismatch=0 unexpected=0 in every case.
docs/model/SCEReBrO.md gains a Channel Counts section separating model.num_channels,
which must equal what the dataset yields, from model.max_channels, which must stay at
the pre-training value or the checkpoint will not match, and noting that the encoder
raises on a mismatch rather than silently reshaping.
model_size
Moved from the experiments into config/model/SCEReBrO_{tiny,small,base}.yaml, so
selecting a group sets the label. It previously sat in the experiment while the
encoder came from the group, so model=SCEReBrO_small without model_size=small
mislabelled the output directory, and now that model_size is interpolated into the
checkpoint path it would have selected the wrong weights.
Two contract tests pin this: a model config declaring model_size must agree with its
own filename, and no experiment may restate a model_size its selected group already
declares. Families that never interpolate it may still omit it entirely, and those
that declare it only in the experiment are unaffected.
Checkpoint location
The fine-tuning experiment gains pretrained_root, defaulting to
${env:CHECKPOINT_DIR}/pretrained/S-CEReBrO, so a checkpoint can be named without an
absolute path:
'pretrained_safetensors_path=${pretrained_root}/SCEReBrO_${model_size}.safetensors'
Switching size then needs one change and the checkpoint follows.
The documented snapshot_download call writes to that directory. ${hydra:runtime.choices.model}
would remove the model_size indirection entirely but does not resolve under
hydra.compose, so it would break the composition tests; the group-declared label
achieves the same guarantee without that cost.
Verified on CPU
Fine-tuning the 64-channel checkpoint onto a 6-channel ISRUC montage through the
interpolated path: loaded=97 shape_mismatch=0 unexpected=0, trains, validates, tests.
All three sizes resolve to their matching checkpoint filename.
29 contract tests, 16 vendored tests, 12 of 12 documented commands compose, and the
10 pre-existing experiments still resolve identically to main.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01TDAmftg4nxMQim2SZFTAha
Fine-tuning with the experiment's own defaults hangs at the first training step on a multi-GPU machine, with no error and no output. The encoder's mask_token and pad_token exist for masked pre-training. A fine-tuning forward pass never substitutes either, so neither receives a gradient. The fine-tuning experiment sets strategy: ddp and find_unused_parameters: False, and DistributedDataParallel's reducer waits for a gradient from every parameter it tracks, so the step never completes. Nothing is printed, which is why this reads as the run doing nothing rather than as a failure. This was not caught earlier because every smoke test forced trainer.strategy=auto and trainer.devices=1, which takes DDP out of the picture entirely. tasks/classification_task.py, tasks/regression_task.py New freeze_pretraining_only_parameters sets requires_grad=False on the encoder's mask and pad tokens, called from both fine-tuning tasks. That removes them from the set DDP tracks, which is cheaper than enabling find_unused_parameters and walking the autograd graph on every step. configure_optimizers already skips parameters that do not require grad, so no parameter group changes. LUNA disables its mask token for classification for the same reason; see models/LUNA.py. tests/model_tests/test_pipeline.py New test asserting that a fine-tuning backward leaves no trainable parameter without a gradient, so a future parameter that the fine-tuning path does not touch fails the suite rather than silently hanging a distributed run. Verified Both fine-tuning tasks now report zero trainable parameters without a gradient, where classification and regression each previously reported model.mask_token and model.pad_token. Fine-tuning from safetensors still trains, validates and tests with loaded=97 shape_mismatch=0 unexpected=0. 29 contract tests, 17 vendored tests, and the 10 pre-existing experiments still resolving identically to main. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01TDAmftg4nxMQim2SZFTAha
…tps://github.com/g12bftd/BioFoundation into claude/biofoundation-scerebrо-integration-a60ylc especially if it merges an updated upstream into a topic branch.
…ides
Switching S-CEReBrO to a different fine-tuning corpus previously meant keeping up to
five overrides consistent by hand: dataset_root, dataset_kind, label_mode,
model.num_channels, and the model_head, task and criterion trio. Getting the trio
partly right is not an error, it is a silently different run, and pairing a
classification head's num_classes with a regression head is a crash.
config/dataset/
New group with one file per prepared corpus: tuab, chb-mit, neonate, physionet-mi,
shu-mi, stew, mumtaz, mental-arithmetic, seed-v, isruc and seed-vig. Each owns its
corpus path, sample layout, label kind, channel count and the matching head, task and
criterion, so the whole combination moves together:
python -u run_train.py +experiment=SCEReBrO_finetune dataset=isruc
Each file sets only the keys its selected head accepts, so seed-vig carries neither
num_classes nor num_patches.
config/experiment/SCEReBrO_finetune.yaml
Gives up ownership of everything the group now holds. This is required rather than
tidy: Hydra applies a config's own values after its defaults list, so a key set in
both places resolves to the experiment's copy and the dataset file is silently
ignored. That is exactly what an earlier probe showed, where a dataset group's head
selection applied but its task, criterion and channel count did not.
tests/test_hydra_composition.py
Two tests. The first composes every dataset option and checks the head, task and
label_mode agree on whether the run is regression, and that a head only receives keys
its constructor accepts. The second fails if the experiment restates anything the
group owns. Both were confirmed non-vacuous by reintroducing the corresponding
mistake; the second initially missed a nested model.num_channels and its pattern was
corrected.
Non-S-CEReBrO models are unaffected. config/defaults.yaml is untouched, so the new
group is reachable only from the S-CEReBrO fine-tuning experiment, and no existing
experiment, model, task, dataset or criterion file is modified anywhere on this branch.
All 10 pre-existing experiments still resolve byte-identically to main.
Verified on CPU
All 11 dataset options compose with the correct head, task and criterion.
Real runs for dataset=tuab, dataset=isruc and dataset=seed-vig each load the
pre-trained encoder with loaded=97 shape_mismatch=0 and complete training,
validation and test, reporting accuracy for the two classification corpora and
rmse, nrmse, r2 and pearson for the regression corpus.
31 contract tests pass with and without hydra installed, 17 vendored unit tests pass,
and all 13 commands in docs/model/SCEReBrO.md compose.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01TDAmftg4nxMQim2SZFTAha
…tegration-a60ylc' into claude/biofoundation-scerebrо-integration-a60ylc
The registry entry for S-CEReBrO now records venue="MICCAI 2026", which the registry-driven citation test checks against docs/CITATIONS.md. The publication-status row and the BibTeX entry are updated to match: an inproceedings entry noting the proceedings are forthcoming, following the LuMamba and PanLUNA pattern, and the closing note now lists MICCAI alongside EUSIPCO and AICAS. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01TDAmftg4nxMQim2SZFTAha
Every file carrying "Author: Glenn Anta Bucagu" now credits only "BioFoundation Contributors", matching the two experiment configs already corrected on this branch. 69 files, one deleted line each; no other content changed. No file that exists on main carried the line, so nothing pre-existing is touched. The line had been added to files created on this branch, including ones in shared directories such as criterion/, tasks/, models/modules/, datasets/ and make_datasets/, where it read as a claim over generic infrastructure. The BibTeX entry in docs/CITATIONS.md keeps "Bucagu, Glenn Anta" in its author list. That is the paper's authorship, not a file-header attribution. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01TDAmftg4nxMQim2SZFTAha
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds S-CEReBrO, a compact EEG encoder using windowed alternating attention, as the
sixth model family. It is the first family here that separates the encoder from its
output layer: the encoder emits token embeddings, and a prediction head turns those into
a reconstruction, class logits, or a scalar. One pre-trained encoder therefore serves
every downstream task without being rebuilt.
The implementation is imported from the S-CEReBrO reference repository. It was landed as
a verbatim commit followed by a separate adaptation commit, so the model code can be
diffed against known-good upstream and the BioFoundation-side changes reviewed on their
own.
Shared contracts (
biofoundation/)All changes are additive; new fields default so that existing callers observe no
difference.
core/batch.py—SignalBatchgainschannel_coords,num_padded_channels,num_padded_timesteps.channel_coords(B, C, 2, 3)keeps both electrodes of achannel separate;
channel_locations(B, C, 3)is unchanged for the families thatconsume midpoints. The two are peers — a model requires one or the other, and
neither is derived from the other. Rationale in
docs/adr/0001.core/protocols.py(new) —SignalEncoderandPredictionHeadastyping.Protocol, so conformance is structural and the bundled families need nochanges.
core/checkpoints.py(new) —SafetensorsCheckpointMixinmaps the twocheckpoint entry points
run_train.pycalls onto a task's ownload_from_checkpoint.model_registry.py—ModelSpecgainsvenue,head_targets,size_variants,all defaulted. The registry is now enforced at runtime: tasks look up their family's
BatchRequirementsvia amodel_familykey and validate every batch.Configuration
config/model_head/group (opt-in; absent fromconfig/defaults.yaml).config/dataset/group with one file per fine-tuning corpus. Each owns itspath, sample layout, label kind, channel count, and the matching head, task and
criterion, so switching corpus is one override:
python -u run_train.py +experiment=SCEReBrO_finetune dataset=isrucmodel_sizeis declared by theconfig/modelgroup rather than the experiment, sothe label always matches the encoder and can be interpolated into checkpoint paths.
Compatibility with existing models
config/defaults.yamlis untouched; the new groups are reachable only from theS-CEReBrO experiments.
modified files are the four contract modules, four test modules,
requirements.txt,and five docs.
main.mainon parameter count,state_dictsignature, and forward-pass output hash from deterministically filledweights. FEMBA and LuMamba need a CUDA-built
mamba_ssmand cannot be imported onCPU, on
mainor here; their sources and resolved configs are unchanged.