💻 computer science

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

DRIFT is a two-stage rectified flow framework for through-plane MRI super-resolution that combines an Anatomical Projection Network for deterministic initialization with a Physics-Aware Difficulty-guided adaptive scheduler to achieve high-fidelity, thickness-consistent reconstructions while significantly reducing inference costs compared to existing methods.

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

Published 2026-07-21

📖 7 min read🧠 Deep dive

CC BY 4.0

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

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to take a perfect, high-definition photo of a bustling city, but your camera lens is a bit fuzzy in one direction. Instead of a crisp, 3D view, you end up with a stack of flat pictures where the details between the layers are blurry and stepped, like a staircase made of foggy glass. This is exactly what happens with a special kind of medical camera called an MRI. Doctors need these "staircase" images to be smooth and sharp to see tiny details in the brain, but getting that perfect picture takes a long time, which can be scary or uncomfortable for patients. So, scientists have been trying to use computers to magically "fill in" the missing steps and turn those blurry stacks into smooth, high-definition 3D volumes. The big challenge has been finding a way to do this quickly without making the brain look like a melted wax sculpture, or taking so long that the computer crashes before the doctor can see the result.

Enter DRIFT, a new computer program designed by researchers to fix these blurry MRI slices. Think of DRIFT as a two-step magic trick for brain images. First, it acts like a skilled architect who quickly sketches a rough, solid outline of the brain based on the blurry input. This "Anatomical Projection Network" (APN) doesn't try to guess every tiny hair or blood vessel yet; it just makes sure the big shapes are in the right place, giving the computer a stable foundation. Then, in the second step, a "Rectified Flow" engine takes over. Imagine this as a team of tiny, hyper-focused artists who zoom in on the rough sketch to add the fine details, textures, and sharp edges. What makes DRIFT special is that it knows exactly how "fuzzy" the original picture was. It uses a clever "difficulty meter" to decide how much effort to spend: if the original slice was very thick and blurry, it works harder and longer; if the slice was already pretty clear, it finishes quickly. This smart balancing act means it creates sharper, more accurate brain maps than previous methods, and it does so much faster, saving time for both the computer and the patient.

The Problem: The Stair-Step Blur

When doctors scan brains with MRI, they often have to choose between speed and quality. To scan fast, they take "thick slices" of the brain. Imagine slicing a loaf of bread with a very thick knife; you get big chunks, but you lose the fine details of the crust between the slices. In MRI, this creates a "stair-step" effect where the image looks blocky and blurry when viewed from certain angles. The goal of Super-Resolution (SR) is to use math to turn these thick, blurry slices back into a smooth, high-resolution 3D loaf.

However, previous methods had a tough time. Some were fast but made the brain look too smooth, like a plastic model with no texture. Others were very accurate but took forever to compute, like trying to paint a masterpiece by hand for every single slice. There was also a problem with "continuous" thickness: real-world MRI scans don't always come in neat, pre-set sizes (like exactly 5mm or 6mm). They can be 5.2mm, 5.8mm, or anything in between. Older computers struggled to handle this variety, often getting confused or producing blurry results when the slice thickness didn't match their training exactly.

The Solution: A Two-Stage Magic Trick

The researchers behind DRIFT (Difficulty-aware Rectified Flows for Through-plane MRI Super-Resolution) realized they needed a smarter approach. Instead of trying to guess the whole picture at once, they broke the job into two distinct stages, like a construction crew that first builds the frame and then does the interior design.

Stage 1: The Rough Sketch (Anatomical Projection)
First, the program uses a network called the Anatomical Projection Network (APN). Think of this as a quick, rough draft. The APN looks at the blurry, thick slice and projects it onto a "coarse high-resolution manifold." In plain English, it draws a solid, stable outline of the brain structures. It doesn't try to recover the tiny, high-frequency details (like the texture of the brain tissue) yet. Its job is to provide a deterministic, reliable starting point. This is crucial because it shortens the distance the computer has to travel in the next step. It's like laying down a solid foundation before building a house; you don't start by painting the walls if the floor isn't even.

Stage 2: The Detail Work (Rectified Flow)
Once the rough sketch is ready, the second stage kicks in. This uses a technique called Rectified Flow, which is a type of generative AI that learns to move from a blurry state to a sharp state along a straight, efficient path. Instead of wandering aimlessly like some older AI models, this engine follows a direct line to add the missing high-frequency details. Because the first stage already handled the big shapes, this second stage only needs to focus on the fine details, making the process much faster and more stable.

The Secret Sauce: Knowing When to Work Hard

The most playful and clever part of DRIFT is how it decides how much effort to put in. The researchers introduced a Physics-Aware Difficulty (PAD) metric. Imagine you are a student taking a test. If the questions are easy, you finish quickly. If they are hard, you spend more time thinking. DRIFT does the same thing.

It calculates a "difficulty score" based purely on the slice thickness (the physical property of the scan), not by looking at the image itself.

Thick slices (very blurry) = High difficulty = The computer takes more steps to refine the image.
Thin slices (less blurry) = Low difficulty = The computer takes fewer steps and finishes fast.

This is managed by an Adaptive Integration Scheduler (AIS). Instead of forcing the computer to do the same amount of work for every single image (which wastes time on easy ones and isn't enough for hard ones), DRIFT dynamically allocates the number of calculation steps. This ensures that difficult, thick-slice scans get the attention they need, while easy scans don't waste computing power.

Keeping Things Consistent

One tricky part of fixing MRI slices is making sure the brain looks smooth from top to bottom, without weird jumps or jumps between slices. The researchers added a special rule called CETA (Consistent Endpoint Trajectory Alignment). This acts like a quality control inspector. It checks that if you have two slices that are very similar in thickness, the computer's "path" to fixing them should be consistent. It forces the AI to produce smooth, continuous results across the whole 3D volume, preventing those annoying "stair-step" artifacts from reappearing in the final image.

What They Found

The team tested DRIFT on several public brain MRI datasets (HCP, MIND, and IDEAS) and also on real-world, thick-slice scans that didn't have a perfect "answer key" to compare against.

Better Quality: DRIFT consistently produced sharper images with better details than other top methods. For example, on one dataset, it improved the image quality score (PSNR) by up to 3.54 points compared to the next best method, which is a huge jump in this field.
Faster: It was significantly faster than complex diffusion models. While some other methods took hundreds of steps to generate an image, DRIFT could do it in as few as 8 to 15 steps, depending on the difficulty.
Robustness: It worked well even on slice thicknesses it hadn't seen before (out-of-distribution), showing that it truly understood the physics of the problem rather than just memorizing patterns.
Zero-Shot Success: When tested on real clinical scans without retraining, DRIFT still managed to reduce the stair-step artifacts and keep the anatomy looking natural, outperforming other models that struggled with the new data.

In short, DRIFT suggests that by splitting the job into a "rough sketch" and a "detail phase," and by letting the computer know exactly how hard the job is based on the slice thickness, we can get high-quality, 3D brain images much faster and more reliably than before. It's a step forward in making MRI scans faster for patients without sacrificing the clarity doctors need to make life-saving diagnoses.

Technical Summary: DRIFT for Through-plane MRI Super-Resolution

Problem Statement

Magnetic Resonance Imaging (MRI) is frequently acquired with anisotropic resolution to minimize scan time, resulting in thick-slice volumes with low resolution along the through-plane direction. This acquisition method introduces "stair-step" artifacts and limits the utility of the data for precise clinical diagnosis and quantitative neuroanatomical analysis.

Existing approaches to through-plane super-resolution (SR) face a fundamental trade-off between efficiency and fidelity:

Feed-forward regressors are computationally fast but tend to oversmooth details, particularly at large slice thicknesses, and often fail to generalize across continuous variations in slice thickness.
Sampling-based generative methods (e.g., diffusion models) offer higher fidelity and better texture recovery but suffer from prohibitive inference costs due to the large number of sampling steps required.
Implicit Neural Representations (INRs) offer arbitrary-scale flexibility but suffer from spectral bias (under-recovering high-frequency textures) and often fail to explicitly model the specific physics of MRI slice acquisition (slice-profile integration).

Furthermore, most existing methods assume discrete scale degradation, whereas clinical practice involves continuous variations in slice thickness.

Methodology: DRIFT

The authors propose DRIFT (Difficulty-aware Rectified Flows for Through-plane MRI Super-Resolution), a two-stage framework designed to handle continuous input slice thicknesses while balancing reconstruction fidelity and inference cost. The framework is conditioned on both the input and target slice thicknesses.

1. Training Data Simulation

To mimic clinical acquisition, DRIFT synthesizes low-resolution (LR) inputs from high-resolution (HR) ground truth using a Slice-Profile based simulation. This process utilizes the Shinnar-Le Roux (SLR) algorithm to model the slice-selection process, applying a thickness-specific 1D slice-profile kernel to integrate signals along the degradation axis. This reproduces the characteristic stair-step artifacts and through-plane signal integration found in real MRI data.

2. Stage 1: Anatomical Projection Network (APN)

The first stage employs an Anatomical Projection Network (APN) to map LR patches to a coarse high-resolution manifold.

Function: The APN provides a deterministic, spatially correlated anatomical initialization. It projects the input from the thick-slice space to a coarse HR space anchored at the target thickness.
Conditioning: The network is conditioned on the inverse thickness ( $\tau = 1/T$ ) of both the input and target slices via lightweight MLP encoders and Adaptive Group Normalization (AdaGN).
Goal: This stage shortens the residual transport distance for the second stage, providing a stable structural estimate that reduces the burden on the generative refinement process.

3. Stage 2: Difficulty-Aware Rectified Flow

The second stage refines high-frequency details using Rectified Flow, a generative model that learns straight-line Ordinary Differential Equation (ODE) trajectories.

Process: Starting from the coarse estimate ( $z$ ) produced by Stage 1, the velocity network predicts a constant velocity vector to transport the image to the target HR distribution ( $y$ ).
Training: The network is trained to match the constant target velocity ( $y - z$ ) using a Huber loss. The training employs a U-shaped timestep distribution to emphasize endpoints where refinement is most sensitive.
Consistency Regularization (CETA): To ensure consistency across neighboring slice thicknesses, the authors introduce the Consistent Endpoint Trajectory Alignment (CETA) loss. This loss minimizes the distance between the predicted endpoints of proximal thickness pairs (e.g., $T$ and $T - 1$ mm), regularizing the flow trajectories to prevent discontinuities.

4. Adaptive Inference: PAD and AIS

A key innovation is the Physics-Aware Difficulty (PAD) metric and the Adaptive Integration Scheduler (AIS).

PAD: Instead of using image-dependent difficulty networks, DRIFT calculates difficulty directly from metadata. The PAD metric is defined as $1 - (T_{hr}/T_i)$ , representing the normalized through-plane bandwidth deficit caused by the input slice thickness.
AIS: The scheduler maps the PAD value to the number of ODE integration steps ( $N$ ). Thicker slices (higher difficulty) receive more integration steps, while thinner slices receive fewer. This allows the model to adapt its computational budget dynamically without auxiliary image analysis modules.

Key Contributions

Two-Stage Formulation: The paper formulates through-plane MRI SR as a two-stage anatomical projection-to-transport problem. The thickness-conditioned APN provides a deterministic initialization that significantly shortens the residual rectified-flow trajectory.
Physics-Aware Conditioning and Inference: Both stages are conditioned on input and target inverse thickness. The proposed PAD/AIS mechanism allocates ODE steps based on the physical degradation severity (slice thickness) rather than image content, eliminating the need for auxiliary difficulty networks.
CETA Loss: The authors propose a proximal-thickness endpoint alignment loss that regularizes rectified-flow trajectories across neighboring slice-thickness conditions, improving consistency over naive random-pair strategies.
Comprehensive Validation: The method is validated on public datasets (HCP, MIND, IDEAS) and real-world thick-slice MRI (in-house and fastMRI), demonstrating robustness across continuous thickness variations and zero-shot transfer.

Experimental Results

Quantitative Performance: On public datasets, DRIFT consistently outperforms state-of-the-art fixed-scale and arbitrary-scale baselines (including SwinIR, TPDM, ResShift, and various INR models). For instance, on the HCP dataset at $\times8$ scale, DRIFT achieves a PSNR of 35.93 dB compared to 33.13 dB for SwinIR.
Qualitative Performance: Visual comparisons show that DRIFT better preserves fine anatomical structures and textures while effectively suppressing stair-step artifacts. It maintains strong continuity along the slice-acquisition direction, reducing inter-slice inconsistencies common in other methods.
Zero-Shot Generalization: DRIFT demonstrates strong zero-shot capabilities on real thick-slice MRI data (without isotropic ground truth), reducing artifacts in reformatted views and outperforming foundation models like BME-X in sharpness and no-reference quality metrics (NIQE, BRISQUE).
Efficiency: DRIFT achieves a superior efficiency-fidelity trade-off. It produces higher fidelity (lower LPIPS, higher PSNR) than feed-forward regressors while requiring significantly fewer function evaluations (NFEs) than multi-step diffusion models. The adaptive scheduler ensures that thin-slice inputs are processed quickly, while thick-slice inputs receive the necessary computational resources.

Significance and Claims

The paper claims that DRIFT addresses the critical gap between the need for high-fidelity, isotropic MRI reconstruction and the practical constraints of inference time and continuous clinical variability. By integrating physical degradation models (SLR) with rectified flows and metadata-driven adaptive scheduling, DRIFT offers a solution that is:

Generalizable: Capable of handling continuous slice-thickness variations without retraining for specific scales.
Efficient: Reducing inference costs by dynamically allocating computation based on physical difficulty rather than using fixed, conservative step counts.
Clinically Relevant: Providing deterministic, anatomically consistent reconstructions that mitigate the stair-step artifacts hindering current clinical workflows.

The authors note a limitation: the main experiments utilize fixed native target thicknesses, and extending the framework to support continuous target-thickness sampling across different resolutions remains a direction for future work.