NUViT: Vision Transformers for Fourier Space Data
Organizations: Princeton University
Abstract
Vision Transformers (ViTs) have demonstrated strong performance in many computer vision tasks but do not naturally extend beyond the natural image domain. One important example is Fourier-space data, which is used to represent a rich and diverse set of modalities, including medical imaging and microscopy. Standard ViT tokenization partitions an image into square patches on a Cartesian grid, which poorly fit the frequency and phase structure of Fourier-space data. We introduce Non-Uniform ViT, a simple, general-purpose vision transformer that redesigns image tokenization, data representation, and positional encoding to respect this structure. Experiments on Fourier-domain image classification, 2D crystal lattice identification, MRI annotation, and cryo-EM protein heterogeneity analysis show that NUViT consistently captures the structure of Fourier data better than standard architectures.
Figures & tables
| Ablation | CIFAR-10 | CIFAR-100 |
|---|---|---|
| Cartesian grid | 43.49 | 14.95 |
| Real + Imaginary | 49.73 | 18.79 |
| Hartley | 46.04 | 18.12 |
| No channel positional encoding | 66.04 | 27.33 |
| NUViT (Ours) | 68.63 | 27.20 |
| Method | Test QWK ( ) |
|---|---|
| ViT (real) | 0.266 |
| ViT (Hartley) | 0.232 |
| NUViT (Ours) | 0.348 |
| Method | IgG-1D | IgG-RL | Ribosembly | Tomotwin-100 | EMPIAR-10076 |
|---|---|---|---|---|---|
| Random chance | 1.00 | 1.00 | 6.25 | 1.00 | 6.25 |
| MLP (real) | 99.33 / 1.25 | 99.60 / 2.10 | 99.61 / 42.40 | 99.50 / 42.12 | 99.65 / 39.94 |
| MLP (Hartley) | 99.10 / 1.43 | 99.06 / 2.61 | 99.58 / 44.58 | 99.18 / 32.71 | 99.77 / 38.55 |
| ViT-S (real) | 100.00 / 1.05 | 100.00 / 1.05 | 97.55 / 45.49 | 99.73 / 75.70 | 98.98 / 46.47 |
| ViT-S (Hartley) | 1.66 / 1.41 | 57.82 / 8.11 | 98.01 / 72.56 | 98.84 / 84.32 | 98.73 / 50.89 |
| NUViT (Ours) | 36.54 / 13.30 | 94.77 / 66.79 | 98.80 / 80.62 | 99.36 / 88.86 | 96.84 / 51.52 |
| Encoder | IgG-1D | Tomotwin-100 | EMPIAR-10076 | |||
|---|---|---|---|---|---|---|
| Mean (std) | Median | Mean (std) | Med | Mean (std) | Med | |
| MLP (Hartley) | 0.350 (0.019) | 0.357 | 0.346 (0.033) | 0.349 | 0.418 (0.027) | 0.422 |
| ViT (Hartley) | 0.363 (0.006) | 0.365 | 0.337 (0.027) | 0.340 | 0.423 (0.021) | 0.397 |
| NUViT (Ours) | 0.367 (0.003) | 0.368 | 0.345 (0.039) | 0.348 | 0.433 (0.020) | 0.435 |
| Encoder Method | IgG-1D | Tomotwin-100 | EMPIAR-10076 | |||
|---|---|---|---|---|---|---|
| Mean (std) | Median | Mean (std) | Med | Mean (std) | Med | |
| MLP (Hartley) (original) | 0.337 (0.018) | 0.336 | 0.139 (0.104) | 0.120 | 0.400 (0.041) | 0.413 |
| ViT (Hartley) | 0.354 (0.016) | 0.359 | 0.293 (0.101) | 0.331 | 0.399 (0.035) | 0.411 |
| NUViT (Ours) | 0.361 (0.011) | 0.365 | 0.301 (0.105) | 0.341 | 0.409 (0.037) | 0.422 |
Appendix figures & tables8 assets
Supplementary material from the paper’s appendix.
Appendix
| Encoder | Params | Time / epoch | GPU Mem. | CPU Mem. |
|---|---|---|---|---|
| MLP | 19.9 M | 2:19.62 | 12.3 GB | 7.9 GB |
| ViT | 24.9 M | 1:40.42 | 13.6 GB | 7.9 GB |
| NUViT | 25.1 M | 2:11.84 | 16.7 GB | 7.8 GB |
| Method | IgG-1D | IgG-RL |
|---|---|---|
| Random chance | 5.00 / 5.00 | 5.00 / 5.00 |
| MLP (real) | 99.99 / 5.57 | 100.00 / 9.08 |
| MLP (Hartley) | 99.99 / 6.59 | 99.98 / 11.54 |
| ViT-S (real) | 100.00 / 4.77 | 100.00 / 5.84 |
| ViT-S (Hartley) | 7.96 / 7.19 | 92.30 / 28.54 |
| NUViT (Ours) | 86.63 / 59.01 | 99.95 / 96.78 |
| Dataset | MLP | ViT | NUViT (Ours) |
|---|---|---|---|
| EMPIAR-10028 | 0.471 | 0.478 | 0.480 |
| CAHRA-ADH | 0.390 | 0.448 | 0.481 |
| CAHRA-Aldolase | 0.363 | 0.437 | 0.477 |
| CAHRA-PKM2 | 0.402 | 0.448 | 0.462 |
| CAHRA-ADH/Aldolase/PKM2 | 0.437 | 0.448 | 0.455 |
| Method | IgG-1D | Tomotwin-100 | ||
|---|---|---|---|---|
| Mean (std) | Median | Mean (std) | Med | |
| 3D classification ( Punjani et al., 2017 ) | 0.297 (0.019) | 0.291 | 0.046 (0.026) | 0.037 |
| 3DVA ( Punjani and Fleet, 2021 ) | 0.349 (0.004) | 0.350 | 0.088 (0.040) | 0.077 |
| RECOVAR ( Gilles and Singer, 2025 ) | 0.386 (0.005) | 0.388 | 0.258 (0.109) | 0.254 |
| CryoDRGN ( Zhong et al., 2020 ) | 0.350 (0.019) | 0.357 | 0.346 (0.033) | 0.349 |
| CryoDRGN-AI ( Levy et al., 2025 ) | 0.364 (0.002) | 0.364 | 0.202 (0.044) | 0.321 |
| Train | Test | |||
|---|---|---|---|---|
| Encoder | Trust. | Cont. | Trust. | Cont. |
| MLP | 0.9260 | 0.8467 | 0.7534 | 0.6910 |
| ViT | 0.9454 | 0.9429 | 0.8791 | 0.8910 |
| NUViT (Ours) | 0.9634 | 0.9692 | 0.9111 | 0.9323 |
| Train | Test | |||
|---|---|---|---|---|
| Encoder | ARI | AMI | ARI | AMI |
| MLP | 0.9952 | 0.9968 | 0.1356 | 0.3601 |
| ViT | 0.9892 | 0.9944 | 0.6924 | 0.8194 |
| NUViT (Ours) | 0.9838 | 0.9942 | 0.7067 | 0.8222 |
| Method | Tomotwin-100 | IgG-1D |
|---|---|---|
| Channel positional encoding | ||
| No channel PE | 90.45 | 7.62 |
| Channel PE | 88.86 | 12.84 |
| Feature representation | ||
| Hartley | 67.68 | 4.91 |
| Real/Imaginary | 90.01 | 11.04 |
Vision Transformer
Vision Transformers (ViTs) use self-attention over image patches to learn visual representations for recognition, segmentation, and other tasks.
See all 364 papers