OSCR

TokenUNet: A new case for transformers integration in efficient and interpretable 3D UNets for brain imaging segmentation.

Code ↔ Paper

7 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 7 matches
  1. [1] § Materials and methods ↔ code/custom/architectures/tokenunet.py, lines 415–529 · score 0.78 · skip connections, MLP Mixer, Transformer encoder, token processing, TokenUNet, bottlenecks
  2. [2] § Materials and methods › Architectures and modules › TokenLearner and TokenFuser. ↔ code/custom/architectures/tokenunet.py, lines 209–268 · score 0.76 · Linear layer, spatial attention masks, feature map, logit, sigmoid, soft
  3. [3] § Materials and methods › Training ↔ code/scripts/profiling_archs.py, lines 333–413 · score 0.74 · deep supervision, Nesterov, mirroring, momentum, nnUNet, SGD
  4. [4] § Materials and methods › Training ↔ code/custom/nnUNetTokenUNetTrainer.py, lines 175–236 · score 0.65 · deep supervision, nnUNet, schedule, SGD, augmented, batch
  5. [5] § Materials and methods › Architectures and modules › UNet. ↔ code/custom/architectures/tokenunet.py, lines 53–84 · score 0.64 · LeakyReLU, trilinear, CONV, nonlinear, volumes, upsampling
  6. [6] § Results › What TokenLearner learns ↔ code/scripts/dice_attention.py, lines 1–24 · score 0.62 · Dice score, spatial attention maps, TokenFuser, TokenLearner, alignment, fold
  7. [7] § Materials and methods › Data ↔ code/scripts/convert_fets_to_nnunet.py, lines 60–85 · score 0.57 · FeTS, BraTS, FLAIR, modalities

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 · 529 lines · 19 KB · no license · 3 matches

  1. import torch
  2. import torch.nn as nn
  3. import math
  5. from .tokenmixers import *
  7. __all__ = [
  8. "Block",
  9. "BlockUp",
  10. "Stage",
  11. "StageUp",
  12. "CNNEnc",
  13. "CNNDec",
  14. "SpatialAttentionMaskMaker3d",
  15. "TokenLearner3d",
  16. "TokenFuser3d",
  17. "TokenUNet"
  18. ]
  20. class Block(nn.Module):
  21. def __init__(
  22. self,
  23. in_channels,
  24. out_channels,
  25. kernel_size,
  26. p=0.0,
  27. downsample=False
  28. ):
  29. super().__init__()
  31. self.block = nn.Sequential(
  32. nn.InstanceNorm3d(in_channels),
  33. nn.Conv3d(in_channels, in_channels*2, kernel_size, stride=2 if downsample else 1, padding=1),
  34. nn.LeakyReLU(0.1),
  35. nn.InstanceNorm3d(in_channels*2),
  36. nn.Conv3d(in_channels*2, out_channels, kernel_size=1,padding="same"),
  37. )
  38. nn.init.kaiming_normal_(self.block[1].weight, mode='fan_in', nonlinearity='relu')
  39. nn.init.kaiming_normal_(self.block[-1].weight, mode='fan_in', nonlinearity='relu')
  41. self.res_alpha = nn.Parameter(torch.tensor([-2.5]))
  42. adjust_volume = nn.AvgPool3d(2,2) if downsample else nn.Identity()
  43. adjust_channels = nn.Conv3d(in_channels, out_channels, 1) if in_channels != out_channels else nn.Identity()
  44. self.adjust = nn.Sequential(adjust_volume, adjust_channels)
  46. def forward(self, x):
  47. res = self.block(x)
  48. alpha = torch.sigmoid(self.res_alpha)
  49. x = self.adjust(x)
  50. x = (1.-alpha)*x + alpha*res
  51. return x
  53. class BlockUp(nn.Module):
  54. def __init__(
  55. self,
  56. in_channels,
  57. out_channels,
  58. kernel_size,
  59. p=0.0,
  60. upsample=False
  61. ):
  62. super().__init__()
  64. self.block = nn.Sequential(
  65. nn.InstanceNorm3d(in_channels),
  66. nn.Conv3d(in_channels, in_channels*2, kernel_size, padding=1 ),
  67. nn.LeakyReLU(0.1),
  68. nn.InstanceNorm3d(in_channels*2),
  69. nn.ConvTranspose3d(in_channels*2, out_channels, kernel_size=2, stride=2) if upsample else nn.Conv3d(in_channels*2, out_channels, kernel_size=1, stride=1),
  70. )
  71. nn.init.kaiming_normal_(self.block[1].weight, mode='fan_in', nonlinearity='relu')
  72. nn.init.kaiming_normal_(self.block[-1].weight, mode='fan_in', nonlinearity='relu')
  74. self.res_alpha = nn.Parameter(torch.tensor([-2.5]))
  75. adjust_volume = nn.Upsample(scale_factor=2, mode="trilinear") if upsample else nn.Identity()
  76. adjust_channels = nn.Conv3d(in_channels, out_channels, 1) if in_channels != out_channels else nn.Identity()
  77. self.adjust = nn.Sequential(adjust_volume, adjust_channels)
  79. def forward(self, x):
  80. res = self.block(x)
  81. alpha = torch.sigmoid(self.res_alpha)
  82. x = self.adjust(x)
  83. x = (1.-alpha)*x + alpha*res
  84. return x
  86. class Stage(nn.Module):
  87. def __init__(
  88. self,
  89. in_channels,
  90. out_channels,
  91. kernel_size,
  92. n_blocks,
  93. downsample=True
  94. ):
  95. super().__init__()
  96. self.blocks = nn.ModuleList(
  97. [Block(in_channels, in_channels, kernel_size) for _ in range(n_blocks-1)]+
  98. [Block(in_channels, out_channels, kernel_size, downsample=downsample)]
  99. )
  101. def forward(self, x):
  102. for block in self.blocks:
  103. x = block(x)
  104. return x
  106. class StageUp(nn.Module):
  107. def __init__(
  108. self,
  109. in_channels,
  110. out_channels,
  111. kernel_size,
  112. n_blocks,
  113. upsample=True
  114. ):
  115. super().__init__()
  116. self.blocks = nn.ModuleList(
  117. [BlockUp(in_channels, in_channels, kernel_size) for _ in range(n_blocks-1)]+
  118. [BlockUp(in_channels, out_channels, kernel_size, upsample=upsample)]
  119. )
  121. def forward(self, x):
  122. for block in self.blocks:
  123. x = block(x)
  124. return x
  126. class CNNEnc(nn.Module):
  127. def __init__(
  128. self,
  129. in_channels,
  130. stage_channels,
  131. kernel_size,
  132. blocks_per_stage
  133. ):
  134. super().__init__()
  136. n_stages = len(stage_channels)
  138. # Normalise blocks_per_stage: accept a single int or a per-stage list
  139. if isinstance(blocks_per_stage, int):
  140. blocks_per_stage = [blocks_per_stage] * n_stages
  141. assert len(blocks_per_stage) == n_stages, (
  142. f"blocks_per_stage length ({len(blocks_per_stage)}) "
  143. f"must match the number of stages ({n_stages})"
  144. )
  146. self.norm = nn.InstanceNorm3d(in_channels)
  148. # Build stages; the output of stage i is the input of stage i+1
  149. in_ch = in_channels
  150. self.stages = nn.ModuleList()
  151. first = 0
  152. for out_ch, n_blocks in zip(stage_channels, blocks_per_stage):
  153. self.stages.append(
  154. Stage(in_ch, out_ch, kernel_size, n_blocks, downsample=first>0)#last!=n_stages)
  155. )
  156. first += 1
  157. in_ch = out_ch
  159. def forward(self, x):
  160. x = self.norm(x)
  161. maps_to_decoder = []
  162. for stage in self.stages:
  163. x = stage(x)
  164. maps_to_decoder.append(x)
  165. #print(x.shape)
  167. return maps_to_decoder
  169. class CNNDec(nn.Module):
  170. def __init__(
  171. self,
  172. in_channels,
  173. stage_channels,
  174. kernel_size,
  175. blocks_per_stage
  176. ):
  177. super().__init__()
  179. n_stages = len(stage_channels)
  181. # Normalise blocks_per_stage: accept a single int or a per-stage list
  182. if isinstance(blocks_per_stage, int):
  183. blocks_per_stage = [blocks_per_stage] * n_stages
  184. assert len(blocks_per_stage) == n_stages, (
  185. f"blocks_per_stage length ({len(blocks_per_stage)}) "
  186. f"must match the number of stages ({n_stages})"
  187. )
  189. #self.norm = nn.InstanceNorm3d(in_channels)
  191. # Build stages; the output of stage i is the input of stage i+1
  192. in_ch = in_channels
  193. self.stages = nn.ModuleList()
  194. #last = 0
  195. for out_ch, n_blocks in zip(stage_channels, blocks_per_stage):
  196. self.stages.append(
  197. StageUp(in_ch, out_ch, kernel_size, n_blocks, upsample=True)
  198. )
  199. in_ch = out_ch
  200. #last += 1
  202. def forward(self, maps_to_decoder):
  203. x = 0.0
  204. for idx,stage in enumerate(self.stages):
  205. x = stage(maps_to_decoder[idx]+x)
  206. #print(maps_to_decoder[idx].shape)
  207. return x
  209. class SpatialAttentionMaskMaker3d(nn.Module):
  210. """
  211. For each spatial location in the feature map, learns a direct linear mapping
  212. from the local feature vector to N token assignment scores (one per token).
  213. A sigmoid then turns each score into an independent soft weight in (0, 1).
  215. Input: feature map of shape (B, C, H, W, D)
  216. Output: attention masks of shape (B, n_tokens, H, W, D), values in (0, 1)
  217. """
  219. def __init__(self, in_channels: int, n_tokens: int = 8, bias: bool = True):
  220. super().__init__()
  222. self.in_channels = in_channels
  223. self.n_tokens = n_tokens
  224. #self.norm = nn.LayerNorm(in_channels, elementwise_affine=False)
  225. # Single linear layer: maps each voxel's C-dim feature vector directly
  226. # to n_tokens scalar logits. Acts like a learned 1x1x1 convolution.
  227. # (B, n_voxels, C) -> (B, n_voxels, n_tokens)
  228. self.linear = nn.Linear(in_features=in_channels, out_features=n_tokens, bias=bias)
  230. def norm(self, x):
  231. return torch.nn.functional.normalize(x, p=2, dim=-1) * math.sqrt(x.shape[-1])
  233. def forward(self, feature_map_bchwd: torch.Tensor):
  234. """
  235. Args:
  236. feature_map_bchwd: (B, C, H, W, D)
  238. Returns:
  239. attention_masks_bnhwd: (B, n_tokens, H, W, D)
  240. """
  241. B, C, H, W, D = feature_map_bchwd.shape
  242. # assert C == self.in_channels, f"Expected in_channels={self.in_channels}, got {C}"
  244. # Flatten spatial dims to treat all voxels uniformly.
  245. # (B, C, H, W, D) -> (B, C, n_voxels)
  246. n_voxels = H * W * D
  247. feature_map_bcv = feature_map_bchwd.reshape(B, C, n_voxels)
  249. # Move channels last so Linear operates on the feature dimension.
  250. # (B, C, n_voxels) -> (B, n_voxels, C)
  251. feature_map_bvc = feature_map_bcv.permute(0, 2, 1)
  252. feature_map_bvc = self.norm(feature_map_bvc)
  254. # Direct linear projection: each voxel's feature vector -> n_tokens scores.
  255. # (B, n_voxels, C) -> (B, n_voxels, n_tokens)
  256. token_scores_bvn = self.linear(feature_map_bvc)
  258. # Sigmoid: independent soft weight per (voxel, token) pair, no competition across voxels.
  259. # (B, n_voxels, n_tokens) -> same shape, values in (0, 1)
  260. attention_weights_bvn = torch.sigmoid(token_scores_bvn)
  262. # Restore spatial structure and move token dim before spatial dims.
  263. # (B, n_voxels, n_tokens) -> (B, n_tokens, n_voxels) -> (B, n_tokens, H, W, D)
  264. attention_masks_bnhwd = attention_weights_bvn.permute(0, 2, 1).reshape(
  265. B, self.n_tokens, H, W, D
  266. )
  268. return attention_masks_bnhwd
  270. class TokenLearner3d(nn.Module):
  271. """
  272. TokenLearner: compresses a volumetric feature map into a small set of learned tokens.
  274. The key idea:
  275. 1. SpatialAttentionMaskMaker produces N soft spatial masks (B, N, H, W, D)
  276. 2. Each mask is element-wise multiplied with the feature map (B, C, H, W, D)
  277. 3. The masked feature map is mean-pooled over (H, W, D) -> one C-dim token per mask
  278. 4. Result: set of N tokens, shape (B, N, C), a compact set-like representation
  280. This replaces the global average pool (which treats all voxels equally) with
  281. N learned, content-adaptive pooling operations.
  283. Args:
  284. in_channels: number of feature channels C of the input
  285. n_tokens: number of output tokens N (the compression factor)
  286. bias: whether Linear layers use bias
  287. out_channels: if set, project each token from C to out_channels via a Linear layer;
  288. otherwise tokens keep their original C channels (Identity)
  289. """
  291. def __init__(
  292. self,
  293. in_channels: int,
  294. n_tokens: int = 8,
  295. bias: bool = True,
  296. out_channels: int = None,
  297. ):
  298. super().__init__()
  300. self.in_channels = in_channels
  301. self.n_tokens = n_tokens
  304. #self.norm = nn.LayerNorm(in_channels, elementwise_affine=False)
  306. self.mask_maker = SpatialAttentionMaskMaker3d(
  307. in_channels=in_channels,
  308. n_tokens=n_tokens,
  309. bias=bias,
  310. )
  312. # Optional linear projection applied independently to each token vector
  313. self.token_projector = (
  314. nn.Linear(in_channels, out_channels) if out_channels else nn.Identity()
  315. )
  316. def norm(self, x):
  317. return torch.nn.functional.normalize(x, p=2, dim=-1) * math.sqrt(x.shape[-1])
  319. def forward(self, feature_map_bchwd: torch.Tensor):
  320. """
  321. Args:
  322. feature_map_bchwd: (B, C, H, W, D) — volumetric feature map
  324. Returns:
  325. tokens_bnc: (B, n_tokens, C_out) — the learned token set
  326. attention_masks_bnhwd: (B, n_tokens, H, W, D) — masks for inspection / auxiliary loss
  327. """
  328. B, C, H, W, D = feature_map_bchwd.shape
  329. n_voxels = H * W * D
  331. # --- Step 1: build one spatial attention mask per token ---
  332. # Each mask encodes which voxels are relevant for that token.
  333. # (B, C, H, W, D) -> (B, n_tokens, H, W, D)
  334. attention_masks_bnhwd = self.mask_maker(feature_map_bchwd)
  335. attention_masks_flat_bnv = attention_masks_bnhwd.reshape(B, self.n_tokens, n_voxels)
  337. # --- Step 2: soft-mask the feature map for each token ---
  338. # Expand feature map along the token dimension so we can broadcast the mask.
  339. # (B, C, H, W, D) -> (B, 1, C, H, W, D)
  340. # feature_map_b1chwd = feature_map_bchwd.unsqueeze(1)
  341. feature_map_bcv = feature_map_bchwd.reshape(B, C, n_voxels)
  343. # (B, n_tokens, H, W, D) -> (B, n_tokens, 1, H, W, D) for channel broadcast
  344. #attention_masks_bn1hwd = attention_masks_bnhwd.unsqueeze(2)
  346. # Element-wise product: each voxel's feature vector is scaled by its token weight.
  347. # (B, n_tokens, 1, H, W, D) * (B, 1, C, H, W, D) -> (B, n_tokens, C, H, W, D)
  348. #masked_features_bnchwd = attention_masks_bn1hwd * feature_map_b1chwd
  351. # --- Step 3: spatial mean-pooling -> one token vector per mask ---
  352. # Average over the three spatial dimensions (H, W, D).
  353. # (B, n_tokens, C, H, W, D) -> (B, n_tokens, C)
  354. #tokens_bnc = masked_features_bnchwd.mean(dim=(3, 4, 5))
  356. # --- True step 2-3 ---
  357. # We do not materialize the full (B, n_tokens, C, H, W, D) tensor!
  358. tokens_bnc = torch.bmm(
  359. attention_masks_flat_bnv / (attention_masks_flat_bnv.sum(dim=2, keepdims=True) + 1e-8),
  360. #self.norm(feature_map_bcv.transpose(1,2))
  361. feature_map_bcv.transpose(1,2)
  362. ) #/ n_voxels
  364. # --- Step 4 (optional): project token channels ---
  365. # (B, n_tokens, C) -> (B, n_tokens, C_out)
  366. tokens_bnc = self.token_projector(tokens_bnc)
  368. return tokens_bnc, attention_masks_bnhwd
  370. class TokenFuser3d(nn.Module):
  371. """
  372. This module weighted-averages N tokens and computes N spatial pertinence masks,
  373. then broadcasts the tokens over the masks to update feature maps.
  374. """
  375. def __init__(self,
  376. n_tokens,
  377. conv_dim,
  378. token_dim,
  379. bias=True
  380. ):
  381. super().__init__()
  382. self.M = nn.Linear(n_tokens, n_tokens)
  383. self.Beta = SpatialAttentionMaskMaker3d(
  384. in_channels=conv_dim,
  385. n_tokens=n_tokens,
  386. bias=bias,
  387. )
  388. self.n_tokens = n_tokens
  389. self.conv_dim = conv_dim
  390. self.token_dim = token_dim
  391. if (token_dim != conv_dim):
  392. self.C = nn.Linear(token_dim, conv_dim)
  393. else:
  394. self.C = nn.Identity()
  396. def forward(self, tokens, feat_maps):
  397. B, C, H, W, D = feat_maps.shape # B,Cc,H,[W,[D]]
  398. n_voxels = H * W * D
  400. # Mix the tokens, and eventually map them to channel dimension of convolutional features
  401. mixed_tokens = self.C(self.M(tokens.transpose(1,2)).transpose(1,2)) # B, N, Ct -> B, N, Cc
  403. # Determine how much each voxel needs each token
  404. pertinence_masks_bnhwd = self.Beta(feat_maps) # B,Cc,H,[W,[D]]
  405. pertinence_masks_flat_bnv = pertinence_masks_bnhwd.reshape(B, self.n_tokens, n_voxels) # B, N, V=HWD
  407. token_broadcast_bcv = torch.bmm(mixed_tokens.transpose(1,2), pertinence_masks_flat_bnv) # B, (Cc, N) @ (N, V) -> (B,Ct,V)
  408. token_broadcast_bchwd = token_broadcast_bcv.reshape(B,C,H,W,D)
  409. out_maps = feat_maps + token_broadcast_bchwd
  410. return out_maps
  415. class TokenUNet(nn.Module):
  416. def __init__(
  417. self,
  418. in_channels,
  419. enc_stage_channels,
  420. kernel_size,
  421. blocks_per_stage,
  422. num_classes,
  423. n_tokens=8,
  424. token_dim=None,
  425. tokenize=True,
  426. attention=False,
  427. process_tokens=True,
  428. token_blocks=2,
  429. bias=True
  430. ):
  431. super().__init__()
  432. self.tokenize = tokenize
  433. self.process_tokens = process_tokens
  435. # Ensure token_dim is set (defaults to the channel size of the bottleneck)
  436. bottleneck_channels = enc_stage_channels[-1]
  437. dec_stage_channels = enc_stage_channels[::-1][1:]
  438. self.token_dim = token_dim if token_dim is not None else bottleneck_channels
  440. # 1. Encoder
  441. self.encoder = CNNEnc(
  442. in_channels=in_channels,
  443. stage_channels=enc_stage_channels,
  444. kernel_size=kernel_size,
  445. blocks_per_stage=blocks_per_stage
  446. )
  448. # 2. Tokenizer (Optional)
  449. if self.tokenize:
  450. self.token_learner = TokenLearner3d(
  451. in_channels=bottleneck_channels,
  452. n_tokens=n_tokens,
  453. bias=bias,
  454. out_channels=self.token_dim
  455. )
  457. # 3. Token Processor (Transformer / MLP Mixer) - Optional
  458. if self.process_tokens:
  459. if attention:
  460. self.token_processor = MyTransformerEncoder(
  461. d_model=self.token_dim,
  462. nhead=4,
  463. dim_feedforward=self.token_dim*2,
  464. dropout=0.0,
  465. num_layers=token_blocks
  466. )
  467. else:
  468. self.token_processor = MyMLPMixer(
  469. n_tokens=n_tokens,
  470. d_model=self.token_dim,
  471. dim_feedforward=self.token_dim*2,
  472. n_blocks=token_blocks,
  473. dropout=0.0
  474. )
  475. else:
  476. self.token_processor = nn.Identity()
  478. # 4. Token Fuser
  479. self.token_fuser = TokenFuser3d(
  480. n_tokens=n_tokens,
  481. conv_dim=bottleneck_channels,
  482. token_dim=self.token_dim,
  483. bias=bias
  484. )
  486. # 5. Decoder
  487. self.decoder = CNNDec(
  488. in_channels=bottleneck_channels, # Starts from the bottleneck size
  489. stage_channels=dec_stage_channels,
  490. kernel_size=kernel_size,
  491. blocks_per_stage=[1,]*len(dec_stage_channels)
  492. )
  494. # 6. Segmentation Head
  495. self.segmentation_head = nn.Conv3d(
  496. in_channels=dec_stage_channels[-1],
  497. out_channels=num_classes,
  498. kernel_size=1
  499. )
  501. def forward(self, x):
  502. # Forward through Encoder
  503. maps_to_decoder = self.encoder(x)
  505. # The bottleneck is the lowest resolution feature map (the last one)
  506. bottleneck = maps_to_decoder[-1]
  508. if self.tokenize:
  509. # Extract Tokens
  510. tokens, attention_masks = self.token_learner(bottleneck)
  512. # Process Tokens (Transformer/Mixer)
  513. if self.process_tokens:
  514. tokens = self.token_processor(tokens)
  516. # Fuse Tokens back into the bottleneck feature map
  517. bottleneck = self.token_fuser(tokens, bottleneck)
  519. # Update the bottleneck in our skip-connection list
  520. maps_to_decoder[-1] = bottleneck
  522. # CRITICAL: Reverse the maps so they match the Decoder's expected ascending resolutions
  523. # Example: if Enc outputs [128x128, 64x64, 32x32], Decoder needs [32x32, 64x64, 128x128]
  524. maps_to_decoder = maps_to_decoder[::-1]
  526. # Forward through Decoder
  527. out = self.segmentation_head(self.decoder(maps_to_decoder))
  529. return out

tokenunet.py at commit 0bc16c2, no license · at the source

Overview

Authors: Louis Fabrice Tshimanga1,2,3, Andrea Zanola1,2, Federico Del Pup2,3, Manfredo Atzori1,2,3,4
ORCID iDs: Louis Fabrice Tshimanga, Federico Del Pup
  1. Department of Neuroscience, University of Padua, Padua, Italy
  2. Padova Neuroscience Center, University of Padua, Padua, Italy
  3. Department of Information Engineering, University of Padua, Padua, Italy
  4. Information Systems Institute, University of Applied Sciences Western Switzerland (HES-SO Valais), Sierre, Switzerland
Institutions: University of Padua (Italy); HES-SO University of Applied Sciences and Arts Western Switzerland (Switzerland); HES-SO Valais-Wallis (Switzerland)
Journal: PloS one, volume 21, issue 8, article e0354511
Dates: received 4 March 2026; accepted 9 July 2026; published online 5 August 2026
Type: Research article · Language: English
License: CC BY
Identifiers: DOI 10.1371/journal.pone.0354511 · PMID 42555631 · PMCID PMC13440844 · OpenAlex W7172538079
Open access: gold, a free copy (OpenAlex)
Status: code verified
Categories: human (organism), other condition (population), methods / tools (subfield)
Methods: Connectivity, Machine learning
MeSH: Brain*, Imaging, Three-Dimensional*, Neuroimaging*, Algorithms, Convolutional Neural Networks, Humans, Image Processing, Computer-Assisted (* major topic)
Journal subjects: Medicine and Health Sciences, Oncology, Cancers and Neoplasms, Research and Analysis Methods, Mathematical and Statistical Techniques, Mathematical Functions, Convolution, Computer and Information Sciences, Computer Architecture, Artificial Intelligence, Machine Learning, Deep Learning, Neural Networks, Biology and Life Sciences, Neuroscience, Cognitive Science, Cognition, Memory, Memory Recall, Learning and Memory, Diagnostic Medicine, Diagnostic Radiology, Magnetic Resonance Imaging, Imaging Techniques, Radiology and Imaging
Topic: Advanced Neural Network Applications (Computer Vision and Pattern Recognition, Computer Science), according to OpenAlex
Funding: EC | HORIZON EUROPE Framework Programme (Horizon Europe) (10113707)
Citations: not cited yet (Europe PMC); 45 references in the paper

Abstract

We present TokenUNet, adopting the TokenLearner and TokenFuser modules to encase Transformers into UNets. While Transformers enable expressive global interactions among input elements in medical imaging, computational challenges hinder their deployment on common hardware. Models like (Swin)UNETR exemplify the integration of (Swin)Transformer encoders into UNets, tokenizing inputs into small subvolumes (83 voxels). The Transformer attention mechanism scales quadratically with the number of tokens, which is tied to the cubic scaling of 3D input resolution. This work reconsiders the role of convolution and attention, introducing TokenUNets, a family of 3D segmentation models better suited to constrained computational environments and time frames. To mitigate computational demands, our approach maintains the convolutional encoder of UNet-like models, and applies TokenLearner to 3D feature maps. This module pools a preset number of tokens from local and global structures, decoupling token number and input size. Our results on the BraTS challenge dataset for glioma segmentation show this tokenization effectively encodes task-relevant information, yielding naturally interpretable attention maps. The memory footprint, computation times at inference, and parameter counts of our heaviest model are reduced to 38%, 10%, and 17% of the SwinUNETR values, with statistically equivalent Dice score performance, for nnunetv2 5-fold cross-validation. This work opens the way to more efficient training in computationally restrained contexts, such as 3D medical imaging. Easing model optimization, fine-tuning, and transfer-learning in limited hardware settings can accelerate and diversify the development of approaches, for the benefit of the research community.

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

Repository

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

MedMaxLab/tokenunet

License: none: the authors keep all their rights
State: the link answers, verified on 27 September 2026
Evidence: files inventoried
Commit: 0bc16c2e27df7e6da0fed35952c2ed951f46e4e8, 17 August 2026
Languages: Python (16), Jupyter (1)
Size: 67 files, 17 scripts
Software Heritage: not archived
Found in: the text, “Training”
Holds: README, 1 notebook
Not found: license file, CITATION.cff, environment file, tests, continuous integration, documentation
Tools: NumPy (9 files), PyTorch (8 files), Matplotlib (6 files), pandas (5 files), nnU-Net (4 files), NiBabel (3 files), MONAI (2 files), seaborn (2 files), Nilearn (1 file), scikit-learn (1 file), SciPy (1 file), Hugging Face Transformers (1 file)
Availability: 1 check, the latest on 27 September 2026: the link answers
  • 27 September 2026: the link answers
18 files
At the source: github.com/MedMaxLab/tokenunet

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:

  • 1 repository of the authors' code, each at its verified commit, with its license and how the link was found in the paper;
  • 17 scripts, each with its path and the digest of its content;
  • 7 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

Datasets cited

Data Availability

Data is publicly accessible upon registration at https://www.synapse.org/Synapse:syn28546456/wiki/633440.

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, issue, pages, dates, 4 authors, 7 MeSH terms, 1 funder, 17 references.

Cite

This paper

Tshimanga, L. F., Zanola, A., Del Pup, F., & Atzori, M. (2026). TokenUNet: A new case for transformers integration in efficient and interpretable 3D UNets for brain imaging segmentation. PloS one, 21(8), e0354511. https://doi.org/10.1371/journal.pone.0354511

BibTeX

@article{tshimanga2026tokenunet,
author = {Tshimanga, Louis Fabrice and Zanola, Andrea and Del Pup, Federico and Atzori, Manfredo},
title = {{TokenUNet: A new case for transformers integration in efficient and interpretable 3D UNets for brain imaging segmentation}},
journal = {PloS one},
year = {2026},
month = aug,
volume = {21},
number = {8},
pages = {e0354511},
publisher = {PLOS},
issn = {1932-6203},
doi = {10.1371/journal.pone.0354511},
url = {https://doi.org/10.1371/journal.pone.0354511},
pmid = {42555631},
pmcid = {PMC13440844}
}

RIS

TY - JOUR
AU - Tshimanga, Louis Fabrice
AU - Zanola, Andrea
AU - Del Pup, Federico
AU - Atzori, Manfredo
TI - TokenUNet: A new case for transformers integration in efficient and interpretable 3D UNets for brain imaging segmentation
T2 - PloS one
J2 - PLoS One
PY - 2026
DA - 2026/08/05
VL - 21
IS - 8
SP - e0354511
SN - 1932-6203
PB - PLOS
DO - 10.1371/journal.pone.0354511
UR - https://doi.org/10.1371/journal.pone.0354511
LA - en
ER -

CSL-JSON

{
"id": "10.1371/journal.pone.0354511",
"type": "article-journal",
"title": "TokenUNet: A new case for transformers integration in efficient and interpretable 3D UNets for brain imaging segmentation",
"container-title": "PloS one",
"author": [
{
"family": "Tshimanga",
"given": "Louis Fabrice"
},
{
"family": "Zanola",
"given": "Andrea"
},
{
"family": "Del Pup",
"given": "Federico"
},
{
"family": "Atzori",
"given": "Manfredo"
}
],
"container-title-short": "PLoS One",
"volume": "21",
"issue": "8",
"page": "e0354511",
"DOI": "10.1371/journal.pone.0354511",
"PMID": "42555631",
"PMCID": "PMC13440844",
"ISSN": "1932-6203",
"publisher": "PLOS",
"URL": "https://doi.org/10.1371/journal.pone.0354511",
"language": "en",
"issued": {
"date-parts": [
[
2026,
8,
5
]
]
}
}

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.21037/qims-2026-0792 [code]
An nnU-Net-based framework with adaptive feature representation for 3D brain tumor segmentation.
Journal: Quantitative imaging in medicine and surgery
In common: nnU-Net, NiBabel, PyTorch, 6 other tools, methods / tools, other condition, 3 references
[2] doi:10.3389/fnins.2026.1870124 [code]
An end-to-end pipeline for automated fetal brain segmentation and biometry from 3D SSFP MRI.
Journal: Frontiers in neuroscience
In common: nnU-Net, MONAI, NiBabel, 7 other tools, 1 reference
[3] doi:10.3389/frai.2026.1771088 [code]
Few-shot deployment of pretrained MRI transformers in brain imaging tasks.
Journal: Frontiers in artificial intelligence
In common: nnU-Net, MONAI, NiBabel, 7 other tools, methods / tools
[4] doi:10.1002/hipo.70124 [code]
Association Between Anterior Hippocampal Gyrification and Episodic Memory Performance in Neurotypical Young Adults.
Journal: Hippocampus
In common: nnU-Net, Nilearn, NiBabel, 7 other tools, 1 reference
[5] doi:10.1371/journal.pcbi.1014555 [code]
Body surface potential driven personalisation of electrophysiological digital twins in hypertrophic cardiomyopathy.
Journal: PLoS computational biology
In common: nnU-Net, MONAI, NiBabel, 7 other tools
[6] doi:10.1016/j.patter.2026.101538 [code]
A multi-modal foundation model for brain disease diagnosis and medical imaging.
Journal: Patterns (New York, N.Y.)
In common: MONAI, Hugging Face Transformers, NiBabel, 6 other tools, 1 reference
[7] doi:10.1038/s41467-026-72253-7 [code]
Spurious alignment between large language models and brains can emerge from non-robust methods and overlooked confounds.
Journal: Nature communications
In common: Hugging Face Transformers, Nilearn, NiBabel, 7 other tools, methods / tools
[8] doi:10.1002/hbm.70469 [code]
VarCoNet: A Variability-Aware Self-Supervised Framework for Functional Connectome Extraction From Resting-State fMRI.
Journal: Human brain mapping
In common: Hugging Face Transformers, Nilearn, NiBabel, 7 other tools, methods / tools
[9] doi:10.1186/s13244-026-02296-3 [code]
A pre-trained foundation model framework for multiplanar MRI classification of extramural vascular invasion and mesorectal fascia invasion in rectal cancer.
Journal: Insights into imaging
In common: nnU-Net, MONAI, NiBabel, 6 other tools, other condition
[10] doi:10.1038/s41467-026-73996-z [code]
Genetic architecture of white matter microstructure captured by unsupervised deep representation learning of fractional anisotropy maps.
Journal: Nature communications
In common: MONAI, Nilearn, NiBabel, 7 other tools

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.