OSCR

MEG-GPT: A transformer-based foundation model for magnetoencephalography data.

Code ↔ Paper

13 matches between paragraphs of the paper and lines of its authors' code, computed by the harvester (lexical-v1). Click a colored paragraph or line to see its counterpart.

The 13 matches
  1. [1] § Methods › The foundation model: MEG-GPT › Transformer and prediction head ↔ osl_foundation/inference/layers.py, lines 237–342 · score 0.70 · receptive field, attention matrix, latent sequence length, attends, unpatched, masked
  2. [2] § Methods › The foundation model: MEG-GPT › Transformer and prediction head ↔ osl_foundation/models/meg_gpt.py, lines 406–444 · score 0.65 · Feed Forward layer, residual connections, MEG GPT
  3. [3] § Results › MEG-GPT captures bursting dynamics in MEG data ↔ examples/paper/analysis/4_bursting_behaviour.py, lines 295–367 · score 0.62 · bursting behaviour, AR model, lifetime, interval, MEG GPT
  4. [4] § Results › MEG-GPT extracts features that enhance decoding performance ↔ examples/paper/decoding/7_decoding_performances.py, lines 87–138 · score 0.62 · decoding accuracy, zero shot, Fine tuning, baseline, model
  5. [5] § Methods › The foundation model: MEG-GPT › Masked multi-head attention ↔ osl_foundation/models/meg_gpt.py, lines 280–322 · score 0.61 · unpatched sequences, attention layers, sequence length, MEG GPT, patching, Decoder
  6. [6] § Methods › The foundation model: MEG-GPT › Masked multi-head attention ↔ osl_foundation/inference/layers.py, lines 237–342 · score 0.58 · receptive field, sequence length, attend, unpatched, patching, layers
  7. [7] § Results › MEG-GPT extracts features that enhance decoding performance ↔ examples/paper/decoding/7_decoding_performances.py, lines 87–138 · score 0.57 · decoding accuracy, zero shot, fine tuned, baseline
  8. [8] § Methods › The foundation model: MEG-GPT › Generating new data ↔ osl_foundation/utils/sampling.py, lines 91–133 · score 0.56 · cumulative probability exceeds, smallest, tokens
  9. [9] § Results › The tokeniser reconstructs MEG data with high accuracy and generalises to unseen data ↔ examples/paper/analysis/1_tokeniser_reconstructs.py, lines 134–183 · score 0.56 · tokeniser reconstructions, variance explained, Wakeman Henson, PVE, Cam, training
  10. [10] § Results › MEG-GPT captures spatial and spectral characteristics of real data ↔ examples/paper/analysis/3_subject_fingerprints.py, lines 166–285 · score 0.54 · power spectral density, spectral features, Welch, GPT, MEG
  11. [11] § Results › The tokeniser reconstructs MEG data with high accuracy and generalises to unseen data ↔ examples/paper/analysis/1_tokeniser_reconstructs.py, lines 134–183 · score 0.53 · variance explained, Wakeman Henson, PVE, reconstruction, tokeniser, Cam
  12. [12] § Results › MEG-GPT captures bursting dynamics in MEG data ↔ examples/paper/analysis/4_bursting_behaviour.py, lines 507–556 · score 0.53 · Wavelet transform, AR model, Location, bursting, HMM, MEG GPT
  13. [13] § Methods › Tokeniser › Training › Annealing ↔ osl_foundation/inference/layers.py, lines 345–386 · score 0.52 · Gumbel Softmax, annealing, argmax, weighted, token, training

Paper

Loaded from Europe PMC by your browser, not stored by OSCR: doi.org · Europe PMC

The paper is loaded when this pane is shown.

The authors' code

Python · 1,134 lines · 40 KB · MIT · 3 matches

  1. from typing import Union
  3. import numpy as np
  4. import tensorflow as tf
  5. import tensorflow_probability as tfp
  8. def rnn_layer(
  9. rnn_type: str, rnn_n_units: int, return_sequences: bool = True
  10. ) -> tf.keras.layers.Layer:
  11. """
  12. Create an RNN layer.
  14. Parameters
  15. ----------
  16. rnn_type : str
  17. Type of RNN layer. Options are 'gru', 'lstm'.
  18. rnn_n_units : int
  19. Number of units in the RNN layer.
  20. return_sequences : bool, optional
  21. Whether to return sequences.
  23. Returns
  24. -------
  25. rnn_layer : tf.keras.layers.Layer
  26. RNN layer.
  27. """
  28. if rnn_type == "gru":
  29. return tf.keras.layers.GRU(
  30. rnn_n_units, return_sequences=return_sequences, stateful=False
  31. )
  32. elif rnn_type == "lstm":
  33. return tf.keras.layers.LSTM(
  34. rnn_n_units, return_sequences=return_sequences, stateful=False
  35. )
  36. else:
  37. raise ValueError(f"Unknown RNN type: {rnn_type}")
  40. class IdentityLayer(tf.keras.layers.Layer):
  41. """
  42. Identity layer.
  43. This layer directly returns the input tensor.
  44. """
  46. def call(self, inputs, **kwargs):
  47. return inputs
  50. class MSELossLayer(tf.keras.layers.Layer):
  51. """
  52. Layer for computing the mean squared error loss.
  53. This is a wrapper around tf.keras.losses.MeanSquaredError.
  54. """
  56. def __init__(self, **kwargs):
  57. super().__init__(**kwargs)
  58. self.loss_fn = lambda y_true, y_pred: tf.reduce_mean(tf.square(y_true - y_pred))
  60. def call(self, y_true, y_pred, **kwargs):
  61. loss = self.loss_fn(y_true, y_pred)
  62. self.add_loss(loss)
  63. return tf.expand_dims(loss, -1)
  66. class SinusoidalPositionalEncodingLayer(tf.keras.layers.Layer):
  67. """
  68. Layer for generating sinusoidal positional encoding for position embeddings.
  69. Implemented as in "Attention is All You Need" (Vaswani et al., 2017), and partly
  70. adpated from the `huggingface/transformers` library and TensorFlow
  71. (https://www.tensorflow.org/text/tutorials/transformer).
  73. Parameters
  74. ----------
  75. sequence_length : int
  76. Length of the sequence.
  77. """
  79. def __init__(self, sequence_length: int, **kwargs):
  80. super().__init__(**kwargs)
  81. self.sequence_length = sequence_length
  83. def call(self, inputs, **kwargs):
  84. # Validation
  85. embedding_dim = tf.shape(inputs)[-1]
  86. tf.debugging.assert_equal(
  87. tf.math.floormod(embedding_dim, 2),
  88. tf.constant(0, dtype=tf.int32),
  89. message="embedding dimension must be even for positional encoding.",
  90. )
  92. # Precompute scaling factor
  93. embedding_dim_float = tf.cast(embedding_dim, tf.float32)
  94. denominator = tf.math.pow(
  95. 10000.0,
  96. 2 * (tf.range(embedding_dim_float) // 2) / embedding_dim_float,
  97. )
  99. # Get position indices
  100. position_indices = tf.range(self.sequence_length, dtype=tf.float32)[
  101. :, tf.newaxis
  102. ]
  103. # position_indices.shape = (sequence_length, 1)
  105. # Compute positional encoding
  106. angle_rads = position_indices / denominator
  107. # angle_rads.shape = (sequence_length, embedding_dim)
  108. pos_encoding = tf.TensorArray(dtype=tf.float32, size=embedding_dim)
  109. for i in range(embedding_dim):
  110. if i % 2 == 0:
  111. pos_encoding = pos_encoding.write(i, tf.sin(angle_rads[:, i]))
  112. else:
  113. pos_encoding = pos_encoding.write(i, tf.cos(angle_rads[:, i]))
  114. pos_encoding = tf.transpose(pos_encoding.stack())
  115. # pos_encoding.shape = (sequence_length, embedding_dim)
  117. # Reshape the positions to match shapes of other embeddings
  118. positions = tf.expand_dims(
  119. tf.expand_dims(pos_encoding, axis=0), axis=0
  120. ) # positions.shape = (1, 1, sequence_length, embedding_dim)
  121. return positions
  124. class RotaryPositionEmbeddingLayer(tf.keras.layers.Layer):
  125. """
  126. Layer for generating Rotary Position Embedding (RoPE). Implemented as in
  127. "RoFormer: Enhanced Transformer with Rotary Position Embedding" (Su et al., 2022), and
  128. adpated from the PyTorch and Keras 3 libraries (see below for the reference).
  130. References:
  131. - https://pytorch.org/torchtune/0.2/_modules/torchtune/modules/position_embeddings.html#RotaryPositionalEmbeddings
  132. - https://github.com/keras-team/keras-hub/blob/v0.17.0/keras_hub/src/layers/modeling/rotary_embedding.py
  134. Parameters
  135. ----------
  136. embedding_dim : int
  137. Dimension of the rotary position embedding. When using multi-head attention blocks,
  138. the dimension is usually `embedding_dim // n_heads`.
  139. max_sequence_length : int
  140. Maximum expected length of the sequence. If the model dimensions of the key and
  141. query vectors differ, the sequence length should be the maximum of the two.
  142. base : int, optional
  143. The base value for computing the rotation angles.
  144. """
  146. def __init__(
  147. self, embedding_dim: int, max_sequence_length: int, base: int = 10000, **kwargs
  148. ):
  149. super().__init__(**kwargs)
  151. # Validation
  152. if embedding_dim % 2 != 0:
  153. raise ValueError(
  154. "embedding dimension must be even for positional encoding."
  155. )
  157. self.embedding_dim = embedding_dim
  158. self.max_sequence_length = max_sequence_length
  159. self.base = base
  160. self.rotation_values = self._compute_rotation_matrix()
  162. def _compute_rotation_matrix(self) -> tf.Tensor:
  163. # Compute pre-defined theta angles
  164. theta = 1.0 / (
  165. self.base
  166. ** (
  167. tf.range(0, self.embedding_dim, 2, dtype=tf.float32)[
  168. : (self.embedding_dim // 2)
  169. ]
  170. / self.embedding_dim
  171. )
  172. )
  174. # Create position indices
  175. seq_idx = tf.range(self.max_sequence_length, dtype=tf.float32)
  176. # NOTE: Having 0 as a starting index is okay, since we don't have to rotate the first position.
  178. # Compute outer product of position index and theta
  179. idx_theta = tf.einsum("i,j->ij", seq_idx, theta)
  180. # idx_theta.shape = (sequence_length, embedding_dim // 2)
  182. # Compute elements of the rotation matrix
  183. rotation_values = tf.stack([tf.cos(idx_theta), tf.sin(idx_theta)], axis=-1)
  184. # rotation_values.shape = (sequence_length, embedding_dim // 2, 2)
  186. return rotation_values
  188. def call(self, inputs, **kwargs):
  189. # Get input sequence length
  190. sequence_length = tf.shape(inputs)[3]
  191. # inputs.shape = (batch_size, n_heads, n_channels, sequence_length, key_dim)
  193. # Validation
  194. tf.debugging.assert_greater_equal(
  195. self.max_sequence_length,
  196. sequence_length,
  197. message="input sequence length should not exceed max_sequence_length.",
  198. )
  199. tf.debugging.assert_equal(
  200. tf.shape(inputs)[-1],
  201. self.embedding_dim,
  202. message="the model dimension of the input tensor must match the embedding dimension.",
  203. )
  205. # Reshape the rotation values to be compabitle with the input shape
  206. rot_vals = tf.reshape(
  207. self.rotation_values[:sequence_length],
  208. (1, 1, 1, sequence_length, self.embedding_dim // 2, 2),
  209. )
  210. # NOTE: We clip the rotation values to the sequence length, in case sequence_length is
  211. # less than max_sequence_length as the dimension for key and query differs.
  213. # Match the input shape to the rope shape
  214. rot_shape = tf.concat(
  215. [tf.shape(inputs)[:-1], [self.embedding_dim // 2, 2]], axis=0
  216. )
  217. x_reshaped = tf.reshape(inputs, rot_shape)
  218. # x_reshaped.shape = (batch_size, n_heads, n_channels, sequence_length, key_dim // 2, 2)
  220. # Calculate the rotary position embedding
  221. rope = tf.stack(
  222. [
  223. rot_vals[..., 0] * x_reshaped[..., 0]
  224. - rot_vals[..., 1] * x_reshaped[..., 1],
  225. rot_vals[..., 1] * x_reshaped[..., 0]
  226. + rot_vals[..., 0] * x_reshaped[..., 1],
  227. ],
  228. axis=-1,
  229. )
  230. # rope.shape = (batch_size, n_heads, n_channels, sequence_length, key_dim // 2, 2)
  231. rope = tf.reshape(rope, tf.shape(inputs))
  232. # rope.shape = (batch_size, n_heads, n_channels, sequence_length, key_dim)
  234. return rope
  237. class ALiBiPositionEmbeddingLayer(tf.keras.layers.Layer):
  238. """
  239. Layer for generating Attention with Linear Biases (ALiBi) position embedding. Implemented
  240. based on "Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation"
  241. (Press et al., 2022) and adapted from the references below.
  243. References:
  244. - https://github.com/ofirpress/attention_with_linear_biases
  245. - https://nn.labml.ai/transformers/alibi/index.html
  247. Parameters
  248. ----------
  249. n_heads : int
  250. Number of attention heads.
  251. n_patches : int
  252. Number of patches to attend to.
  253. patch_length : int
  254. Patch length.
  255. unpatched_length : int
  256. Number of unpatched elements to attend to.
  257. """
  259. def __init__(
  260. self,
  261. n_heads: int,
  262. n_patches: int,
  263. patch_length: int,
  264. unpatched_length: int,
  265. **kwargs,
  266. ):
  267. super().__init__(**kwargs)
  268. self.n_heads = n_heads
  269. self.n_patches = n_patches
  270. self.patch_length = patch_length
  271. self.unpatched_length = unpatched_length
  273. def _get_slopes(self) -> tf.Tensor:
  274. # Compute the slopes for the linear biases (static, non-learned)
  275. h = tf.cast(self.n_heads, tf.float32)
  276. n = 2 ** tf.math.floor(tf.math.log(h) / tf.math.log(2.0)) # nearest power of 2
  277. ratio = tf.pow(2.0, -8.0 / n)
  278. m = tf.pow(
  279. ratio, tf.range(1, tf.cast(n, tf.int32) + 1, dtype=tf.float32)
  280. ) # head-specific slopes
  282. # When n_heads is not a power of 2, then we add additional slopes
  283. if n < h:
  284. ratio_hat = tf.pow(2.0, -4.0 / n)
  285. m_hat = tf.pow(ratio_hat, tf.range(1, 1 + 2 * (h - n), 2, dtype=tf.float32))
  286. m = tf.concat([m, m_hat], axis=0)
  288. return m
  290. def _get_relative_bias_matrix(self, mask: tf.Tensor) -> tf.Tensor:
  291. # Get head-specific slopes
  292. m = self._get_slopes()
  294. # Get a binary masking matrix
  295. mask = 1 - mask
  296. cum_mask = tf.cumsum(mask, axis=0) # cumulative sum of the mask
  297. # NOTE: In the mask matrix, columns are where you attend to and rows are the targets.
  299. # Compute the relative distance matrix for the patches
  300. patch_indices = tf.cast(tf.range(self.n_patches), dtype=tf.float32)
  301. avg_idx = (
  302. (self.patch_length * patch_indices + 1)
  303. + self.patch_length * (patch_indices + 1)
  304. ) / 2.0
  305. distance_patched = (
  306. -(cum_mask[:, : self.n_patches] + self.patch_length * (patch_indices + 1))
  307. + avg_idx
  308. )
  309. # NOTE: Considering a patch as a single receptive field, we calculate relative distance between the patch
  310. # and the current sequence index (i.e., the distance between the center of the patch and the current index).
  311. # This amounts to the average distance between the current index and all the indices in the patch.
  313. # Compute the relative distance matrix for the unpatched elements
  314. distance_unpatched = -cum_mask[
  315. :, self.n_patches : self.n_patches + self.unpatched_length
  316. ]
  318. # Concatenate two matrices
  319. distance = tf.concat([distance_patched, distance_unpatched], axis=1)
  320. # distance.shape = (latent_sequence_length, n_patches + unpatched_length)
  322. # Compute head-specific relative bias matrices
  323. masked_distance = distance * mask
  324. alibi_matrix = masked_distance[None, :, :] * m[:, None, None]
  325. # alibi_matrix.shape = (n_heads, latent_sequence_length, n_patches + unpatched_length)
  327. return alibi_matrix
  329. def call(self, inputs, mask: tf.Tensor, **kwargs):
  330. # Get the relative bias matrix
  331. alibi_matrix = self._get_relative_bias_matrix(mask)
  332. # alibi_matrix.shape = (n_heads, latent_sequence_length, n_patches + unpatched_length)
  334. # Reshape the relative bias matrix to match the attention matrix
  335. alibi_matrix = tf.reshape(
  336. alibi_matrix, (1, self.n_heads, 1, tf.shape(mask)[0], tf.shape(mask)[1])
  337. )
  339. # Add the relative bias matrix to the attention matrix
  340. attention_with_alibi = inputs + alibi_matrix
  341. # attention_with_alibi.shape = (batch_size, n_heads, n_channels, out_sequence_length, in_sequence_length)
  342. return attention_with_alibi
  345. class TokenWeightsLayer(tf.keras.layers.Layer):
  346. """
  347. Layer for computing token weights.
  349. Parameters
  350. ----------
  351. output_dim : int
  352. Dimension of the output.
  353. activation : str, optional
  354. Activation function to use.
  355. """
  357. def __init__(self, output_dim: int, activation: str = "linear", **kwargs):
  358. super().__init__(**kwargs)
  360. self.output_dim = output_dim
  361. self.dense_layer = tf.keras.layers.Dense(output_dim, activation=activation)
  362. self.activation_layer = tf.keras.layers.Activation(activation)
  363. self.norm_layer = tf.keras.layers.LayerNormalization()
  364. self.temperature = tf.Variable(0.0, trainable=False)
  366. def call(self, inputs, training=None, **kwargs):
  367. ell = self.activation_layer(self.dense_layer(inputs))
  368. ell = self.norm_layer(ell) / 0.1
  369. # Shape: (batch_size * n_channels, sequence_length, n_tokens)
  371. if training:
  372. # Sample from gumbel softmax parameterized by ell
  373. theta_sample = tf.argmax(ell, axis=2)
  374. theta_sample = tf.one_hot(theta_sample, self.output_dim)
  376. # Annealing
  377. theta_weight = tf.nn.softmax(ell, axis=2)
  378. # Shape: (batch_size * n_channels, sequence_length, n_tokens)
  380. token_weight = (
  381. self.temperature * theta_weight + (1 - self.temperature) * theta_sample
  382. )
  383. # Shape: (batch_size * n_channels, sequence_length, n_tokens)
  384. else:
  385. token_weight = tf.one_hot(tf.argmax(ell, axis=2), self.output_dim)
  386. return token_weight
  389. class PositionEmbedding(tf.keras.layers.Layer):
  390. """
  391. Layer for learning position embeddings.
  393. Parameters
  394. ----------
  395. sequence_length : int
  396. Sequence length.
  397. initializer : str, optional
  398. Initializer for the position embeddings.
  399. """
  401. def __init__(
  402. self,
  403. sequence_length: int,
  404. initializer: str = "glorot_uniform",
  405. **kwargs,
  406. ):
  407. super().__init__(**kwargs)
  408. self.sequence_length = sequence_length
  409. self.initializer = tf.keras.initializers.get(initializer)
  411. def build(self, inputs_shape):
  412. feature_size = inputs_shape[-1]
  413. self.position_embeddings = self.add_weight(
  414. name="embeddings",
  415. shape=[self.sequence_length, feature_size],
  416. initializer=self.initializer,
  417. trainable=True,
  418. )
  419. self.built = True
  421. def call(self, inputs, start_index=0):
  422. inputs_shape = tf.shape(inputs)
  423. feature_length = inputs_shape[-1]
  424. sequence_length = inputs_shape[-2]
  425. # trim to match the length of the input sequence, which might be less
  426. # than the sequence_length of the layer.
  427. position_embeddings = tf.convert_to_tensor(self.position_embeddings)
  428. position_embeddings = tf.slice(
  429. position_embeddings,
  430. (start_index, 0),
  431. (sequence_length, feature_length),
  432. )
  433. return tf.broadcast_to(position_embeddings, inputs_shape)
  436. class NormalizationLayer(tf.keras.layers.Layer):
  437. """
  438. Layer for performing normalization.
  440. Parameters
  441. ----------
  442. norm_type : str
  443. Type of normalization to perform.
  444. Options: "layer", "batch", "group".
  445. n_groups : int, optional
  446. Number of groups for group normalization.
  447. Required if norm_type is "group".
  448. """
  450. def __init__(self, norm_type: str = "layer", n_groups=None, **kwargs):
  451. super().__init__(**kwargs)
  452. if norm_type == "layer":
  453. self.norm_layer = tf.keras.layers.LayerNormalization()
  454. elif norm_type == "batch":
  455. self.norm_layer = tf.keras.layers.BatchNormalization()
  456. elif norm_type == "group":
  457. if n_groups is None:
  458. raise ValueError("n_groups must be specified for group normalization")
  459. self.norm_layer = tf.keras.layers.GroupNormalization(groups=n_groups)
  460. else:
  461. raise ValueError(f"Unknown normalization type: {norm_type}")
  463. def call(self, inputs, **kwargs):
  464. return self.norm_layer(inputs)
  467. class TimeAttentionLayer(tf.keras.layers.Layer):
  468. """
  469. Layer for performing time attention.
  471. Parameters
  472. ----------
  473. n_heads : int
  474. Number of heads.
  475. latent_sequence_length : int
  476. Sequence length of latent space.
  477. key_dim : int
  478. Key dimension.
  479. n_patches : int
  480. Number of patches to attend to.
  481. patch_length : int
  482. Patch length.
  483. unpatched_length : int
  484. Number of unpatched elements to attend to.
  485. transform_attention : str, optional
  486. Type of attention transformation to apply.
  487. """
  489. def __init__(
  490. self,
  491. n_heads: int,
  492. latent_sequence_length: int,
  493. key_dim: int,
  494. n_patches: int,
  495. patch_length: int,
  496. unpatched_length: int,
  497. transform_attention: str = None,
  498. **kwargs,
  499. ):
  500. super().__init__(**kwargs)
  501. self.n_heads = n_heads
  502. self.latent_sequence_length = latent_sequence_length
  503. self.key_dim = key_dim
  504. self.n_patches = n_patches
  505. self.patch_length = patch_length
  506. self.unpatched_length = unpatched_length
  507. self.transform_attention = transform_attention
  509. # Create a position embedding layer
  510. if transform_attention == "rope":
  511. self.rotary_position_embedding_layer = RotaryPositionEmbeddingLayer(
  512. embedding_dim=self.key_dim,
  513. max_sequence_length=tf.cast(
  514. tf.math.maximum(
  515. self.latent_sequence_length,
  516. self.n_patches + self.unpatched_length,
  517. ),
  518. dtype=tf.int32,
  519. ),
  520. )
  521. if transform_attention == "alibi":
  522. self.alibi_position_embedding_layer = ALiBiPositionEmbeddingLayer(
  523. n_heads=self.n_heads,
  524. n_patches=self.n_patches,
  525. patch_length=self.patch_length,
  526. unpatched_length=self.unpatched_length,
  527. )
  529. def call(self, inputs, mask=None, **kwargs):
  530. # q (Query): (batch_size, n_heads, out_sequence_length, out_n_channels, key_dim)
  531. # k (Key): (batch_size, n_heads, in_sequence_length, out_n_channels, key_dim)
  532. # v (Value): (batch_size, n_heads, out_sequence_length, out_n_channels, in_sequence_length, key_dim)
  533. q, k, v = inputs
  535. # Transpose inputs for time attention
  537. q = tf.transpose(q, perm=(0, 1, 3, 2, 4))
  538. # q: (batch_size, n_heads, out_n_channels, out_sequence_length, key_dim)
  540. k = tf.transpose(k, perm=(0, 1, 3, 2, 4))
  541. # k: (batch_size, n_heads, out_n_channels, in_sequence_length, key_dim)
  543. if self.transform_attention == "rope":
  544. # Apply rotary position embedding to the query and key vectors
  545. q = self.rotary_position_embedding_layer(q)
  546. k = self.rotary_position_embedding_layer(k)
  548. # Compute attention
  549. attention = tf.matmul(q, k, transpose_b=True) / tf.math.sqrt(
  550. tf.cast(self.key_dim, tf.float32)
  551. )
  552. # attention: (batch_size, n_heads, out_n_channels, out_sequence_length, in_sequence_length)
  554. if self.transform_attention == "alibi":
  555. # Apply ALiBi position embedding to the attention matrix
  556. attention = self.alibi_position_embedding_layer(attention, mask=mask)
  558. # Apply mask
  559. if mask is not None:
  560. attention += -1e9 * mask
  562. # Normalise attention with softmax
  563. attention = tf.nn.softmax(attention, axis=-1)
  565. if mask is not None:
  566. attention = attention * (1 - mask)
  568. # Apply attention to value
  569. attention = tf.expand_dims(
  570. tf.transpose(attention, perm=(0, 1, 3, 2, 4)), axis=-2
  571. )
  572. # attention: (batch_size, n_heads, out_sequence_length, out_n_channels, 1, in_sequence_length)
  574. output = tf.matmul(attention, v)
  575. # output: (batch_size, n_heads, out_sequence_length, out_n_channels, 1, key_dim)
  577. output = tf.squeeze(output, axis=-2)
  578. # output: (batch_size, n_heads, out_sequence_length, out_n_channels, key_dim)
  580. return output
  583. class ChannelAttention(tf.keras.layers.Layer):
  584. """
  585. Layer for performing channel attention.
  587. Parameters
  588. ----------
  589. key_dim : int
  590. Key dimension.
  591. """
  593. def __init__(self, key_dim: int, **kwargs):
  594. super().__init__(**kwargs)
  595. self.key_dim = key_dim
  597. def call(self, inputs, mask=None, **kwargs):
  598. # q (Query): (batch_size, n_heads, out_sequence_length, out_n_channels, key_dim)
  599. # k (Key): (batch_size, n_heads, out_sequence_length, in_n_channels, key_dim)
  600. # v (Value): (batch_size, n_heads, in_sequence_length, in_n_channels, key_dim)
  601. q, k, v = inputs
  603. batch_size = tf.shape(q)[0]
  604. n_heads = tf.shape(q)[1]
  605. out_sequence_length = tf.shape(q)[2]
  606. out_n_channels = tf.shape(q)[3]
  607. in_n_channels = tf.shape(k)[3]
  608. in_sequence_length = tf.shape(v)[2]
  610. # Compute attention
  611. attention = tf.matmul(q, k, transpose_b=True) / tf.math.sqrt(
  612. tf.cast(self.key_dim, tf.float32)
  613. )
  614. # attention: (batch_size, n_heads, out_sequence_length, out_n_channels, in_n_channels)
  616. # Apply mask
  617. if mask is not None:
  618. attention += -1e9 * mask
  620. # Normalise attention with softmax
  621. attention = tf.nn.softmax(attention, axis=-1)
  623. if mask is not None:
  624. attention = attention * (1 - mask)
  626. # Apply attention to value
  627. v = tf.transpose(v, perm=(0, 1, 3, 2, 4))
  628. # v: (batch_size, n_heads, in_n_channels, in_sequence_length, key_dim)
  630. v = tf.reshape(v, shape=(batch_size, n_heads, 1, in_n_channels, -1))
  631. # v: (batch_size, n_heads, 1, in_n_channels, in_sequence_length * key_dim)
  633. output = tf.matmul(attention, v)
  634. # output: (batch_size, n_heads, out_sequence_length, out_n_channels, in_sequence_length * key_dim)
  636. output = tf.reshape(
  637. output,
  638. shape=(
  639. batch_size,
  640. n_heads,
  641. out_sequence_length,
  642. out_n_channels,
  643. in_sequence_length,
  644. self.key_dim,
  645. ),
  646. )
  647. # output: (batch_size, n_heads, out_sequence_length, out_n_channels, in_sequence_length, key_dim)
  649. return output
  652. class PASSTALayer(tf.keras.layers.Layer):
  653. """
  654. The Perceiver AR Separable Space-Time self-Attention (PASSTA) layer.
  655. This layer performs space-time attention on the input tensor.
  657. Parameters
  658. ----------
  659. n_heads : int
  660. Number of heads.
  661. n_channels : int
  662. Number of channels.
  663. latent_sequence_length : int
  664. Sequence length of latent space.
  665. n_patches : int
  666. Number of patches to attend to.
  667. patch_length : int
  668. Patch length.
  669. unpatched_length : int
  670. Number of unpatched elements to attend to.
  671. key_dim : int
  672. Key dimension.
  673. pos_embedding_type : str
  674. Type of positional embedding to use.
  675. channel_attention_dropout : float
  676. Dropout rate for channel attention.
  677. Values greater than 1.0 means no channel attention.
  678. Values less than 0.0 means no dropout.
  679. within_channel_attention_dropout : float
  680. Dropout rate for within-channel attention.
  681. Values greater than 1.0 means no within-channel attention.
  682. Values less than 0.0 means no dropout.
  683. """
  685. def __init__(
  686. self,
  687. n_heads: int,
  688. n_channels: int,
  689. latent_sequence_length: int,
  690. n_patches: int,
  691. patch_length: int,
  692. unpatched_length: int,
  693. key_dim: int,
  694. pos_embedding_type: str,
  695. channel_attention_dropout: float,
  696. within_channel_attention_dropout: float,
  697. **kwargs,
  698. ):
  699. super().__init__(**kwargs)
  700. self.n_heads = n_heads
  701. self.n_channels = n_channels
  702. self.latent_sequence_length = latent_sequence_length
  703. self.key_dim = key_dim
  704. self.pos_embedding_type = pos_embedding_type
  705. self.n_patches = n_patches
  706. self.patch_length = patch_length
  707. self.unpatched_length = unpatched_length
  709. # Time attention layer
  710. self.time_attention_layer = TimeAttentionLayer(
  711. n_heads,
  712. latent_sequence_length,
  713. key_dim,
  714. n_patches,
  715. patch_length,
  716. unpatched_length,
  717. pos_embedding_type,
  718. )
  719. # Mask for time attention (This is fixed).
  720. self.time_attention_mask = self._compute_time_attention_mask()
  722. # Channel attention layer
  723. self.channel_attention_layer = ChannelAttention(key_dim)
  725. # Channel attention dropouts
  726. self.channel_attention_dropout = tf.Variable(
  727. channel_attention_dropout, trainable=False
  728. )
  729. self.within_channel_attention_dropout = tf.Variable(
  730. within_channel_attention_dropout, trainable=False
  731. )
  733. def _compute_time_attention_mask(self) -> tf.Tensor:
  734. """
  735. Compute the mask for time attention.
  737. Returns
  738. -------
  739. mask : tf.Tensor
  740. Mask for time attention.
  741. Shape: (latent_sequence_length, n_patches + unpatched_length).
  742. """
  743. mask = np.zeros(
  744. (self.latent_sequence_length, self.n_patches + self.unpatched_length)
  745. )
  747. # Patch masking
  748. for i in range(self.n_patches):
  749. m_indx = max(
  750. 0,
  751. (i + 1 - self.n_patches) * self.patch_length
  752. + self.latent_sequence_length
  753. - 1,
  754. )
  755. mask[:m_indx, i] = 1
  757. # Unpatched masking
  758. for i in range(self.n_patches, self.n_patches + self.unpatched_length):
  759. mask[
  760. : self.latent_sequence_length
  761. - (self.n_patches + self.unpatched_length)
  762. + i,
  763. i,
  764. ] = 1
  765. mask = tf.constant(mask, dtype=tf.float32)
  767. return mask
  769. def _compute_channel_attention_mask(self, training: bool) -> Union[tf.Tensor, None]:
  770. """
  771. Compute the mask for channel attention.
  773. Parameters
  774. ----------
  775. training : bool
  776. Whether the model is training.
  777. If False, no dropout is applied.
  779. Returns
  780. -------
  781. mask : tf.Tensor or None
  782. Mask for channel attention.
  783. Shape: (n_channels, n_channels).
  784. If None, full attention is applied.
  785. """
  786. if self.channel_attention_dropout < 1.0:
  787. # We apply channel attention dropout only during training
  788. if not training:
  789. return tf.zeros((self.n_channels, self.n_channels), dtype=tf.float32)
  790. else:
  791. uniform_sampler = tfp.distributions.Uniform()
  793. # Sample whether to apply channel attention
  794. if uniform_sampler.sample() < self.channel_attention_dropout:
  795. # Does not apply channel attention and mask all off-diagonal elements
  796. mask = 1 - np.eye(self.n_channels)
  797. return tf.constant(mask, dtype=tf.float32)
  799. # Sample whether to apply within-channel attention
  800. elif uniform_sampler.sample() < self.within_channel_attention_dropout:
  801. # Does not apply within-channel attention and mask all diagonal elements
  802. mask = np.eye(self.n_channels)
  803. return tf.constant(mask, dtype=tf.float32)
  805. # Apply channel attention and does not mask any elements, return None
  806. else:
  807. return tf.zeros(
  808. (self.n_channels, self.n_channels), dtype=tf.float32
  809. )
  810. else:
  811. # If channel_attention_dropout >= 1.0, no channel attention is applied
  812. # Mask all off-diagonal elements
  813. mask = 1 - np.eye(self.n_channels)
  814. return tf.constant(mask, dtype=tf.float32)
  816. def call(self, inputs, training=None, **kwargs):
  817. # ---------- Unpack Inputs ---------- #
  819. # q (Query): (batch_size, n_heads, latent_sequence_length, out_n_channels, key_dim)
  820. # k (Key): (batch_size, n_heads, n_patches + unpatched_length, out_n_channels, key_dim)
  821. # v (Value): (batch_size, n_heads, n_patches + unpatched_length, in_n_channels, key_dim)
  822. # c_q (Channel Query): (batch_size, n_heads, latent_sequence_length, out_n_channels, key_dim)
  823. # c_k (Channel Key): (batch_size, n_heads, latent_sequence_length, in_n_channels, key_dim)
  824. q, k, v, c_q, c_k = inputs
  826. # ---------- Channel Attention ---------- #
  828. # First sample channel attention mask
  829. channel_attention_mask = self._compute_channel_attention_mask(training=training)
  831. # Apply channel attention
  832. output = self.channel_attention_layer(
  833. [c_q, c_k, v], mask=channel_attention_mask, training=training, **kwargs
  834. )
  836. # ---------- Time Attention ---------- #
  837. output = self.time_attention_layer(
  838. [q, k, output], mask=self.time_attention_mask, training=training, **kwargs
  839. )
  841. return output
  844. class MultiHeadPASSTALayer(tf.keras.layers.Layer):
  845. """
  846. The Multi-head PASSTA layer.
  848. Parameters
  849. ----------
  850. n_heads : int
  851. Number of heads.
  852. model_dim : int
  853. Model dimension.
  854. n_channels : int
  855. Number of channels.
  856. sequence_length : int
  857. Sequence length.
  858. latent_sequence_length : int
  859. Latent sequence length.
  860. n_patches : int
  861. Number of patches to attend to.
  862. patch_length : int
  863. Patch length.
  864. unpatched_length : int
  865. Number of unpatched elements to attend to.
  866. pos_embedding_type : str
  867. Type of positional embedding to use.
  868. channel_attention_dropout : float
  869. Dropout rate for channel attention.
  870. Values greater than 1.0 means no channel attention.
  871. Values less than 0.0 means no dropout.
  872. within_channel_attention_dropout : float
  873. Dropout rate for within-channel attention.
  874. Values greater than 1.0 means no within-channel attention.
  875. Values less than 0.0 means no dropout
  876. """
  878. def __init__(
  879. self,
  880. n_heads: int,
  881. model_dim: int,
  882. n_channels: int,
  883. sequence_length: int,
  884. latent_sequence_length: int,
  885. n_patches: int,
  886. patch_length: int,
  887. unpatched_length: int,
  888. pos_embedding_type: str,
  889. channel_attention_dropout: float,
  890. within_channel_attention_dropout: float,
  891. **kwargs,
  892. ):
  893. super().__init__(**kwargs)
  894. self.n_heads = n_heads
  895. self.model_dim = model_dim
  896. self.key_dim = model_dim // n_heads
  897. self.n_channels = n_channels
  898. self.sequence_length = sequence_length
  899. self.latent_sequence_length = latent_sequence_length
  900. self.n_patches = n_patches
  901. self.patch_length = patch_length
  902. self.unpatched_length = unpatched_length
  903. self.pos_embedding_type = pos_embedding_type
  904. self.channel_attention_dropout = channel_attention_dropout
  905. self.within_channel_attention_dropout = within_channel_attention_dropout
  907. # Patch projection
  908. self.patch_projection = tf.keras.layers.Dense(n_heads)
  910. # Input projections
  911. self.time_patched_projection = tf.keras.layers.Dense(2 * self.model_dim)
  912. self.time_unpatched_projection = tf.keras.layers.Dense(2 * self.model_dim)
  913. self.time_query_projection = tf.keras.layers.Dense(self.model_dim)
  914. self.channel_projection = tf.keras.layers.Dense(2 * self.model_dim)
  916. # PASSTA layer for time and channel attention
  917. self.passta_layer = PASSTALayer(
  918. n_heads,
  919. n_channels,
  920. latent_sequence_length,
  921. n_patches,
  922. patch_length,
  923. unpatched_length,
  924. self.key_dim,
  925. pos_embedding_type,
  926. channel_attention_dropout,
  927. within_channel_attention_dropout,
  928. )
  930. # Output projection
  931. self.output_projection = tf.keras.layers.Dense(self.model_dim)
  933. def _patch_x(self, x: tf.Tensor) -> tf.Tensor:
  934. # x.shape: (batch_size, sequence_length, n_channels, model_dim)
  936. x = tf.transpose(x, perm=(0, 2, 3, 1))
  937. # x.shape: (batch_size, n_channels, model_dim, sequence_length)
  939. x = tf.reshape(
  940. x,
  941. (
  942. tf.shape(x)[0],
  943. self.n_channels,
  944. self.model_dim,
  945. self.n_patches,
  946. self.patch_length,
  947. ),
  948. )
  949. # x.shape: (batch_size, n_channels, model_dim, n_patches, patch_length)
  951. x = tf.transpose(x, perm=(0, 3, 1, 2, 4))
  952. # x.shape: (batch_size, n_patches, n_channels, model_dim, patch_length)
  954. return x
  956. def _perceiver_x(self, x: tf.Tensor) -> tf.Tensor:
  957. """Get the last latent_sequence_length elements of the sequence."""
  958. # x.shape: (batch_size, sequence_length, n_channels, model_dim)
  960. x = tf.slice(
  961. x,
  962. [0, self.sequence_length - self.latent_sequence_length, 0, 0],
  963. [-1, -1, -1, -1],
  964. )
  965. # x.shape: (batch_size, latent_sequence_length, n_channels, model_dim)
  967. return x
  969. def _unpatch_x(self, x: tf.Tensor) -> tf.Tensor:
  970. """Get the last unpatched_length elements of the sequence."""
  971. # x.shape: (batch_size, sequence_length, n_channels, model_dim)
  973. x = tf.slice(
  974. x,
  975. [0, self.sequence_length - self.unpatched_length, 0, 0],
  976. [-1, -1, -1, -1],
  977. )
  978. # x.shape: (batch_size, unpatched_length, n_channels, model_dim)
  980. return x
  982. def _split_heads(self, x: tf.Tensor) -> tf.Tensor:
  983. # x.shape: (batch_size, time_length, n_channels, model_dim)
  984. # Here time_length is either n_patches + unpatched_length or latent_sequence_length
  986. x = tf.reshape(
  987. x,
  988. (
  989. tf.shape(x)[0],
  990. tf.shape(x)[1],
  991. self.n_channels,
  992. self.n_heads,
  993. self.key_dim,
  994. ),
  995. )
  996. # x.shape: (batch_size, time_length, n_channels, n_heads, key_dim)
  998. x = tf.transpose(x, perm=(0, 3, 1, 2, 4))
  999. # x.shape: (batch_size, n_heads, time_length, n_channels, key_dim)
  1001. return x
  1003. def _combine_heads(self, x: tf.Tensor) -> tf.Tensor:
  1004. # x.shape: (batch_size, n_heads, latent_sequence_length, n_channels, key_dim)
  1006. x = tf.transpose(x, perm=(0, 2, 3, 1, 4))
  1007. # x.shape: (batch_size, latent_sequence_length, n_channels, n_heads, key_dim)
  1009. x = tf.reshape(
  1010. x,
  1011. (
  1012. tf.shape(x)[0],
  1013. self.latent_sequence_length,
  1014. self.n_channels,
  1015. self.model_dim,
  1016. ),
  1017. )
  1018. # x.shape: (batch_size, latent_sequence_length, n_channels, model_dim)
  1020. return x
  1022. def call(self, inputs, training=None, **kwargs):
  1023. x = inputs
  1024. # x.shape: (batch_size, sequence_length, n_channels, model_dim)
  1026. # ---------- Process inputs ---------- #
  1027. # Input is processed into 3 parts:
  1028. # 1. Patched input: patched_x
  1029. # 2. Unpatched input: unpatched_x
  1030. # 3. Perceiver input: perceiver_x
  1032. patched_x = self._patch_x(x)
  1033. # patched_x.shape: (batch_size, n_patches, n_channels, model_dim, patch_length)
  1035. patched_x = tf.reshape(
  1036. patched_x,
  1037. (
  1038. tf.shape(patched_x)[0],
  1039. self.n_patches,
  1040. self.n_channels,
  1041. self.key_dim,
  1042. self.n_heads,
  1043. self.patch_length,
  1044. ),
  1045. )
  1046. # patched_x.shape: (batch_size, n_patches, n_channels, key_dim, n_heads, patch_length)
  1047. patched_x = tf.reshape(
  1048. patched_x,
  1049. (
  1050. tf.shape(patched_x)[0],
  1051. self.n_patches,
  1052. self.n_channels,
  1053. self.key_dim,
  1054. self.n_heads * self.patch_length,
  1055. ),
  1056. )
  1057. # patched_x.shape: (batch_size, n_patches, n_channels, key_dim, n_heads * patch_length)
  1059. patched_x = self.patch_projection(patched_x)
  1060. # patched_x.shape: (batch_size, n_patches, n_channels, key_dim, n_heads)
  1062. patched_x = tf.reshape(
  1063. patched_x,
  1064. (
  1065. tf.shape(patched_x)[0],
  1066. self.n_patches,
  1067. self.n_channels,
  1068. self.model_dim,
  1069. ),
  1070. )
  1071. # patched_x.shape: (batch_size, n_patches, n_channels, model_dim)
  1073. perceiver_x = self._perceiver_x(x)
  1074. # perceiver_x.shape: (batch_size, latent_sequence_length, n_channels, model_dim)
  1076. unpatched_x = self._unpatch_x(x)
  1077. # unpatched_x.shape: (batch_size, unpatched_length, n_channels, model_dim)
  1079. # ---------- Project inputs to Q, K, V ---------- #
  1080. q = self.time_query_projection(perceiver_x)
  1081. # q.shape: (batch_size, latent_sequence_length, n_channels, model_dim)
  1083. k_patched, v_patched = tf.split(
  1084. self.time_patched_projection(patched_x), 2, axis=-1
  1085. )
  1086. # k_patched.shape: (batch_size, n_patches, n_channels, model_dim)
  1087. # v_patched.shape: (batch_size, n_patches, n_channels, model_dim)
  1089. k_unpatched, v_unpatched = tf.split(
  1090. self.time_unpatched_projection(unpatched_x), 2, axis=-1
  1091. )
  1092. # k_unpatched.shape: (batch_size, unpatched_length, n_channels, model_dim)
  1093. # v_unpatched.shape: (batch_size, unpatched_length, n_channels, model_dim)
  1095. # Concatenate k and v for time attention
  1096. k = tf.concat([k_patched, k_unpatched], axis=1)
  1097. v = tf.concat([v_patched, v_unpatched], axis=1)
  1098. # k.shape: (batch_size, n_patches + unpatched_length, n_channels, model_dim)
  1099. # v.shape: (batch_size, n_patches + unpatched_length, n_channels, model_dim)
  1101. c_q, c_k = tf.split(self.channel_projection(perceiver_x), 2, axis=-1)
  1102. # c_q.shape: (batch_size, latent_sequence_length, n_channels, model_dim)
  1103. # c_k.shape: (batch_size, latent_sequence_length, n_channels, model_dim)
  1105. # ---------- Split heads ---------- #
  1107. q = self._split_heads(q)
  1108. # (batch_size, n_heads, latent_sequence_length, n_channels, key_dim)
  1110. k = self._split_heads(k)
  1111. # (batch_size, n_heads, n_patches + unpatched_length, n_channels, key_dim)
  1113. v = self._split_heads(v)
  1114. # (batch_size, n_heads, n_patches + unpatched_length, n_channels, key_dim)
  1116. c_q = self._split_heads(c_q)
  1117. # (batch_size, n_heads, latent_sequence_length, n_channels, key_dim)
  1119. c_k = self._split_heads(c_k)
  1120. # (batch_size, n_heads, latent_sequence_length, n_channels, key_dim)
  1122. # ---------- PASSTA Layer ---------- #
  1123. output = self.passta_layer([q, k, v, c_q, c_k], training=training, **kwargs)
  1124. # output.shape: (batch_size, n_heads, latent_sequence_length, n_channels, key_dim)
  1126. # ---------- Combine heads ---------- #
  1127. output = self._combine_heads(output)
  1128. # output.shape: (batch_size, latent_sequence_length, n_channels, model_dim)
  1130. # ---------- Output projection ---------- #
  1131. output = self.output_projection(output)
  1132. # output.shape: (batch_size, latent_sequence_length, n_channels, model_dim)
  1134. return output

layers.py at commit ffbeff3, under MIT · at the source

Overview

Authors: Rukuang Huang1,2, SungJun Cho1,3, Chetan Gohil1,2, Oiwi Parker Jones1,4, Mark Woolrich1,2
ORCID iDs: Chetan Gohil
  1. Oxford Centre for Integrative Neuroimaging (OxCIN), University of Oxford, Oxford, United Kingdom
  2. Department of Psychiatry, University of Oxford, Oxford, United Kingdom
  3. Nuffield Department of Clinical Neurosciences, University of Oxford, Oxford, United Kingdom
  4. Department of Engineering Science, University of Oxford, Oxford, United Kingdom
Institutions: University of Oxford (United Kingdom); Wellcome Centre for Integrative Neuroimaging (United Kingdom)
Journal: Imaging neuroscience (Cambridge, Mass.), volume 4, article IMAG.a.1301
Dates: received 10 January 2026; accepted 16 June 2026; published online 24 July 2026
Type: Research article · Language: English
License: CC BY
Identifiers: DOI 10.1162/imag.a.1301 · PMID 42516171 · PMCID PMC13403652 · OpenAlex W4416055129
Open access: diamond, a free copy (OpenAlex)
Preprint: arxiv:2510.18080
Status: code verified
Categories: MEG (modality), methods / tools (subfield)
Methods: Spectral & time-frequency, Preprocessing, Connectivity, Smoothing, state filtering, decompositions, Machine learning
Keywords: electrophysiology, MEG, GPT, transformer, foundation model, tokenisation
Topic: EEG and Brain-Computer Interfaces (Cognitive Neuroscience, Neuroscience), according to OpenAlex
Funding: Wellcome Trust (203139/Z/16/Z, 106183/Z/14/Z, 203139/A/16/Z, 215573/Z/19/Z); National Institute for Health Research (NIHR) (NIHR203316)
Citations: not cited yet (Europe PMC); 71 references in the paper

Abstract

Modelling the complex spatio-temporal patterns of large-scale brain dynamics is crucial for neuroscience, but traditional methods fail to capture the rich structure in modalities such as magnetoencephalography (MEG). Recent advances in deep learning have enabled significant progress in other domains, such as language and vision, by using foundation models at scale. Here, we introduce MEG-GPT, a transformer-based foundation model that uses time-attention and next time-point prediction. To facilitate this, we also introduce a novel data-driven tokeniser for continuous MEG data, which preserves the high temporal resolution of continuous MEG signals without lossy transformations. We trained MEG-GPT on tokenised brain region time courses extracted from a large-scale MEG dataset (N = 612, eyes-closed rest, Cam-CAN data), and show that the learnt model can generate data with realistic spatio-spectral properties, including transient events and population variability. Critically, it performs well in downstream decoding tasks, improving downstream supervised prediction task, showing improved zero-shot generalisation across sessions (improving accuracy from 0.56 to 0.59) and subjects (improving accuracy from 0.45 to 0.49) compared with a PCA baseline method. Furthermore, we show the model can be efficiently fine-tuned on a smaller labelled dataset to boost performance in cross-subject decoding scenarios. This work establishes a powerful foundation model for electrophysiological data, paving the way for applications in computational neuroscience and neural decoding.

Reproduced under the paper's license (CC BY), from the paper cited above.

Repositories

Its files are read in the Code ↔ Paper reader above, with 13 matches between paragraphs and lines of code.

OHBA-analysis/osl-foundation

License: MIT
State: the link answers, verified on 27 September 2026
Evidence: files inventoried
Commit: ffbeff318811de3adea33d7dfeaa0af15a5b902b, 8 July 2026
Languages: Python (80)
Size: 103 files, 80 scripts
Software Heritage: not archived
Found in: “Data and Code Availability”
Holds: README, license file, environment (pyproject.toml)
Not found: CITATION.cff, tests, continuous integration, documentation
Tools: NumPy (45 files), Matplotlib (17 files), TensorFlow (15 files), pandas (6 files), seaborn (6 files), scikit-learn (4 files), MNE-Python (3 files), SciPy (2 files), NiBabel (1 file), Nilearn (1 file), UMAP (1 file)
Availability: 1 check, the latest on 27 September 2026: the link answers
  • 27 September 2026: the link answers
82 files
At the source: github.com/OHBA-analysis/osl-foundation

Zenodo 11099418

License: CC-BY-4.0
State: the link answers, verified on 27 September 2026
Evidence: files inventoried
Size: 9 files
Software Heritage: not checked
Found in: the references
Not found: README, license file, CITATION.cff, environment file, tests, continuous integration, documentation
Availability: 1 check, the latest on 27 September 2026: the link answers (HTTP 200)
  • 27 September 2026: the link answers (HTTP 200)
At the source: zenodo.org/records/11099418

The paper's code and data availability statement is in the Data section.

Tracing map

Proposed by the machine: these links were found in the paper and verified at the source, without human review. The map will receive a Zenodo DOI once one of the paper's authors has validated it with their ORCID.

What the map holds:

  • 2 repositories of the authors' code, each at its verified commit, with its license and how the link was found in the paper;
  • 80 scripts, each with its path and the digest of its content;
  • 13 matches between paragraphs of the paper and lines of the code (method lexical-v1);
  • neither the text of the paper nor the code itself.

Its JSON (tracing-map.json) is deposited on Zenodo with its DOI once the map is validated.

Data

No dataset and no data link were found in the paper.

Data and Code Availability

Data used are publicly available. For the Wakeman–Henson dataset, we refer the readers to the original paper (Wakeman & Henson, 2015). For the Cam-CAN dataset, we refer the readers to the original paper (Taylor et al., 2017).

Source code and scripts for reproducing results in the paper using TensorFlow are available on GitHub: https://github.com/OHBA-analysis/osl-foundation. Example code and tutorials for training a foundation model and applying it to new data are also provided.

A PyTorch implemention of the tokenizer is available here: https://github.com/OHBA-analysis/EphysTokenizer, and a PyTorch implementation of MEG-GPT is available here: https://github.com/OHBA-analysis/MEG-GPT.

Reproduced under the paper's license (CC BY), from the paper cited above.

Versions

The history of this record: each version stored by the harvester or made by a correction of its authors or of the maintainers of its code, and what changed in its facts. The texts of the paper (its abstract, its availability statements) are not part of it; versions that changed only those are not listed.

Version 1, 27 September 2026: the first record

Recorded: type, language, journal, volume, pages, dates, 5 authors, 6 keywords, 2 funders, 64 references.

Cite

This paper

Huang, R., Cho, S., Gohil, C., Jones, O. P., & Woolrich, M. (2026). MEG-GPT: A transformer-based foundation model for magnetoencephalography data. Imaging neuroscience (Cambridge, Mass.), 4, IMAG.a.1301. https://doi.org/10.1162/imag.a.1301

BibTeX

@article{huang2026meg,
author = {Huang, Rukuang and Cho, SungJun and Gohil, Chetan and Jones, Oiwi Parker and Woolrich, Mark},
title = {{MEG-GPT: A transformer-based foundation model for magnetoencephalography data}},
journal = {Imaging neuroscience (Cambridge, Mass.)},
year = {2026},
month = jul,
volume = {4},
pages = {IMAG.a.1301},
publisher = {MIT Press},
issn = {2837-6056},
doi = {10.1162/imag.a.1301},
url = {https://doi.org/10.1162/imag.a.1301},
pmid = {42516171},
pmcid = {PMC13403652}
}

RIS

TY - JOUR
AU - Huang, Rukuang
AU - Cho, SungJun
AU - Gohil, Chetan
AU - Jones, Oiwi Parker
AU - Woolrich, Mark
TI - MEG-GPT: A transformer-based foundation model for magnetoencephalography data
T2 - Imaging neuroscience (Cambridge, Mass.)
J2 - Imaging Neurosci (Camb)
PY - 2026
DA - 2026/07/24
VL - 4
SP - IMAG.a.1301
SN - 2837-6056
PB - MIT Press
DO - 10.1162/imag.a.1301
UR - https://doi.org/10.1162/imag.a.1301
LA - en
ER -

CSL-JSON

{
"id": "10.1162/imag.a.1301",
"type": "article-journal",
"title": "MEG-GPT: A transformer-based foundation model for magnetoencephalography data",
"container-title": "Imaging neuroscience (Cambridge, Mass.)",
"author": [
{
"family": "Huang",
"given": "Rukuang"
},
{
"family": "Cho",
"given": "SungJun"
},
{
"family": "Gohil",
"given": "Chetan"
},
{
"family": "Jones",
"given": "Oiwi Parker"
},
{
"family": "Woolrich",
"given": "Mark"
}
],
"container-title-short": "Imaging Neurosci (Camb)",
"volume": "4",
"page": "IMAG.a.1301",
"DOI": "10.1162/imag.a.1301",
"PMID": "42516171",
"PMCID": "PMC13403652",
"ISSN": "2837-6056",
"publisher": "MIT Press",
"URL": "https://doi.org/10.1162/imag.a.1301",
"language": "en",
"issued": {
"date-parts": [
[
2026,
7,
24
]
]
}
}

The tracing map gets a citation of its own once an author has validated it and it has a DOI.

Similar papers

The papers with a page that share the most with this one: the tools found in their code, their categories, datasets, cited references and authors, the rarest counting most.

[1] doi:10.1162/imag.a.1188 [code]
Modelling variability in functional brain networks using embeddings.
Journal: Imaging neuroscience (Cambridge, Mass.)
In common: NiBabel, seaborn, scikit-learn, 4 other tools, 15 references
[2] doi:10.1162/imag.a.1190 [code]
Canonical Hidden Markov Model Networks for studying M/EEG.
Journal: Imaging neuroscience (Cambridge, Mass.)
In common: MNE-Python, Nilearn, NiBabel, 5 other tools, MEG, 9 references, author Chetan Gohil
[3] doi:10.1162/imag.a.1237 [code]
Modelling discrete states and long-term dynamics in functional brain networks.
Journal: Imaging neuroscience (Cambridge, Mass.)
In common: MNE-Python, TensorFlow, seaborn, 5 other tools, MEG, 8 references, author Chetan Gohil
[4] doi:10.1002/hbm.70516 [code]
Effects of Age on Resting-State Cortical Networks.
Journal: Human brain mapping
In common: MNE-Python, Nilearn, NiBabel, 5 other tools, MEG, 8 references, author Chetan Gohil
[5] doi:10.1038/s41531-026-01372-1 [code]
Varying patterns of association between cortical large-scale networks and subthalamic nucleus activity in Parkinson's disease.
Journal: NPJ Parkinson's disease
In common: MNE-Python, Nilearn, NiBabel, 6 other tools, 8 references
[6] doi:10.1162/imag.a.1269 [code]
From early to contemporary normative modeling: Mapping individual differences in neurophysiological signals.
Journal: Imaging neuroscience (Cambridge, Mass.)
In common: MNE-Python, Nilearn, NiBabel, 6 other tools, MEG, 3 references
[7] doi:10.1093/braincomms/fcag236 [code]
Dynamic, state-dependent characteristics of cognitive fluctuations in Lewy body dementia: a magnetoencephalography study.
Journal: Brain communications
In common: MNE-Python, pandas, NumPy, MEG, 7 references
[8] doi:10.7554/elife.100605 [code]
Age-related changes in ‘cortical’ 1/f dynamics are linked to cardiac activity
Journal: —
In common: MNE-Python, seaborn, scikit-learn, 4 other tools, MEG, 5 references
[9] doi:10.1038/s41597-026-07350-9 [code]
An open multi-center MEG-EEG dataset for studying conscious visual perception.
Journal: Scientific data
In common: MNE-Python, Nilearn, NiBabel, 6 other tools, MEG, 2 references
[10] doi:10.1093/nc/niag029 [code]
A data-driven approach to identifying and evaluating connectivity-based neural correlates of conscious visual perception.
Journal: Neuroscience of consciousness
In common: MNE-Python, Nilearn, NiBabel, 6 other tools, MEG, 1 reference

Contribute

The authors of this paper can claim it, correct its record and validate its tracing map, and the maintainers of its code (its owner, or a public member of its organization) correct what it says of their repository; anyone signed in can ask for its removal. Every request goes to OSCR's own machine, which answers it; your account page follows them.

Sign in with ORCID to claim this paper as one of its authors, correct its record or validate its tracing map: when the paper's metadata lists your ORCID iD, you are recognized at once. Maintainers of its code: sign in with GitHub, then claim the repository on your account page.

Request its removal

To ask OSCR to remove this record, the copies of its authors' scripts or its tracing map, use the removal request page: signed in, you say who you are, what to remove and why, then review and confirm the request. Published rules decide every request (how).

Discussion, reproductions, activity

Discussion: questions and error reports about this paper and its code, from signed-in readers and its authors. It opens with sign-in.

Reproductions: reports from readers who ran the authors' code: what they reproduced, with which environment, commit and data. It opens with sign-in.

Activity: what happens around this paper: new versions of its record, its map's validation, discussions and reproductions. It opens with sign-in.