File size: 9,990 Bytes
e6708da
 
a3c71e8
0f78592
 
 
a3c71e8
 
 
 
 
 
e6708da
a3c71e8
0f78592
 
 
a3c71e8
 
 
 
 
 
2309f6d
 
 
a3c71e8
 
 
 
 
 
 
 
 
 
 
 
 
0f78592
 
 
 
 
 
 
 
 
 
 
 
a3c71e8
 
 
 
2309f6d
 
 
 
 
 
 
 
a3c71e8
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
2309f6d
 
 
 
a3c71e8
 
 
2309f6d
 
 
 
 
 
 
 
 
 
 
 
 
a3c71e8
 
 
 
2309f6d
 
 
 
 
 
 
 
a3c71e8
2309f6d
a3c71e8
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
2309f6d
a3c71e8
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
6ec8092
a3c71e8
 
2309f6d
a3c71e8
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
---
license: apache-2.0
tags:
  - physiowave
  - physiowave-t
  - foundation-model
  - eeg
  - biosignal
  - self-supervised
  - masked-autoencoder
  - pretrained
library_name: pytorch
---

# PhysioWave-T-EEG: multi-route pretrained EEG encoder

One of the three **PhysioWave-T** encoders (EEG, [ECG](https://huggingface.co/ForeverBlue/ECG), [sEMG](https://huggingface.co/ForeverBlue/EMG)); the [RALF fusion model](https://huggingface.co/ForeverBlue/RALF-SeizeIT2) combines them. Code: <https://github.com/ForeverBlue816/PhysioWave>.

Masked dual-reconstruction pretraining on multi-montage EEG. One shared
transformer serves four electrode layouts through route-specific wavelet
frontends, and channel identity enters as a learned embedding of the channel
NAME rather than of its index.

Trained for 60 epochs
(46080 optimizer updates) on a mixture of public EEG
corpora, 30.3 M parameters.

## Objective

```
0.5*L_spec + 0.5*L_raw + 0.001*L_foldKL
```

Two heads at equal weight. One predicts folded wavelet patches, the other
predicts the preprocessed EEG waveform of the same masked patches. Masking is
applied BEFORE the wavelet frontend (True), at ratio
0.7, and the spec target is per-patch normalised
(True).

## Training curves

Pretraining ran for 60 epochs (46,080 optimizer updates) on 16 GPUs over the seven corpora. Validation follows training on every term to the last epoch, so the model is not over-fitting its 15 M windows; the released checkpoint is the epoch with the lowest validation total loss.

![Training and validation per epoch](training_curves.png)

*Train and validation per epoch: the total loss, the two reconstruction terms (folded wavelet patch, raw waveform) and the masked spectral correlation.*

![Validation loss per route and per corpus](route_convergence.png)

*Validation total loss per route (left) and per corpus (right). All four routes converge; the dense 64-channel route reaches the lowest loss and the 19-channel clinical route (TUEG, the largest corpus) the highest. The dashed line is the route-macro mean, which weights the routes equally.*

## Validation, final epoch

| metric | value | what 1.0 would mean |
|---|---|---|
| total loss | 0.2488 | — |
| spec MSE | 0.3046 | — |
| spec NMSE | 0.305 | predicting zero |
| spec Pearson r | 0.831 | perfect |
| raw SmoothL1 | 0.1930 | — |
| raw NMSE | 0.278 | predicting zero |
| raw Pearson r | 0.847 | perfect |
| route-macro total | 0.2403 | — |

NMSE is dimensionless: MSE(prediction, target) / MSE(0, target), so 1.0 is
what predicting zero scores on that head's own target. The two losses are NOT
comparable to each other -- one is an MSE on normalised wavelet coefficients,
the other a SmoothL1 on z-scored volts. The columns that compare the heads are
the correlations and the NMSEs.

The route-macro row weights the four routes equally. The plain rows weight by
sample count, and validation is dominated by the two largest corpora, so a
model specialising toward the smaller ones moves the two rows in opposite
directions.

## Files

### Which file do you want

| your situation | download |
|---|---|
| fine-tune on any montage | `eeg_c1_encoder.pth` |
| your montage IS one of the four routes below | `pretrain_eeg_c1_moe_best.pth`, then cut the route out (one command, below) |
| continue pretraining, or take the model apart yourself | `pretrain_eeg_c1_moe_best.pth` |

**`pretrain_eeg_c1_moe_best.pth` is the complete model** -- all four wavelet
frontends, both patchers, both reconstruction decoders. Nothing is held back;
only the optimiser state was stripped, which no downstream use needs. Take it
if you want to keep or strip pieces yourself.

`eeg_c1_encoder.pth` is the same weights minus everything route-bound and
pretraining-only. It exists because it cannot be loaded wrong.

| file | what it is |
|---|---|
| `eeg_c1_encoder.pth` | encoder alone: channel embedding, gate, transformer |
| `pretrain_eeg_c1_moe_best.pth` | complete model, lowest validation **total** loss (epoch 57) |
| `pretrain_eeg_c1_moe_final.pth` | complete model, last epoch (epoch 60) |
| `pretrain_eeg_c1_moe_best_spec.pth` | complete model, lowest spec MSE (epoch 59) |
| `pretrain_eeg_c1_moe_best_raw.pth` | complete model, lowest raw SmoothL1 (epoch 57) |
| `metrics_epoch.jsonl` | the full training curve |
| `config_resolved.yaml` | every resolved hyperparameter |

### Why several checkpoints

Each is the minimum of a different validation metric, and they are published
together so a downstream comparison can ask which criterion transfers instead
of assuming the one that names `best.pth`. The four bars are within validation
noise of each other here, which is exactly why the question is worth asking and
why reconstruction loss cannot answer it.

`best.pth` is selected on the TOTAL loss -- the loss that was trained. It is
not the best spec loss, and not the best raw loss.

`pretrain_eeg_c1_moe_best_macro.pth` would have been byte-identical to `pretrain_eeg_c1_moe_best.pth` -- the same epoch won both bars -- so it is not published separately.

### What is in the complete model

| component | parameters | share | |
|---|---|---|---|
| shared transformer | 25.22 M | 83.2% | encoder, always transfers |
| channel-name embedding | 0.02 M | 0.1% | always transfers |
| channel projection | 0.03 M | 0.1% | always transfers |
| wavelet frontends (all 4 routes) | 3.37 M | 11.1% | route-bound; downstream usually builds its own |
| patchers (both rates) | 0.20 M | 0.7% | route-bound |
| spec decoder | 0.72 M | 2.4% | pretraining only |
| raw decoder | 0.72 M | 2.4% | pretraining only |
| **total** | **30.30 M** | 100% | |

The shared transformer is 83% of the model and the four frontends together are 11%, which is why there is no per-route encoder file here: four of them would duplicate that transformer four times to deliver frontends that are a small fraction of it. Cut the one you want out of the complete checkpoint instead.

### What transfers, and what does not

`downstream.py` splits a checkpoint into two sets. The channel embedding, its
gate and the shared transformer **always** load. The wavelet frontend and the
patcher load **only when your montage is the route they were trained for** --
same electrode count, same patch length -- and are skipped otherwise rather
than reshaped.

So a fresh frontend is the normal case, not a fallback. Both downstream tasks
in the source repository build their own.

The four routes the shared transformer was pretrained across:

| route | channels | rate |
|---|---|---|
| `E128_512` | 128 | 512 Hz |
| `E19_256` | 19 | 256 Hz |
| `E32_512` | 32 | 512 Hz |
| `E64_256` | 64 | 256 Hz |

To fine-tune ON one of them exactly, take that route's frontend and patcher out
of the complete checkpoint:

```bash
python scripts/export_eeg_pretrained_encoder.py \
    --checkpoint pretrain_eeg_c1_moe_best.pth --route E64_256 \
    --output encoder_E64_256.pth
```

Omit `--route` and you get `eeg_c1_encoder.pth` again. Both pretraining
decoders are dropped either way: fine-tuning against the pretext objective by
accident is the failure that prevents.

**`best.pth` and the final checkpoint are within validation noise of each
other.** Which one transfers better is a downstream question, not a
reconstruction-loss question -- measure both.

## The channel vocabulary travels with the weights

An embedding row means whatever channel held that id when the row was learned.
Every checkpoint here carries the vocabulary and its hash. Resolving channel
names against a different vocabulary silently trains on relabelled electrodes,
which is why the hash is checked rather than assumed.

## Architecture

```
embed_dim   512      depth 8      heads 8
norm        rmsnorm     ffn swiglu     qk_norm True
wavelet     max_level 3, kernel 16
channel     id embedding, injected at the token site
mask        frequency_guided, ratio 0.7
```

## Loading

`EEGC1Downstream` builds the frontend and patcher for your montage and loads
this file into everything else. Nothing here is route-bound, so the montage,
the window length and the patch length are yours to choose.

```python
import torch
from physiowave.eeg_c1.downstream import EEGC1Downstream

ENC = "eeg_c1_encoder.pth"
ck = torch.load(ENC, map_location="cpu", weights_only=False)
mc = ck["model_config"]                  # the width the weights were trained at
CH = ["Fz", "Cz", "Pz", "Oz", "C3", "C4", "P3", "P4"]   # your electrodes

model = EEGC1Downstream(
    in_channels=len(CH), window_samples=1024, sampling_rate=256.0,
    patch_samples=64, num_classes=2, channel_names=CH,
    embed_dim=mc["embed_dim"], depth=mc["depth"], num_heads=mc["num_heads"],
    channel_embed_dim=mc.get("channel_embed_dim", 64),
    channel_vocab_size=ck["channel_vocab_size"])

report = model.load_pretrained(ENC)
print(model.describe_transfer(report))
# loaded: transformer N, channel encoder N, gate yes, frontend 0, patcher 0

class Meta:                               # or your dataset's ChannelMeta
    channel_names = CH
    channel_mask = None

out = model(torch.randn(2, len(CH), 1024), Meta())   # {"logits": [2, 2]}
```

`frontend 0, patcher 0` is the expected line, not a warning: those are the two
things your montage supplies.

**The frontend you build is not randomly initialised.** Its filters are seeded
from real wavelet families -- `sym4, sym5, db6, sym8, db8` -- and learned from
there during fine-tuning. What a fresh frontend costs you is the adaptation
those filters underwent during pretraining, not the wavelet prior itself.

Channel names, not indices, are what this encoder and your data agree on. An
embedding row means whichever electrode held it when it was learned, so pass
real names and check `channel_vocab_sha256` -- `load_pretrained` refuses a
mismatch rather than silently training on relabelled electrodes.

Code: <https://github.com/ForeverBlue816/PhysioWave>

---
Prepared by `scripts/prepare_hf_release.py` from `pretrain_eeg_c1_moe`
on 2026-09-09.