DRIFT: Difficulty-aware Rectified Flows for Through-plane MRI Super-Resolution

Choi, Yoonseok; Ha, Eun-Gyu; Kim, Daniel; Al-masni, Mohammed A.; Yang, Ming-Hsuan; Kim, Dong-Hyun

DRIFT concept: two-stage pipeline and Adaptive Integration Scheduler — **(a)** DRIFT reconstructs high-resolution through-plane detail in two stages — Structural Manifold Projection (APN) followed by a Difficulty-Aware Rectified Flow — avoiding the long sampling of conventional diffusion. **(b)** An Adaptive Integration Scheduler spends more steps on harder regions (higher PAD) and fewer on easy ones.

Abstract

Magnetic Resonance Imaging (MRI) is often acquired with anisotropic resolution to reduce scan time, producing stair-step artifacts along the through-plane direction. In through-plane MRI super-resolution, an efficiency–fidelity trade-off arises: feed-forward regressors are fast but oversmooth at large slice-thicknesses, while sampling-based methods improve fidelity at high inference cost. We propose DRIFT, a two-stage thickness-conditioned rectified flow framework for through-plane MRI super-resolution with continuous input slice-thickness. Stage 1 employs an Anatomical Projection Network (APN) to map low-resolution patches to a coarse high-resolution manifold, providing a deterministic anatomical initialization that shortens the residual transport of Stage 2. Stage 2 refines details via rectified flow and introduces a Physics-Aware Difficulty (PAD) metric, derived from the slice-thickness-induced through-plane bandwidth deficit, to guide an Adaptive Integration Scheduler (AIS) that allocates ODE steps by thickness. A Consistent Endpoint Trajectory Alignment (CETA) loss enforces thickness-consistent reconstructions. DRIFT outperforms super-resolution baselines while reducing inference cost.

Key Contributions

What makes DRIFT work.

Anatomical initialization (APN)

An Anatomical Projection Network maps thick-slice LR patches onto a coarse HR manifold in a single deterministic step, shortening the residual transport the flow must cover.

Thickness-conditioned rectified flow

A rectified flow refines the coarse prediction, conditioned on continuous input slice-thickness and regularized by a Consistent Endpoint Trajectory Alignment (CETA) loss for thickness-consistent outputs.

Adaptive compute (PAD → AIS)

A Physics-Aware Difficulty score, derived from the through-plane bandwidth deficit of each slice-thickness, drives an Adaptive Integration Scheduler that spends more ODE steps on thicker (harder) inputs and fewer on thin ones.

Method

Training pipeline: Shinnar–Le Roux degradation simulation, structural manifold projection (APN), and a difficulty-aware rectified flow with CETA.

DRIFT training architecture — **Stage 1** projects LR patches onto a coarse HR manifold via the APN under a slice-thickness condition. **Stage 2** trains a rectified flow on the coarse→HR trajectory, regularized by Consistent Endpoint Trajectory Alignment (CETA).

Demo

Sweeping through the volume — thick-slice low-resolution input (left) vs. DRIFT reconstruction to isotropic resolution (right). Each panel is annotated with its acquisition resolution and slice thickness.

HCP — sagittal (T1)

HCP · ×8

0.7 mm in-plane × 5.6 mm slice → 0.7 mm isotropic.

HCP · ×6

0.7 mm in-plane × 4.2 mm slice → 0.7 mm isotropic.

IDEAS — coronal (T1)

IDEAS · ×6

1.0 mm in-plane × 6.0 mm slice → 1.0 mm isotropic.

IDEAS · ×4

1.0 mm in-plane × 4.0 mm slice → 1.0 mm isotropic.

MIND — coronal (T2)

MIND · ×6

0.9 mm in-plane × 5.4 mm slice → 0.9 mm isotropic.

MIND · ×4

0.9 mm in-plane × 3.6 mm slice → 0.9 mm isotropic.

Results

+3.54 dB

PSNR over best baseline (HCP)

0.043

LPIPS @ ~20 s/volume (HCP ×8)

8–13

Adaptive NFEs (×2–×8)

Qualitative comparison

Multi-planar consistency

Multi-planar reformatted views — Although DRIFT is slice-wise, deterministic APN initialization yields coherent reformatted views with low through-plane gradient error.

Efficiency vs. accuracy

PSNR vs runtime — PSNR vs. runtime per volume. DRIFT (★) reaches the highest PSNR at a fraction of the runtime of diffusion / INR baselines.

LPIPS vs runtime — LPIPS vs. runtime per volume.

Ablation: CETA weight λ_CETA

Ablation on CETA weight across slice thickness — PSNR / SSIM / LPIPS vs. λ_CETA across slice-thickness regimes (1.5 mm easy, 3.0 mm normal, 5.0 mm hard); λ_CETA = 1.0 is consistently optimal.

Downstream segmentation

0.910

DRIFT Dice

0.816

SA-INR Dice

0.801 / 0.772

TPDM / ArSSR Dice

Downstream SynthSeg segmentation comparison on HCP x8 — SynthSeg applied to HCP T1w reconstructions (×8), with Dice computed against the HR reference. DRIFT preserves anatomical boundaries most consistently in the reformatted view, translating its through-plane gains into segmentation accuracy.

Zero-shot on real thick-slice MRI

Applied to real clinical scans without retraining or fine-tuning — no isotropic ground truth available.

Zero-shot reconstruction on in-house clinical scans — In-house 3 T scans (5 mm slice): FLAIR with IDEAS-pretrained weights, T2w with HCP-pretrained weights. DRIFT reduces stair-step artifacts and preserves anatomical continuity in reformatted views.

Zero-shot comparison on fastMRI vs BME-X foundation model — fastMRI clinical T2w (5 mm), HCP-pretrained DRIFT vs. the BME-X foundation model. DRIFT yields sharper coronal reformats — sharpness 0.291 vs. 0.238, NIQE 5.30 vs. 7.22, BRISQUE 21.16 vs. 62.46.

BibTeX

@inproceedings{choi2026drift,
  title     = {DRIFT: Difficulty-aware Rectified Flows for Through-plane MRI Super-Resolution},
  author    = {Choi, Yoonseok and Ha, Eun-Gyu and Kim, Daniel and
               Al-masni, Mohammed A. and Yang, Ming-Hsuan and Kim, Dong-Hyun},
  booktitle = {European Conference on Computer Vision (ECCV)},
  year      = {2026}
}