Overview
Data & figure notes
One-second inference example using the full recording montage (128 contacts), not the benchmark’s at-most-120-contact subset. Each colored cell is a token: one contact, one band, one timepoint. Its 7 / 6 / 7 input bins are projected to 384 features. Shafts A–M are display aliases for the real shafts, in the order listed in the linked data; contact numbers are unchanged. The layout uses real contact metadata; cell colors identify bands, not measured amplitudes. Data & provenance ↗
Frontend
Voltage → normalized STFT clips
STFT: 32 frames/s in every band.
Windows, normalization & recording source
Slow: 500 ms / 2–14 Hz; mid: 125 ms / 16–56 Hz; fast: 62.5 ms / 64–160 Hz. Outlines show Hann-window support, drawn to scale. Only retained centers move here; all three STFTs are computed every 31.25 ms before subsampling. Offline input recipe: shaft mean-reference → notch/high-pass filtering → magnitude STFT → session median/MAD → normalized clip. Statistics are fitted per contact/bin over session time; z = (x − median) / (1.4826 × MAD), with a floor for near-constant bins. This is session preprocessing, not a causal streaming frontend. This figure shows raw BTB voltage (subject 1, trial 0, F3aOFa2, 30–31 s), centered/scaled vertically; marks show retained frame times, not computed spectra. Samples & provenance ↗ · Brain Treebank, CC BY 4.0
Clip & contact layout
Keep following Shaft A: 13 retained contacts, 1 second. Clinical contact numbers remain unchanged, including gaps. Batch axis omitted unless shown.
Referencing and filtering
250 # 3. Subtract each shaft's mean, then apply harmonic notches and a high-pass filter.251 model_voltage = _filter(_reference(x.astype(np.float64), sidecar), sidecar.labels).astype(252 np.float32253 )254 waveform = torch.from_numpy(model_voltage)Reference and filter implementations
63def _reference(data, sidecar):64 data = np.array(data, copy=True)65 for array in range(sidecar.n_arrays):66 ids = np.flatnonzero(sidecar.array_id.numpy() == array)67 data[ids] -= data[ids].mean(axis=0, keepdims=True)68 return data57 freqs = np.arange(60.0, min(SAMPLE_RATE / 2, 301), 60.0)58 raw.notch_filter(freqs, phase="zero", filter_length="auto", verbose=False)59 raw.filter(0.5, None, verbose=False)60 return raw.get_data()Three STFT windows
31SAMPLE_RATE = 204832FRAME_RATE = 3233# name, FFT length, first bin, last bin (inclusive). Hz = bin * 2048 / FFT length.34# Bin centers (inclusive): slow 2-14 Hz by 2; mid 16-56 by 8; fast 64-160 by 16.35# Window lengths: 500/125/62.5 ms; all three use hop=64 samples (31.25 ms).36BANDS = (("slow", 1024, 1, 7), ("mid", 256, 2, 7), ("fast", 128, 4, 10))28 win = torch.hann_window(nperseg, device=waveform.device)35 spec = torch.stft(36 wf,37 n_fft=nperseg,38 hop_length=hop_length,39 win_length=nperseg,40 window=win,41 return_complex=True,42 normalized=False,43 center=True,44 )45 band = spec[..., k0 : k1 + 1, :] # (..., n_bins, T) complex51 mag = band.abs() # (..., n_bins, T)Session STFT and normalization
256 # 4. Voltage (C,samples) -> magnitude STFT (C,F_b,session_frames), no log.257 # Slow: 2-14 Hz, 7 bins, 500 ms window; mid: 16-56 Hz, 6 bins, 125 ms;258 # fast: 64-160 Hz, 7 bins, 62.5 ms. Ranges are bin centers, inclusive.259 # All bands use a 31.25 ms hop (32 frames/s); the stem subsamples later.260 for _, fft, lo, hi in BANDS:261 band = _single_stft_raw_view_chunked(262 waveform,263 sample_rate=SAMPLE_RATE,264 nperseg=fft,265 hop_length=64,266 k0=lo,267 k1=hi,268 log_eps=1e-6,269 chunk_frames=chunk_frames,270 )271 # 5. Fit once over session time, separately for each contact/frequency bin:272 # median and sigma = 1.4826 * median(abs(band - median)), both (C,F_b,1).273 # Keep raw magnitudes + these statistics; window() applies normalization.274 norm = SessionRobustZNormalizer(sigma_floor=1e-6).fit(band)275 bands.append(band)276 normalizers.append(norm)44 # Time is the last axis: (C,F,T_session) -> statistics (C,F,1).45 # Each contact/bin gets its own median and scaled median absolute deviation.46 self.median = x.median(dim=-1, keepdim=True).values47 self.sigma = SCALE_TO_SIGMA * (x - self.median).abs().median(dim=-1, keepdim=True).values48 self._valid_bin_mask = valid_bin_mask49 return selfZoom in: selecting a normalized clip
129 frames = round(duration_s * FRAME_RATE)130 if frames < 8 or frames % 8 or not math.isclose(frames / FRAME_RATE, duration_s):131 raise ValueError("duration must be a positive multiple of 0.25 seconds")132 start = round(start_s * FRAME_RATE)140 clip_bands = []141 for band, norm in zip(self.bands, self.normalizers):142 clip = band[..., start : start + frames] # select cached STFT frames143 normalized = norm.transform(clip) # use stored session statistics144 batched = normalized.unsqueeze(0) # add the one-clip batch axis145 clip_bands.append(batched.to(device))146 return clip_bandsApply stored session statistics
51 def transform(self, x: torch.Tensor) -> torch.Tensor:52 if self.median is None or self.sigma is None:53 raise RuntimeError("SessionRobustZNormalizer.transform called before fit()")54 x = x if x.dtype in (torch.float16, torch.float32, torch.float64) else x.float()55 # Broadcast stored session statistics across the clip's time axis.56 # Near-constant bins are set to zero below; clipping is done later in the stem.57 z = (x - self.median) / self.sigma.clamp(min=self.sigma_floor)58 z = torch.where(self.sigma >= self.sigma_floor, z, torch.zeros_like(z))59 if self._valid_bin_mask is not None:60 z = torch.where(self._valid_bin_mask.to(dtype=torch.bool), z, torch.zeros_like(z))61 return zTokenization
7 / 6 / 7 frequency bins → 384 features per token
Columns = timepoints · rows = frequency bins
Shaft A: 13 contacts × 52 tokens = 676 tokens.
Shapes & coordinates
All frontend bands initially use a 32 Hz clock. For one second: 4 + 16 + 32 = 52 tokens/contact. For Shaft A: 13 × 52 = 676 tokens.
Retain every eighth, second, or first frame
62 def decimate(x: Tensor, stride: int) -> Tensor:63 """Strided time slice ``x[..., ::stride]`` — the band's own-rate frames. Requires64 the 32 Hz length to be an exact multiple of ``stride`` (64 % {1,2,8} == 0)."""65 t32 = x.shape[-1]66 if t32 % stride != 0:67 raise ValueError(f"32 Hz length {t32} not a multiple of stride {stride}")68 return x[..., ::stride]Subsample, project, add band identity
24PER_BAND_SPECS: tuple[tuple[int, int], ...] = ((7, 8), (6, 2), (7, 1))25INPUT_CLIP_Z = (15.0, 15.0, 20.0) # Slow, Mid, Fast: published Guard 3 caps.79 for b, (x, (n_bins, stride), proj) in enumerate(80 zip(band_inputs, self.specs, self.projs)81 ):86 # 1. Subsample time and clip normalized inputs.87 xd = self.decimate(x, stride) # (..., F_b, T_b)88 xd = clip_band(xd, b) # same shape; cap extreme normalized values89 # 2. Put frequency last, then project F_b -> d and add band identity.90 xd = xd.transpose(-1, -2) # (..., T_b, F_b)91 tok = proj(xd) # (..., T_b, F_b) -> (..., T_b, d)92 tok = tok + self.band_type_emb[b] # add the same band vector to each token93 t_b = xd.shape[-2]94 # 3. Keep positions on the shared clock: slow 0,8,...; mid 0,2,...; fast 0,1,...95 pos = torch.arange(t_b, device=x.device, dtype=torch.long) * stride # (T_b,)96 tokens.append(tok)97 positions.append(pos)Projection definitions & input clipping
28def clip_band(x: Tensor, band: int) -> Tensor:29 """Bound normalized model inputs without modifying the caller's tensor."""30 cap = INPUT_CLIP_Z[band]31 return x.clamp(-cap, cap)54 self.specs = tuple((int(nb), int(st)) for nb, st in bands)55 self.projs = nn.ModuleList(nn.Linear(nb, d_model) for nb, _ in self.specs)56 # One learned width-d vector per band, broadcast over contacts and time.57 self.band_type_emb = nn.Parameter(torch.empty(len(self.specs), d_model))Token coordinates
72 S, max_c = geom.n_arrays, geom.max_c73 lengths = band_token_counts(n_time) # (T_slow, T_mid, T_hga)74 k_full = int(sum(lengths))7576 # canonical (array-major, slot) contact order.77 rows = torch.arange(S, device=device)[:, None].expand(S, max_c)[geom.valid] # (N,)78 canon = geom.gather_idx[geom.valid] # (N,) contact index into N79 depth_canon = geom.depth[geom.valid] # (N,) clinical depth80 n = int(canon.shape[0])8182 # the per-contact (band, bandpos, time_pos) block — identical for every contact.83 band_blk = torch.cat(84 [torch.full((L,), b, dtype=torch.long, device=device) for b, L in enumerate(lengths)]85 ) # (k_full,)86 pos_blk = torch.cat(87 [torch.arange(L, device=device) for L in lengths]88 ) # (k_full,) band-token index89 lat_blk = torch.cat(90 [torch.arange(L, device=device) * st for L, st in zip(lengths, BAND_STRIDES)]91 ) # (k_full,) lattice position = j * stride9293 # tile over the N contacts (contact-major ⇒ array-contiguous).94 depth = depth_canon.repeat_interleave(k_full)95 contact = canon.repeat_interleave(k_full)96 array = rows.repeat_interleave(k_full)97 band = band_blk.repeat(n)98 bandpos = pos_blk.repeat(n)99 time_pos = lat_blk.repeat(n)Pack contacts in array order
133 # 1. Read the array-grouped contact order from the metadata plan.134 canon = grid.contact.reshape(n, grid.k_full)[:, 0] # (n,) contact index, canonical order135 blocks: list[Tensor] = []136 for band, (tok, length) in enumerate(zip(band_tokens, grid.band_lengths)):137 if tok.shape[2] != length:138 raise ValueError(139 f"band {band} has {tok.shape[2]} tokens, grid expects {length}"140 )141 blocks.append(tok[:, canon]) # (B, n, T_b, d) reindexed to canonical contacts142 # 2. Concatenate each contact's bands, then flatten contacts and token positions.143 contact_tokens = torch.cat(blocks, dim=2) # (B,n,k_full,d): slow, mid, fast per contact144 flat_tokens = contact_tokens.reshape(B, grid.total, d) # (B,n*k_full,d)145 return flat_tokensSpatial encodings
Anatomical region + relative position
Person A
Person B
Identify the region in each person. Reuse its learned embedding.
Where this intuition came from · μECoG
Original trial labels retained. μECoG context, separate from the MAPA shaft example.
Figure notes · native anatomy and local geometry
The two cortical surfaces are actual Cogan participants, reconstructed in native space with FreeSurfer. For comparable left-lateral views, each surface is centered and rigidly rotated to level its area-weighted principal axis in the display plane. This display adjustment preserves shape and parcel boundaries; it is not AC-PC registration or a template warp. Colors come from each participant's saved DKT annotations. These participants are anatomical examples, not MAPA's pretraining or held-out cohort. Implant coverage is not depicted.
The shaft is a separate MAPA example: Brain Treebank subject 2, shaft LT1bIb, session 4. Retained contacts 1–3 are in left insula; contacts 4–10 are in left STG. Indices and region assignments are real. The movable activity pattern, including grey context bars around the two highlighted contacts, is illustrative, not recorded activity; context bars are clipped at the shaft endpoints; the shaft line and parcel box are schematic. RoPE uses clinical contact indices, not millimeter distances. Region identity supplies a shared anatomical cue; relative position distinguishes contacts inside an array without asserting cross-person subparcel correspondence. Signal features, time and band are held fixed while both contact positions shift together within STG. The dial shows the first actual contact RoPE frequency (1 radian per contact) and identical unit starting Q/K vectors. Its highlighted angle is always 1 radian, so the illustrated dot product stays cos(1). Other feature pairs have their own frequencies; their relative rotations also remain unchanged under a common shift. The contact-index differences and relative RoPE rotations are unchanged; this does not assert that real neural responses or the full encoder are translation invariant.
This is the design rationale. These anatomy images alone do not establish functional equivalence of parcels or superiority to template-coordinate conditioning.
Encoding details · region lookup & RoPE
Real contact indices; schematic embedding vector. Data & provenance ↗
Illustration: identical starting Q/K, one fixed frequency. Rotate Q/K; leave V unchanged.
Region: once per stack. RoPE: every attention block. Encoder + decoder.
Exact coordinate convention
The region table has 74 atlas labels plus an unknown row. Contact position is the integer from the clinical label, with gaps preserved—not millimeters or a whole-brain coordinate. Time is on the shared 32 Hz lattice. The released RoPE bases are 8 (contact) and 64 (time). For one frequency, R(p)q · R(r)k = qᵀR(r−p)k. A common coordinate shift preserves this dot product for fixed Q/K content; this is not a claim that the entire model ignores anatomy. Paper: positional encoding and region appendix ↗
Add the shared region embedding
113 # 1. Add the atlas-region embedding to each token.114 region_bm = region_packed[None, :].expand(b_size, m)115 x_flat = x_flat + self.region_embed(region_bm).to(x_flat.dtype)Region lookup & layer definition
130 grid = build_r4_grid(geom, n_time=int(n_time))131 region_id = sidecar.region_id.to(device)132 contact_order = grid.contact.reshape(-1, grid.k_full)[:, 0]133 return SessionGrid(134 grid=grid,135 region_packed=region_id[grid.contact],136 contact_order=contact_order.cpu(),137 region_of_contact=region_id[contact_order].cpu(),138 k_full=int(grid.k_full),139 n_time=int(n_time),140 )122 embed = nn.Embedding(n_regions, d_model)123 nn.init.trunc_normal_(embed.weight, std=init_std)124 self.embed = embed126 def forward(self, region_id: Tensor) -> Tensor:127 """region_id: (..., seq) long → (..., seq, d_model)."""128 if self.embed is None:129 return region_id.new_zeros((*region_id.shape, self.d_model), dtype=torch.float32)130 return self.embed(region_id)Fixed frequencies for the two axes
53 base_index: float = 8.0,54 base_time: float = 64.0,65 pairs = head_dim // 4 # head_dim=64: 16 pairs for contact + 16 for time66 idx_freq = 1.0 / (base_index ** (torch.arange(pairs).float() / pairs))67 # Disabling space makes contact rotations identity; time rotations remain.68 if not self.space:69 idx_freq = torch.zeros_like(idx_freq)70 t_freq = 1.0 / (base_time ** (torch.arange(pairs).float() / pairs))71 self.register_buffer("idx_freq", idx_freq, persistent=False)72 self.register_buffer("t_freq", t_freq, persistent=False)Turn coordinates into rotation angles
82 # angle = coordinate * frequency (radians); frequencies are fixed buffers.83 ang_idx = idx[..., None].float() * self.idx_freq # (..., seq, pairs)84 ang_t = t[..., None].float() * self.t_freq85 ang = torch.cat([ang_idx, ang_t], dim=-1) # (..., seq, head_dim/2)86 cos = ang.cos().repeat_interleave(2, dim=-1) # (..., seq, head_dim)87 sin = ang.sin().repeat_interleave(2, dim=-1)88 return cos, sinRotate Q and K
40def _rotate_half(x: Tensor) -> Tensor:41 # Each adjacent feature pair (a,b) becomes (-b,a), a 90-degree rotation.42 x_paired = x.unflatten(-1, (-1, 2)) # (..., hd/2, 2)43 return torch.stack([-x_paired[..., 1], x_paired[..., 0]], dim=-1).flatten(-2)96 cos = cos.unsqueeze(-3) # (..., 1, seq, head_dim) → broadcast over heads97 sin = sin.unsqueeze(-3)98 q_out = q * cos + _rotate_half(q) * sin99 k_out = k * cos + _rotate_half(k) * sin100 return q_out, k_outTransformer
Shaft A · 13 contacts × 52 tokens · 12 blocks
Inside one block
One diagonal square = one shaft, one clip. All contacts, bands, and times. Two blocks shown.
Attention details
Q/K use fixed contact/time RoPE rotations; V is unchanged. Learned anatomical region embeddings are added once before the stack. No cross-array or cross-clip attention. All arrays use the same encoder weights. Inference uses the raw layer-12 tap (384 features/token), not the four-tap pretraining fusion.
Start with the two residual updates
115 # Block overview: keep x, learn an attention update, then learn an MLP update.116 # LayerNorm acts across features per token; both additions preserve (M,d).117 normalized = self.norm1(x_flat)118 attention_update = self._attn_flat(119 normalized, depth, time_pos, cu_seqlens, cu_seqlens_drop,120 max_seqlen, rope_cs=rope_cs,121 )122 x_flat = x_flat + attention_update123 # Normalize -> Linear(d,4d) -> GELU -> Linear(4d,d) -> add to x.124 normalized = self.norm2(x_flat)125 mlp_update = self.mlp(normalized)126 x_flat = x_flat + mlp_update127 return x_flatLayerNorm definitions
60 self.norm1 = nn.LayerNorm(d_model, eps=LN_EPS)61 self.norm2 = nn.LayerNorm(d_model, eps=LN_EPS)Zoom in: QKV, rotations, attention
82 # 1. Project Q, K, V: (M,384) -> (M,1152) -> (M,3,6,64).83 qkv = self.qkv(x_flat) # (M,3*d): three learned projections in one linear layer84 qkv = qkv.reshape(m, 3, self.n_heads, self.head_dim)85 q, k, v = qkv.unbind(dim=1) # (M, H, hd)86 # 2. Rotate Q/K. RoPE expects heads before tokens; V is unchanged.87 qh, kh = q.transpose(0, 1), k.transpose(0, 1) # (H, M, hd)88 if rope_cs is None:89 qh, kh = self.rope(qh, kh, depth, time_pos)90 else:91 qh, kh = self.rope.rotate(qh, kh, rope_cs[0], rope_cs[1])92 q, k = qh.transpose(0, 1), kh.transpose(0, 1) # (M, H, hd)93 q, k = q.to(v.dtype), k.to(v.dtype) # align to v for SDPA (no-op under fp32)94 # 3. Attend within each array/clip block; output remains (M,H,hd).95 # Per head: softmax(Q @ K.T / sqrt(hd), over keys) @ V; see varlen.py.96 if q.is_cuda:97 ctx = gpu_block_diag(q, k, v, cu_seqlens_drop, max_seqlen) # (M, H, hd)98 else:99 ctx = _reference_block_diag(q, k, v, cu_seqlens) # (M, H, hd)100 # 4. Concatenate heads, then mix them with a learned d -> d projection.101 ctx = ctx.reshape(m, d) # (M,H,hd) -> (M,d)102 return self.out(ctx)Projection definitions
54 self.n_heads = n_heads55 self.head_dim = d_model // n_heads56 # Released model: d=384, six heads of width 64; no separate Q/K norm.57 self.qkv = nn.Linear(d_model, 3 * d_model, bias=True)58 self.out = nn.Linear(d_model, d_model, bias=True)Attention within each array
1. Build boundaries: [13, 8] contacts × 52 tokens → lengths [676, 416] → offsets [0, 676, 1092].
103 n_per_array = geom.valid.sum(dim=1).to(torch.long) # (S,)104 seg = (n_per_array * k_full).to(torch.int32) # (S,)105 cu = torch.zeros(S + 1, dtype=torch.int32, device=device)106 cu[1:] = seg.cumsum(0).to(torch.int32)2. GPU: those offsets separate Q/K/V into independent sequences. Attention stays inside each sequence; no dense mask is built.
74 def _nt(x: Tensor) -> Tensor:75 return torch.nested.nested_tensor_from_jagged(76 x, offsets=cu_seqlens_drop, min_seqlen=1, max_seqlen=max_seqlen77 ).transpose(1, 2)7879 # Same within-block attention, without allocating a full total-by-total mask.80 ctx = F.scaled_dot_product_attention(_nt(q), _nt(k), _nt(v))81 return ctx.transpose(1, 2).values()Offset each clip and remove empty segments
117 # 2. Give every clip its own array blocks before flattening the batch.118 # E.g. [0,1040,2080] per clip -> [0,1040,2080,3120,4160] for two clips.119 offsets = torch.arange(b_size, device=device, dtype=torch.int64) * m120 array_starts = grid.cu_seqlens[:-1].to(torch.int64)121 batch_starts = array_starts[None, :] + offsets[:, None] # (B,n_arrays)122 final_end = torch.tensor([b_size * m], dtype=torch.int64, device=device)123 cu_b = torch.cat([batch_starts.reshape(-1), final_end])124 keep = torch.ones_like(cu_b, dtype=torch.bool)125 keep[1:] = cu_b[1:] != cu_b[:-1] # drop empty arrays and boundary coincidences126 cu_drop_b = cu_b[keep]127 cu_b = cu_b.to(grid.cu_seqlens.dtype)CPU reference: explicit block-diagonal mask
35 lengths = (cu_seqlens[1:] - cu_seqlens[:-1]).to(torch.long)36 seg = torch.arange(lengths.shape[0], device=cu_seqlens.device)37 return seg.repeat_interleave(lengths)45 total = q.shape[0]46 seg = _segment_ids(cu_seqlens, total)47 same = seg[:, None] == seg[None, :]48 bias = torch.where(same, 0.0, NEG_INF_MASK).to(q.dtype)49 # Per head/block: softmax(Q @ K.T / sqrt(head_dim), dim=-1) @ V.50 # Rows are queries; softmax normalizes over key tokens.51 ctx = F.scaled_dot_product_attention(52 q.transpose(0, 1)[None], k.transpose(0, 1)[None], v.transpose(0, 1)[None],53 attn_mask=bias[None, None],54 )55 return ctx[0].transpose(0, 1).contiguous()Zoom in: the token-wise MLP
37 # Apply the same feature transformation independently to every token.38 # Shapes below use the released encoder width d=384.39 hidden = self.fc1(x) # (M,384) -> (M,1536)40 hidden = self.act(hidden) # GELU; shape unchanged41 output = self.fc2(hidden) # (M,1536) -> (M,384)42 return outputMLP layer definitions
31 hidden = d_model * mlp_ratio32 self.fc1 = nn.Linear(d_model, hidden) # expand features: 384 -> 153633 self.fc2 = nn.Linear(hidden, d_model) # restore width: 1536 -> 38434 self.act = nn.Open code: Zoom in: the token-wise MLP GELU()Twelve blocks; raw layer-12 features
31DEPTH = 12 # blocks; not a knob139 for i, blk in enumerate(self.blocks):140 xf = blk.forward_flat(141 xf, depth_b, time_b, cu_b, cu_drop_b, grid.max_seqlen, rope_cs=rope_cs,142 )143 # Inference taps keep width d; the main multi-level output has width 4d.144 if (i + 1) in tap_blocks:145 taps[i + 1] = xf.reshape(b_size, m, d)Block construction
62 self.blocks = nn.ModuleList(63 [L1Block(d_model, n_heads, space_rope=space_rope) for _ in range(DEPTH)]64 )Readout
One-second trials · 4 + 16 + 32 = 52 tokens/contact · encoder frozen
C = total retained contacts. Keep contacts separate. Cross-session:
Within-session · train on one half, evaluate on the otherCross-session · train on one session, evaluate on another
One bar = one retained contact. All seven shared regions are right-hemisphere; hover for full names. Brain Treebank ↗
Anchor · fit scaling and ridge weightsTarget validation ·
Evaluation protocol
Neuroprobe trials span 0–1 s after word onset. Keep band and time coordinates; flatten rather than average them. One ridge is fit per binary task. Cross-session intersects contact labels when sessions retain different contacts. Cross-subject averages contacts by atlas region across all arrays, intersects the anchor/target regions (3–7 in the benchmark), and excludes the unknown label. The figure uses the actual Neuroprobe-Lite montage for anchor subject 2, session 4 and target subject 1, session 1: seven shared regions after session-specific exclusions.
Ridge weights and feature scaling use training trials only. Select λ by validation AUROC from 25 log-spaced multipliers (10⁻⁴–10⁴), scaled by the mean squared training-feature norm; ties choose the smallest λ. Cross-subject validation uses target labels; the separate target test split is used for reporting. Within-session and cross-session evaluate both directions. Report mean AUROC across sessions for each task, then macro AUROC over 15 tasks.
Matched control: frontend → readout versus frontend → frozen encoder → readout. Same trials, splits, preprocessing and ridge recipe; the baseline uses raw clipped/subsampled bins, not learned stem tokens. Only pretraining learns encoder weights; the decoder is discarded for evaluation. Results and spatial ablations on the project page ↗
Collect array features
247 for s in range(0, n, batch_size):248 e = min(s + batch_size, n)249 bb = [b[s:e].to(device) for b in bands]250 Bb = e - s251 with torch.autocast(device_type=device.type, dtype=torch.bfloat16,252 enabled=(device.type == "cuda")):253 _z, taps = encoder.forward(bb, grid, region_packed, tap_blocks=GPU_TAPS)254 for t in GPU_TAPS:255 enc = taps[t].float().reshape(Bb, -1, k, taps[t].shape[-1]).cpu() # (Bb, n, k, d)Freeze encoder weights
48def _freeze(m, *, device):49 m.eval().to(device)50 for p in m.parameters():51 p.requires_grad_(False)52 return mCross-session: shared contact labels
320 a = train_rec.get("elec_labels")321 t = test_rec.get("elec_labels")322 if a is None or t is None:323 return None, None, 0324 a = np.asarray(a)325 t = np.asarray(t)326 common = np.intersect1d(a, t)327 if common.size == 0:328 return None, None, 0329 a_idx = [int(np.where(a == c)[0][0]) for c in common]330 t_idx = [int(np.where(t == c)[0][0]) for c in common]331 return a_idx, t_idx, int(common.size)Cross-subject: mean-pool contacts by region
164 B = x.shape[0]165 blocks = []166 for p in present:167 cols = np.where(region_canon == p)[0]168 sub = x[:, cols] # (B, |cols|, *feat)169 blocks.append(sub.mean(1).reshape(B, -1)) # (B, prod(feat))170 return torch.stack(blocks, dim=1).to(torch.float16) # (B, |P|, F)Cross-subject: align shared atlas IDs
238 a_p = np.asarray(anchor_rec["present_parcels"], dtype=np.int64)239 t_p = np.asarray(test_rec["present_parcels"], dtype=np.int64)240 common = np.intersect1d(a_p, t_p)241 if UNKNOWN_REGION_ID in common:242 print(f"[check] dropping unknown region {UNKNOWN_REGION_ID} from intersection", flush=True)243 common = common[common != UNKNOWN_REGION_ID]244 if common.size == 0:245 return None, None, common246 a_idx = [int(np.where(a_p == c)[0][0]) for c in common]247 t_idx = [int(np.where(t_p == c)[0][0]) for c in common]248 return a_idx, t_idx, common306 z_tr = _feat(anchor_rec, enc, tr, a_idx)307 z_va, z_te = _feat(test_rec, enc, va, t_idx), _feat(test_rec, enc, te, t_idx)Flatten aligned features
143def _feat(rec, enc, rows, col_idx=None) -> np.ndarray:144 """Gather rows and aligned contact/region columns, then flatten to fp32."""145 with _timed("gather_fp16"):146 x = rec["feats"][enc]["raw"][np.asarray(rows, dtype=np.int64)]147 if col_idx is not None:148 x = x[:, np.asarray(col_idx, dtype=np.int64)]149 with _timed("to_fp32"):150 x = x.to(torch.float32).numpy()151 return x.reshape(x.shape[0], -1)Standardize using training trials
z_tr: training trials × flattened features. Compute mean/std over axis=0: across trials, separately for each contact/region × token position × embedding coordinate. Apply those training statistics to train, validation, and test; no averaging across tokens.
119 for lo, hi in zip(edges, edges[1:]):120 b = z_tr[:, lo:hi]121 mu = b.mean(axis=0)122 sd = b.std(axis=0)123 sd[sd == 0] = 1.0124 b -= mu125 b /= sd126 for z in others:127 zb = z[:, lo:hi]128 zb -= mu129 zb /= sd130 return z_tr, othersCaller: use training statistics for both evaluation splits
252 a, (b, c) = _standardize_inplace(z_tr, [z_va, z_te])Memory-bounded column blocks
115 d = z_tr.shape[1]116 edges = list(range(0, d, blk)) + [d]117 if len(edges) > 2 and edges[-1] - edges[-2] == 1:118 edges.pop(-2) # avoid a width-one trailing blockSolve ridge over the lambda grid
78LAM_MULTS = tuple(np.logspace(-4.0, 4.0, 25))153def _linear_grams(z_tr, evals):154 """fp32 Gram products promoted to fp64 for the ridge solve."""155 with _timed("gram_gemm"):156 g = np.asarray(z_tr @ z_tr.T, dtype=np.float64) # fp32 GEMM → fp64 Gram157 with _timed("eval_kernels"):158 kern = {name: np.asarray(z @ z_tr.T, dtype=np.float64)159 for name, (z, _) in evals.items()}160 return g, kern168 g, kern = _linear_grams(z_tr, evals)169 n = g.shape[0]170 with _timed("eigh"):171 w, V = np.linalg.eigh(g) # G symmetric PSD ⇒ w >= 0172 c = V.T @ np.asarray(y_tr, dtype=np.float64)173 base = float(np.sum(w) / max(n, 1)) # trace(G)/n — the λ scale174 out: dict = {name: {} for name in evals}175 with _timed("lam_sweep"):176 for m in LAM_MULTS:177 alpha = V @ (c / (w + m * base))178 for name, (_, y) in evals.items():179 s = kern[name] @ alpha180 out[name][m] = (auroc(s, y) if len(y) >= 2 else float("nan"))181 return outSelect lambda on validation AUROC
209 finite = [(m, va) for m, va in d["val"].items() if not np.isnan(va)]210 if not finite:211 return {"val": float("nan"), "test": float("nan"), "lam_mult": float("nan"),212 "lam_pin": "", "lam_pinned": False, "n_tied": 0}213 best_val = max(va for _, va in finite)214 # Insertion order is LAM_MULTS ascending, so tied[0] is the smallest tied λ — exactly what the215 # strict-`>` loop this replaced selected. That equivalence is what keeps "argmax" byte-faithful.216 tied = [m for m, va in finite if va == best_val]217 m = tied[0]Returned metrics
218 pin = "lo" if m == LAM_MULTS[0] else ("hi" if m == LAM_MULTS[-1] else "")219 out = {"val": best_val, "test": d["test"][m], "lam_mult": float(m),220 "lam_pin": pin, "lam_pinned": pin == "lo", "n_tied": len(tied)}221 return outPretraining
How the encoder learns · two-second clips · encode 25%, reconstruct 75%
Figure 2 in the paper ↗. Outlined = masked targets · faded = no loss. 75% masked per shaft/clip. Mean over valid bins → mean over masked tokens. Schematic bars; counts use 128 contacts.
Gray positions are withheld from the encoder. The selected locations are illustrative; training samples an exact 75% mask per array/clip.
Masking & loss details
Randomly mask exactly 75% of tokens per shaft/clip; the fraction per contact or band can vary. Average squared error over each token’s valid bins, then over masked tokens. The mid band’s padded seventh bin is excluded. Bands are not equally weighted: they contribute 8 / 32 / 64 tokens per contact. Targets are session-normalized STFT magnitudes, with no extra per-token normalization. No teacher. Decoder attention carries gradients back to visible encoder outputs.
The paper’s run uses 27.9 hours from 7 subjects, 2-second clips, AdamW for 55,000 steps, and a 5k warmup / 45k constant / 5k cooldown schedule. Subjects 7 and 10 are excluded from pretraining; checkpoint timing is chosen on pretraining sessions, not benchmark test scores. Training excerpts are selected from the research implementation; they are not a standalone training package.
Sample an exact mask per shaft / clip
729 vpos = valid.reshape(-1).nonzero(as_tuple=True)[0] # (N,) valid grid positions730 vcontact = geom.gather_idx.reshape(-1)[vpos] # (N,) permutation of 0..N−1731 unit_of_contact = torch.zeros(n, dtype=torch.long, device=dev)732 unit_of_contact[vcontact] = vpos // c733 slow_m, mid_m, hga_m = _sample_global_iid_band_time(734 n, (t // SLOW_STRIDE, t // MID_STRIDE, t),735 frac=cfg.iid_mask_frac, n_rows=r, unit_of_contact=unit_of_contact,736 generator=generator, device=dev,737 )473 masked, in_loss = token_flags_r6(grid, masks) # (B, total); in_loss == maskedMap band masks into packed token order
318 band_masks = (masks.slow_mask, masks.mid_mask, masks.hga_mask) # each (B, N, T_b)319 B = masks.contact_mask.shape[0]320 device = grid.contact.device321322 contact_masked = masks.contact_mask[:, grid.contact, grid.band] # (B, total)323324 temporal_masked = torch.zeros(B, grid.total, dtype=torch.bool, device=device)325 for b, bm in enumerate(band_masks): # bm (B, N, T_b)326 sel = grid.band == b # (total,) tokens of band b327 ct = grid.contact[sel] # (n_b,) each token's contact328 pos = grid.bandpos[sel] # (n_b,) band-token index329 temporal_masked[:, sel] = bm[:, ct, pos] # (B, n_b) own-contact mask at (contact, pos)330331 masked = contact_masked | temporal_masked332 return masked, masked # score EVERY masked token — no margin gateExact-count masking implementation
658 total = n_units * sum(lengths)659 # Per-cell unit id, in the SAME flat layout the band split below undoes: each band segment is660 # (n_units, length) row-major = contact-major, so a contact's `length` cells are contiguous.661 unit_of_cell = torch.cat([unit_of_contact.repeat_interleave(length) for length in lengths])662 n_unit = int(unit_of_contact.max()) + 1663 cells = torch.bincount(unit_of_cell, minlength=n_unit) # (U,) cells per attention unit664 cnt = torch.round(frac * cells.float()).long() # (U,) exact masked count per unit665 start = torch.cat([torch.zeros(1, dtype=torch.long, device=device), cells.cumsum(0)[:-1]])666 # Offsetting the key by the unit id makes the sort group by unit and order by key WITHIN unit,667 # so unit u occupies sorted positions [start[u], start[u]+cells[u]) and the global rank minus668 # start[u] IS the rank within the unit. One argsort pair, no python loop over units.669 key = torch.rand(n_rows, total, generator=generator, device=device)670 key = key + unit_of_cell[None, :].to(key.dtype) * 2.0671 rank = key.argsort(-1).argsort(-1) - start[unit_of_cell][None, :]672 flat = rank < cnt[unit_of_cell][None, :] # (R, total) exactly cnt[u] per (row, unit)673 out, off = [], 0674 for length in lengths:675 seg = n_units * length676 out.append(flat[:, off:off + seg].reshape(n_rows, n_units, length))677 off += seg678 return out[0], out[1], out[2]Gather visible tokens before the encoder
478 pack = build_visible_pack(479 grid, masked, parcel_packed, m_vis=m_vis, max_seqlen=pack_max_seqlen480 )183 tokens, _ = self.stem(bands) # per-band tuple (B, N, T_b, d) on the shared 32 Hz clock184 x_full = pack_band_tokens(tokens, grid) # (B, total, d) flat grid order189 d = x_full.shape[-1]190 x_vis = x_full.gather(1, pack.idx[:, :, None].expand(-1, -1, d)) # (B, M_vis, d)191 return self.encoder.forward_flat_pack(x_vis, pack, tap_blocks=tap_blocks)Visible-index construction
201 B, total = masked.shape202 device = masked.device203 visible = ~masked204 if m_vis is None:205 m_vis = int(visible.sum(1).max().item())206 ar = torch.arange(total, device=device)207 key = torch.where(visible, ar[None, :].expand(B, total), total)208 idx = key.sort(dim=1).values[:, :m_vis] # (B, M_vis) visible columns, grid orderNormalize and concatenate four taps
35ENC_SUP_TAPS: tuple[int, ...] = (3, 6, 9, 12)36N_LEVELS = len(ENC_SUP_TAPS)146 if self.deep_sup and (i + 1) in self.sup_taps:147 level_index = self.sup_taps.index(i + 1)148 normalized_level = self.norms_block[level_index](xf)149 levels.append(normalized_level)150 # Inference uses raw taps (d each); pretraining used normalized taps (4*d).151 if self.deep_sup:152 out = torch.cat(levels, dim=-1) # (B*total,4*d)LayerNorm definitions
72 self.norms_block = nn.ModuleList(73 [nn.LayerNorm(d_model, eps=LN_EPS) for _ in self.sup_taps]74 )Fuse four levels into decoder width
338 self.enc_to_pred = nn.Sequential(339 nn.Linear(N_LEVELS * self.d_model, self.d_model),340 nn.GELU(),341 nn.Linear(self.d_model, self.d_pred),342 )Restore the full sequence
500 pred_in = scatter_visible(501 self.enc_to_pred(z), pack.idx, grid.total, self.mask_token502 ) # (B, total, 192)241 out = fill.to(h_vis.dtype).expand(B, total, d).clone()242 out.scatter_(1, idx[:, :, None].expand(B, m_vis, d), h_vis)243 return outLearned mask-token definition
415 self.mask_token = nn.Parameter(torch.zeros(1, 1, self.d_pred))Decoder: band, region, and RoPE
514 pred_in = pred_in + self.pred_band_emb[grid.band][None] # (B, total, d_pred)515 # forward_flat adds the parcel embed to every token itself (towers._run_flat), so with the516 # band embed above the predictor now sees parcel + band on ALL tokens — the same identity517 # pair the encoder gets before its first L1 block.518 h = self.predictor.forward_flat(pred_in, grid, parcel_packed) # (B, total, 192)302 B, M, d = x_flat.shape303 device = x_flat.device304 x_flat = x_flat + self.parcel_embed(parcel_bm).to(x_flat.dtype) # (B, M, d)317 depth_b = depth_bm.reshape(B * M)318 time_b = time_bm.reshape(B * M)319 # shared L1 rope table: build once over the flattened per-clip coords, reused by every block.320 first_l1 = next((blk for blk in self.blocks if isinstance(blk, L1Block)), None)321 rope_cs = None322 if first_l1 is not None:323 cos, sin = first_l1.rope.cos_sin(depth_b, time_b) # (B·M, head_dim)324 rope_cs = (cos, sin)Decoder block loop and final normalization
326 xf = x_flat.reshape(B * M, d)327 taps: dict[int, Tensor] = {}328 levels: list[Tensor] = []329 for i, blk in enumerate(self.blocks):330 # r4 is L1-only (Design B); an L2 in the layout has no flat path (guard).331 if not isinstance(blk, L1Block):332 raise TypeError("forward_flat is L1-only (Design B); found a non-L1 block")333 xf = blk.forward_flat(334 xf, depth_b, time_b, cu_b, cu_drop_b, max_seqlen, rope_cs=rope_cs,335 )336 if (i + 1) in tap_blocks:337 taps[i + 1] = xf.reshape(B, M, d)338 if self.deep_sup and (i + 1) in self.sup_taps:339 levels.append(self.norms_block[self.sup_taps.index(i + 1)](xf))340 out = torch.cat(levels, dim=-1) if self.deep_sup else self.norm_out(xf)341 out = out.reshape(B, M, out.shape[-1])Decoder construction & width configuration
50PRED_LAYOUT: tuple[str, ...] = ("L1",) * 6 # Design B (Ben 2026-07-15): L1-only, 6 layers; class W dropped ⇒ no cross-shaft predictor75 n_heads = enc_d_model // HEAD_DIM76 return enc_d_model, n_heads, enc_d_model // 2, n_heads370 _, _, pred_d, pred_h = width_spec(enc_d_model)371 return V3Tower(372 PRED_LAYOUT, d_model=pred_d, n_heads=pred_h, n_parcels=n_parcels,373 parcel_embed=parcel_embed, space_rope=space_rope,374 )Band heads: 192 → 7 / 6 / 7
384 self.mae_heads = nn.ModuleList(385 nn.Linear(self.d_pred, nb) for nb, _ in PER_BAND_SPECS386 )387 self.mae_heads.apply(init_transformer_weights)902 f_max = max(nb for nb, _ in PER_BAND_SPECS)903 outs: list[Tensor] = []904 for head, (nb, _) in zip(self.mae_heads, PER_BAND_SPECS):905 o = head(h) # (B, total, F_b)906 if nb < f_max:907 o = F.pad(o, (0, f_max - nb))908 outs.append(o) # (B, total, F_MAX)909 stack = torch.stack(outs, dim=-1) # (B, total, F_MAX, N_BANDS)910 onehot = F.one_hot(grid.band, len(PER_BAND_SPECS)).to(h.dtype) # (total, N_BANDS)911 return (stack * onehot[None, :, None, :]).sum(-1) # (B, total, F_MAX)Targets and valid bins
850 B, N = bands[0].shape[0], bands[0].shape[1]851 T = bands[0].shape[-1]852 f_max = max(nb for nb, _ in PER_BAND_SPECS)853 padded: list[Tensor] = []854 for b, (nb, _) in enumerate(PER_BAND_SPECS):855 xb = bands[b].transpose(-1, -2) # (B, N, T, F_b)856 if nb < f_max:857 xb = F.pad(xb, (0, f_max - nb)) # zero-pad freq to F_MAX858 padded.append(xb)859 stack = torch.stack(padded, dim=1) # (B, N_BANDS, N, T, F_MAX)860 flat = stack.reshape(B, len(PER_BAND_SPECS) * N * T, f_max)861 lin = (grid.band * N + grid.contact) * T + grid.time_pos # (total,)862 target = flat[:, lin, :] # (B, total, F_MAX)863 device = grid.band.device864 f_by_band = torch.tensor([nb for nb, _ in PER_BAND_SPECS], device=device)865 feat_count = f_by_band[grid.band] # (total,)866 feat_valid = torch.arange(f_max, device=device)[None, :] < feat_count[:, None]867 return target, feat_valid, feat_countReconstruction call site
535 return self._mae_output(536 bands, grid, h, in_loss,537 enc_taps=enc_taps if collect_taps else None,538 collect_taps=collect_taps,539 norm_pix=self.force_norm_pix or not self.r6, # r6's ONLY loss-side delta from r4540 )Masked-only MSE
604 pred = self._mae_pred(h, grid) # (B, total, F_MAX); pad slot = 0605 target = target.to(pred.dtype)606 w = in_loss.to(pred.dtype) # (B, total) masked weight607 fv = feat_valid[None].to(pred.dtype) # (1, total, F_MAX)608 # per-token MSE over that token's own valid bins (pad excluded), then masked-mean.609 se_tok = (((pred - target) ** 2) * fv).sum(-1) / feat_count[None].to(pred.dtype)610 mae_loss = (se_tok * w).sum() / w.sum().clamp(min=1.0)Throughput
Batch the pretraining objective · fixed token budget, variable shaft lengths
Optimizer & checkpoint
Use the final 55k-step checkpoint. Cooldown timing was fixed from an earlier pretraining-only probe; no benchmark-score checkpoint selection.
Compute choices
Exact-count 75% masking keeps the visible total fixed. Shaft count and lengths still vary, so compilation uses dynamic=True. The final run uses two NVIDIA H100 GPUs with BF16 mixed precision. Counts refer to sampled contact-windows, not unique electrodes. The run sets shaft_alpha=0: subject weights are equal, shafts are drawn uniformly within subject, and clip starts are uniform over valid aligned starts. Artifact-overlapping clips are excluded. The schedule axis counts optimizer updates, not microbatches. In code, warmup uses (step + 1) / 5000; the first update is at 1.2e-6 and the drawn zero is the conceptual starting boundary. The scheduler class is named WarmupCosine, but the run selects decay_shape="linear". stable_steps=50000 is the plateau end, not its duration. The active sampler uses no filler; no separate padding to the longest shaft is needed.
Uniform subject → shaft → clip
287 --batch-unit shaft --contact-budget $CB --shaft-alpha 0.0 \127 self._subj_w = pool.subject_n_shafts.to(torch.float64) ** self.alpha128129 def draw(self, generator: torch.Generator) -> ShaftRef:130 u = int(torch.multinomial(self._subj_w, 1, generator=generator).item())131 idxs = self.pool.subject_ref_idxs[u]132 j = int(torch.randint(len(idxs), (1,), generator=generator).item())133 return self.pool.refs[idxs[j]]95 t0 = sample_clip_start(96 n_frames=s.n_frames, clip_frames=self.clip_frames, bad_spans_s=s.bad_spans_s,97 fps=self.fps, generator=g, start_align=self.start_align,98 )105 counts = torch.tensor([c for _, c in aligned], dtype=torch.float64)106 which = int(torch.multinomial(counts / counts.sum(), 1, generator=generator).item())107 a_lo, count = aligned[which]108 step = int(torch.randint(0, count, (1,), generator=generator).item())109 return a_lo + step * start_alignFill the contact budget; trim the final shaft
148 drawn: list[tuple[ShaftRef, int]] = []149 total = 0150 # each iter adds ≥1 contact ⇒ the budget closes within contact_budget draws (draw is151 # WITH replacement, so the pool never "runs dry" — the loop always reaches the budget).152 for _ in range(contact_budget):153 r = self.draw(generator)154 room = contact_budget - total155 if r.n_contacts >= room: # this shaft closes (or overshoots) the budget → trim it156 drawn.append((r, room))157 return drawn158 drawn.append((r, r.n_contacts))159 total += r.n_contacts160 return drawnEffective batch per optimizer step
144CB=8250145ACC=4146WORLD=2292 --batch-size 32 --accumulate-grad-batches $ACC \293 --accelerator gpu --devices $WORLD --precision bf16-mixed \GPU attention uses packed values and boundaries
74 def _nt(x: Tensor) -> Tensor:75 return torch.nested.nested_tensor_from_jagged(76 x, offsets=cu_seqlens_drop, min_seqlen=1, max_seqlen=max_seqlen77 ).transpose(1, 2)7879 # Same within-block attention, without allocating a full total-by-total mask.80 ctx = F.scaled_dot_product_attention(_nt(q), _nt(k), _nt(v))81 return ctx.transpose(1, 2).values()Compilation handles changing shaft geometry
876 module.model = torch.compile(module.model, dynamic=True) # type: ignore[assignment]AdamW configuration
198 "optimizer": {199 "name": "AdamW",200 "lr": lr,201 "kwargs": {"weight_decay": weight_decay, "betas": [0.9, adam_beta2]},202 },288 --lr 6e-3 --weight-decay 0.04 --grad-clip 3.0 --adam-beta2 0.95 --seed 33 \Final-run learning-rate schedule
288 --lr 6e-3 --weight-decay 0.04 --grad-clip 3.0 --adam-beta2 0.95 --seed 33 \289 --min-lr-ratio 0.0 --stable-steps 50000 --decay-shape linear \301 --ssl-max-steps 55000572 p.add_argument("--warmup-steps", dest="warmup_steps", type=int, default=5000)Scheduler implementation
81 return 1.0 - progress # linear8283 def _multiplier(self, step: int) -> float:84 if self.warmup_steps > 0 and step < self.warmup_steps:85 return min(1.0, (step + 1) / self.warmup_steps)86 if step < self.stable_steps:87 return 1.088 decay_total = max(1, self.total_steps - self.stable_steps)89 progress = min(1.0, max(0.0, (step - self.stable_steps) / decay_total))90 shaped = self._decay(progress)91 return self.min_lr_ratio + (1.0 - self.min_lr_ratio) * shaped9293 def get_lr(self) -> list[float]: # type: ignore[override]94 mult = self._multiplier(self.last_epoch)95 return [base_lr * mult for base_lr in self.base_lrs]