PaperlayerView exact revision on arXiv

Source-linked AI summary

NUViT: Vision Transformers for Fourier Space Data

Jeffrey Gu, Robert C. Heeter, Ellen D. Zhong

arXiv:2610.32041v1cs.CVcs.LG

TL;DR

Fourier-space data is common in applications such as medical imaging and microscopy, but general-purpose transformer architectures remain underexplored and standard ViT tokenization poorly fits its structure. NUViT uses non-uniform Fourier sampling and amplitude-phase representation; across varied tasks, it outperforms standard ViT and MLP baselines. Its design assumes that low spatial frequencies carry information best sampled by a polar grid.

  • Problem

    General-purpose architectures for Fourier-space data remain underexplored, although such data is used in applications including MRI and cryo-EM.

    Source: §1 · p. 1

  • Method

    NUViT samples Fourier transforms on a non-uniform polar grid, represents coefficients by amplitude and phase, and applies positional and channel encodings before transformer processing.

    Source: §1 · p. 2

  • Results

    Across Fourier-space image classification, MRI disease classification, cryo-EM classification and reconstruction, and 2D crystal identification, NUViT consistently outperforms standard ViT and MLP baselines.

    Source: §1 · p. 2

  • Takeaways & Limitations

    NUViT is reported as a more expressive, generalizable, and robust encoder for Fourier data across the evaluated settings.

    Source: §5 · p. 9

    NUViT assumes low spatial frequencies carry information content better sampled with a polar grid, an assumption stated to fit data obeying the Fourier power law.

    Source: §2.2.4 · p. 4

Abstract

from arXiv · show

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.

1 INTRODUCTION

Fourier-space data is common across modalities, but general-purpose architectures for it remain underexplored and standard ViT tokenization fits its frequency structure poorly. NUViT redesigns Fourier sampling and representation, and the paper reports consistent improvements over standard ViT and MLP baselines across varied tasks.

1 INTRODUCTION

  • General-purpose Fourier-space transformers remain underexplored, while existing frequency-space ViTs are largely domain-specific.Potential applications include cryo-EM reconstruction, MRI, and methods using Fourier-space token mixing.Source: §1 · p. 1
  • Natural signals can exhibit 1/f α spectral energy decay, while resonance and dispersion also produce frequency-dependent structure.Examples span images, video, audio, MRI, atomic spectra, and crystal structures.Source: §1 · p. 1
  • Standard Cartesian-grid patching poorly fits Fourier frequency and phase structure, oversampling high frequencies while undersampling information-dense low frequencies when data follows a power law.The figure notes that low frequencies capture broad semantic information in natural and cryo-EM images.Source: §1 · p. 2, Figure 1 · p. 2
  • NUViT samples Fourier transforms on a non-uniform grid with NUFFT, represents complex coefficients by amplitude and phase, and applies positional and channel encodings before a transformer.The grid provides uniform frequency and phase coordinates.Source: §1 · p. 2
  • Across Fourier-space natural-image classification, MRI disease classification, cryo-EM classification and reconstruction, and high-frequency 2D crystal identification, NUViT consistently outperforms standard ViT and MLP baselines.Improvements are clear on tasks requiring fine-grained discrimination under high noise.Source: §1 · p. 2

2 METHODS

NUViT adapts Fourier-space tokenization to the signal’s frequency structure, then encodes amplitude, phase, and position in a standard transformer. Its polar grid allocates more samples to low frequencies and organizes patches by frequency shell.

2.1 MOTIVATION

  • The polar grid addresses Cartesian patching’s mismatch with frequency bands and its undersampling of information-dense low frequencies.NUViT allocates more samples to the low-frequency center.Source: §2.1 · pp. 2–3

2.2.1 TOKENIZATION

  • NUFFT samples the Fourier transform at polar coordinates, forming a radial-angular grid that is patched into tokens.The grid has N_r radial shells and N_ϕ angular samples per shell; patching uses P × P patches.Source: §2.2.1 · p. 3

2.2.2 DATA REPRESENTATION

  • Each complex Fourier value is represented by its amplitude and cosine and sine channels for phase.The amplitude-and-phase channels align with the polar grid: amplitude captures power-law decay across shells, while phase captures within-shell variation.Source: §2.2.2 · p. 3

2.2.3 POSITIONAL ENCODING

  • NUViT adds spatial and channel positional encodings, using random Fourier features for continuous 2D patch coordinates and learned embeddings for channels.The spatial encoding can generalize to different grid resolutions at test time.Source: §2.2.3 · p. 4

2.2.4 FORWARD PASS

  • Projected channel patches receive spatial and channel encodings, are prepended by a [CLS] token, and pass through a standard transformer encoder.The [CLS] output is used for classification or as the cryo-EM reconstruction encoder representation.Source: §2.2.4 · p. 4
  • The polar grid assumes low frequencies carry more information, favors isotropy, and aligns frequency shells with coarse-to-fine structure.This assumption is described as appropriate for Fourier power-law data, including natural images, MRI, and cryo-EM images.Source: §2.2.4 · p. 4

3 RESULTS

Across five settings, the results test NUViT on Fourier-domain classification, lattice identification, MRI, and cryo-EM tasks. NUViT improves cryo-EM reconstruction and amortized inference, and outperforms alternative encoders in the classification settings reported here.

3.1 NATURAL IMAGE CLASSIFICATION

  • NUViT achieves the best overall performance on CIFAR-10 and CIFAR-100; the polar grid is its most important ablated design choice, while channel encoding has little effect.Removing the polar grid costs 25.14 points on CIFAR-10; amplitude-and-phase representation is most effective but requires more parameters.Source: §3.1 · p. 5

3.2 2D CRYSTAL LATTICE CLASSIFICATION

  • NUViT substantially outperforms alternative encoders on crystal-lattice classification despite every encoder reaching 100% training accuracy, indicating a generalization gap.Random perturbations dampen higher-frequency Bragg peaks more strongly, making subtle peak information harder to identify.Source: §3.2 · p. 6

3.3 MRI CLASSIFICATION

  • Under realistic MRI conditions, NUViT gains QWK over alternative encoders on a task requiring fine-grained distinctions between adjacent severity levels.The authors suggest NUViT could operate on raw MRI frequency-domain data, leaving that evaluation to future work.Source: §3.3 · p. 6

3.4 CRYO-EM IMAGE CLASSIFICATION

  • NUViT outperforms Fourier-space baselines across cryo-EM classification, reaching 13.30% on IgG-1D and 66.79% on IgG-RL, and is more noise-robust on IgG-1D and Tomotwin-100.At SNR 0.01, the MLP encoder fails almost entirely; NUViT’s robustness is attributed tentatively to greater capacity for low-frequency components.Source: Table 5 · p. 7, §3.4 · p. 7, Figure 3 · p. 7

3.5 CRYO-EM 3D RECONSTRUCTION

  • Replacing cryoDRGN’s encoder with NUViT improves reconstruction quality on all three datasets; on IgG-1D, only RECOVAR exceeds it among deep-learning methods.RECOVAR is a linear subspace method whose adaptive kernel regression decodes multi-views at inference time.Source: Table 6 · p. 8, §3.5 · p. 8
  • NUViT achieves the best held-out test performance on all three reconstruction datasets, with better organized train and test latent spaces than the MLP encoder.On Tomotwin-100, MLP FSCAUC drops from 0.346 to 0.139; NUViT’s low-frequency focus may reduce overfitting to high-frequency noise.Source: §3.5 · pp. 8–9, Figure 4 · p. 8, Table 7 · p. 9

4 RELATED WORK

Prior work uses Fourier transforms for data representation and several computational or encoding roles, but Fourier-space architectures remain comparatively underexplored. NUViT extends irregular-grid Fourier methods from operator learning to general transformer architectures.

Fourier transforms in deep learning

  • Fourier transforms have served as data representations and as tools for efficient convolution or token mixing, positional encoding, and 3D reconstruction via the Fourier Slice Theorem.The cited examples span audio, medical imaging, microscopy, and other deep-learning applications.Source: §4 · p. 9
  • FIT unrolls Fourier images into frequency-ordered pixel tokens for tomography and super-resolution, while Alshaabi et al. (2025) retain Cartesian ViT patching for aberration correction.Neither approach applies a polar grid and NUFFT as a general-purpose architecture; complex-valued alternatives also face optimization and validation challenges.Source: §4 · p. 9

Non-uniform grids for Fourier space in deep learning

  • Fourier Neural Operators use low-frequency modes on Cartesian grids; Lingsch et al. (2023) show NUFFT-based irregular-grid FNOs can outperform equispaced ones, and NUViT extends this direction to transformers.Source: §4 · p. 9

A.1 IMPLEMENTATION DETAILS

The experiments use NUViT-S as the default configuration and process polar grids with 128 radial and angular samples unless otherwise stated.

A.1 IMPLEMENTATION DETAILS

  • NUViT-S uses 12 layers, hidden dimension 384, 12 attention heads, and an MLP dimension of 1536.The configuration is analogous to ViT-S.Source: §A.1 · p. 14
  • The default polar grid is 128 × 128 with patch size 16, yielding 64 patches per channel and 129 tokens.The RFF bandwidth is σ = 1; the maximum radial-shell radius R is controlled.Source: §A.1 · p. 14
  • Experiments use PyTorch and implement NUFFTs with TorchKbNUFFT.Unless otherwise stated, NUViT models process polar grids with Nr = Nϕ = 128.Source: §A.1 · p. 14

A.2 COMPUTE DETAILS

The supplementary sections describe compute use and experiment-specific data preparation and training settings across natural images, MRI, crystal lattices, cryo-EM classification, and reconstruction.

A.2 COMPUTE DETAILS

  • 31% more wall-clock time per epoch and 23% more memory than a ViT: NUViT’s ADH cryo-EM reconstruction training cost was measured on an H100.Compared with an MLP, NUViT was slightly faster in wall-clock time but used more GPU memory; the encoders had similar CPU memory usage.Source: Table S1 · p. 14, §A.2 · p. 14

A.3 NATURAL IMAGE CLASSIFICATION

  • Natural-image inputs are resized to 128 × 128, augmented, normalized, converted to grayscale, Fourier-transformed, and placed in the target representation.Training used batch size 64, Adam at 1e-4 for 100 epochs; CIFAR-10 used log amplitude and phase with R = 2, while CIFAR-100 used amplitude and phase with R = 4.Source: §A.3 · p. 14

A.4 MRI CLASSIFICATION

  • The OASIS MRI dataset contains scans from 416 subjects, with patient data confined to a single split and four CDR severity classes.The class counts are 316/70/28/2, making the dataset unbalanced toward non-demented patients.Source: §A.4 · pp. 14–15

A.5 2D CRYSTAL LATTICE CLASSIFICATION

  • The synthetic crystal-lattice dataset has four classes, with 1,000 training and 250 test examples per class; lattice points were perturbed by 30% of lattice spacing.The perturbation level was calibrated on ViT-S (real); models trained for 100 epochs without data augmentation.Source: §A.5 · p. 15, Figure S1 · p. 15

A.6 CRYO-EM IMAGE CLASSIFICATION

  • Cryo-EM heterogeneity classification uses random 90/10 train/test splits, batch size 64, and Adam with learning rate 1e-4.The passage lists dataset-specific radial-shell radii and uses Nϕ = 128 for all datasets.Source: §A.6 · p. 15

A.7 CRYO-EM 3D RECONSTRUCTION

  • For cryo-EM reconstruction, the synthetic IgG-1D and Tomotwin-100 models train on their full datasets, while an additional 100K CryoBench-based images assess amortization.Training durations are 20, 50, and 100 epochs for IgG-1D, EMPIAR-10076, and Tomotwin-100, respectively; cryoDRGN-MLP uses default parameters and cryoDRGN-NUViT uses the best prior-section parameters.Source: §A.7 · p. 15

B METRICS

The evaluation uses per-image Fourier Shell Correlation to assess reconstruction quality and heterogeneity, with FSCAUC summarizing the correlation curve.

B METRICS

  • Per-image FSC measures how closely each individual reconstruction matches the ground truth volume.The protocol jointly evaluates reconstruction quality and heterogeneity, following CryoBench.Source: B METRICS · p. 15
  • FSCAUC summarizes performance as the area under the FSC curve.The FSC curve compares two volumes by correlation across concentric Fourier-space shells.Source: B METRICS · p. 15
  • The protocol samples one image per conformation to capture reconstruction-quality distributions; the maximum FSCAUC is 0.5.Source: B METRICS · p. 15

C ADDITIONAL RESULTS

Top-5 accuracy results show NUViT significantly outperforming baselines on both IgG-1D and IgG-RL. The comparisons also show how classification difficulty differs across the two datasets.

C.1 CRYO-EM IMAGE CLASSIFICATION

  • NUViT significantly outperforms the baselines on top-5 accuracy for both IgG-1D and IgG-RL.Table S2 reports train / test accuracy.Source: Table S2 · p. 16, §C.1 · p. 16
  • Most baselines perform near random chance on IgG-1D, whose conformations have more similar neighbors than IgG-RL’s.Source: §C.1 · p. 16
  • On IgG-RL, most methods perform significantly better than random chance in top-5 accuracy, unlike in top-1 accuracy.Source: §C.1 · p. 16

C.2 CRYO-EM 3D RECONSTRUCTION

Across additional cryo-EM reconstruction datasets, NUViT encoders outperform ViT encoders, which outperform MLP encoders. Further comparisons and latent-space analyses report strong NUViT performance while distinguishing continuous from clustered heterogeneity.

C.2 CRYO-EM 3D RECONSTRUCTION

  • NUViT encoders outperform ViT encoders, which outperform MLP encoders on EMPIAR-10028 and CAHRA subsets.The ADH, Aldolase, PKM2 dataset combines the three proteins.Source: §C.2 · p. 16, Table S3 · p. 16
  • On IgG-1D, NUViT achieves state-of-the-art performance among deep-learning methods, trailing only RECOVAR; on Tomotwin-100, it reaches the best methods’ performance.Performance is measured by per-image FSCAUC.Source: §C.2 · p. 16, Table S4 · p. 17
  • NUViT’s train and test latent spaces are closer to IgG-1D’s ideal ring structure than those of other encoders.Trustworthiness and continuity assess neighborhood preservation against the idealized ring.Source: §C.2 · pp. 16–17, Table S5 · p. 17
  • On Tomotwin-100, MLP latent space collapses at test inference despite strong training performance, while NUViT has the best-performing latent space.ARI and AMI assess clustering quality; the ViT latent space performs only slightly less well than NUViT’s.Source: §C.2 · p. 17, Table S6 · p. 17
  • Incorrect conformation assignment is one of the main drivers of reconstruction error, partially explaining the reconstruction results.Source: §C.2 · p. 17

C.3 ABLATIONS

Cryo-EM ablations on IgG-1D and Tomotwin-100 show that design choices can vary in value across datasets. The default configuration is presented as a starting point rather than an optimum for every dataset.

C.3 ABLATIONS

  • The cryo-EM ablations repeat the study on IgG-1D and Tomotwin-100.Table S7 reports top-1 accuracy.Source: §C.3 · p. 18
  • Channel positional encoding and amplitude-phase representation are not always optimal, but perform better on difficult datasets such as IgG-1D.Source: §C.3 · p. 18
  • The default method may be a strong starting point, while features may need tuning for optimal performance on an individual dataset.Source: §C.3 · p. 18
Loading 2610.32041v1…