Vision transformer fuses x-ray cameras to beat the space-time trade-off at light sources

Synchrotron light sources can film the invisible. At facilities such as the Advanced Photon Source (APS) at Argonne National Laboratory, high-speed x-ray imaging experiments routinely capture millions of frames per second, revealing processes that unfold far too quickly for any conventional camera: pores forming in molten metal during additive manufacturing, cracks propagating through materials, turbulent combustion, and electric discharges. But there is a catch that has frustrated experimentalists for years. Modern detectors face a fundamental trade-off between spatial and temporal resolution, a contradiction well known in remote sensing. Push a camera to film at megahertz frame rates and you sacrifice spatial detail; demand crisp, high-resolution images and you under-sample the very dynamics you are trying to study. A new study published in Frontiers in High Performance Computing shows how a vision transformer, trained on a supercomputer, can have it both ways.

The research, led by Songyuan Tang of the APS in collaboration with Tekin Bicer, Kamel Fezzaa, and Samuel Clark of Argonne National Laboratory, tackles the problem by fusing two x-ray image streams of the same physical process. One stream comes from a high-speed (HS) camera with excellent spatial resolution but a modest frame rate. The other comes from an ultra-high-speed (UHS) camera that films at megahertz rates but with far fewer pixels. The team’s deep learning model, called SWIN-XVR, takes a handful of consecutive low-resolution (LR) frames from the UHS camera and two high-resolution (HR) frames from the HS camera, and reconstructs a new image sequence in which every frame has the pixel size of the HR images and the timing of the LR images. In effect, the model consolidates the capabilities of both cameras into a single virtual detector.

The architecture builds on the Shifted Window (SWIN) transformer, a vision transformer design that distributes self-attention across spatially local, shifting patch groups rather than relating every pixel to every other pixel. Standard vision transformers scale quadratically with image size in two dimensions, which makes them impractical for dense vision tasks such as video restoration. SWIN transformers sidestep this bottleneck by computing attention within local spatio-temporal windows and alternately shifting those windows between transformer blocks to expand the effective receptive field. SWIN-XVR consists of three modules: an input feature extraction module that converts multi-resolution inputs into shallow features through convolution branches, an attention-based deep feature enhancement module that stacks SWIN transformer blocks in multiple stages, and an output reconstruction module that selects the target feature map and upscales spatial dimensions by a factor of four using pixel shuffle layers.

Inside each SWIN transformer block, a multi-head self-attention (MSA) layer jointly attends to multiple feature subspaces and aggregates features within each local spatial window and across all feature maps. A two-layer feed-forward network cascades with the MSA for fine adjustment, layer normalization regularizes features throughout the deep network, and a learnable relative position bias is added to the self-attention maps to better model spatial relationships among features. Stochastic depth randomly drops samples within each training batch for regularization, and residual connections between subsequent transformer stages enhance feature propagation. The convolution outputs are used directly as tokens with a tokenization patch size of just one by one, allowing the model to work at fine granularity.

Training data came from 1,358 videos recorded with a Photron FastCam SA-Z camera operating at frame rates of 1 to 140 kilohertz and a Shimadzu HPV-X2 camera operating at 1 to 5 megahertz, both deployed during high-speed synchrotron x-ray imaging experiments at the 32-ID beamline of the APS. The videos cover a wide range of advanced manufacturing processes. The Photron videos contain 101 to 8,901 frames, while the Shimadzu videos contain 128 frames. Low-resolution images were synthesized by binning high-resolution images with a four-by-four window, and random crops of 96 by 96 and 24 by 24 pixels were made on corresponding HR and LR images to cover spatially consistent regions. The model was trained to minimize the L1 loss with the Adam optimizer and an initial learning rate of 0.0002, updated by a cosine annealing scheduler with warmup.

The decisive step was moving model training from beamline-local computing resources to the Polaris supercomputer at the Argonne Leadership Computing Facility (ALCF). The team implemented data parallelization with PyTorch’s distributed data parallelization (DDP) library, initiating multiple processes across four Nvidia A100 SXM4 GPUs on a single node and across multiple nodes, with the Nvidia collective communications library (NCCL) and message passing interface (MPI) synchronizing GPU devices across the cluster. Using 40 Nvidia A100 GPUs, the end-to-end training time was sped up by roughly a factor of ten compared to training with beamline-local resources. The team monitored GPU utilization using the HTA software package and found that utilization exceeded 90 percent with more than five samples in a batch, so they fixed the batch size at ten for the scalability analysis.

Strong scaling experiments, in which a fixed training job of 1,206,240 images was distributed across increasing numbers of GPUs from four to 256, revealed that scaling efficiency remained above 90 percent for up to 32 GPUs. Beyond that threshold, communication overhead among GPUs began to erode the gains. Kernel time analysis showed that the percentage of computation kernel time generally decreased with the number of GPUs, with the opposite trend in communication kernel time, confirming that distributed training faces a lower bound in wall time as GPU counts continue to rise.

Hyperparameter tuning on Polaris identified the optimal configuration: three input LR frames, a spatial attention window of four by four, a hidden token size of 384, and six attention heads per transformer block. Doubling the hidden size significantly improved reconstruction fidelity, by up to 0.8 decibels in peak signal-to-noise ratio (PSNR) on one test case, at the proportional cost in inference time. Doubling model depth yielded only a slight improvement of up to 0.12 decibels while nearly doubling inference time, so the team doubled the hidden size and used six attention heads for the final architecture. With three subsequent LR frames and two HR frames differing in spatial and temporal resolution by factors of four and twenty respectively, the algorithm achieved average PSNRs of 37.40 and 35.60 decibels on the two testing datasets.

The validated model was tested on two independent x-ray videos from real applications, one of additive manufacturing and one of friction stir welding, both previously uploaded to the tomobank repository. Testing conditions varied the PSNR of the input LR images from roughly 20 to 60 decibels by simulating Poisson noise, and varied the temporal down-sampling of the HR sequence by factors of 2, 10, 20, and 450. Compared against four independent algorithms, namely bicubic interpolation, Bayesian fusion, EDVR, and EDVR-STF, SWIN-XVR performed close to or significantly better than EDVR-STF on the two cases. Bayesian fusion demonstrated higher accuracy only when input LR images were nearly noise-free and in the absence of scene motion, but its performance dropped significantly in the presence of complex motion dynamics. The deep learning-based algorithms consistently outperformed bicubic interpolation.

The practical implications are substantial. Currently, reconstructing 500 high-resolution images of dimension 400 by 1,024 takes about 115 seconds on four Nvidia A100 GPUs. Because the UHS and HS cameras operate in parallel, users avoid frequent camera switching, which could take a prohibitive amount of time during experiments. With total average experiment durations ranging from about ten minutes to an hour, the proposed image acquisition and reconstruction framework could double the data throughput within the same time budget. The algorithm is implemented in Python and publicly available on GitHub. Looking ahead, the team envisions imaging data accumulated at the beamline being transferred to HPC systems to train foundation models at scale, which could then be deployed back to local application servers and fine-tuned for specific tasks, supporting user services with improved accuracy, automation, and robustness. The pipeline is intended to be reproducible during routine operations of ultra-high-speed imaging beamlines such as 32-ID at the APS, and techniques such as domain adaptation could generalize the method to other detectors and modalities, including computerized tomography and magnetic resonance imaging.

Subject of Research: A SWIN-based vision transformer for high-fidelity and high-speed imaging experiments at light sources

Article Title: A SWIN-based vision transformer for high-fidelity and high-speed imaging experiments at light sources

Article References: Tang, S., Bicer, T., Fezzaa, K., & Clark, S. (2025). A SWIN-based vision transformer for high-fidelity and high-speed imaging experiments at light sources. Frontiers in High Performance Computing, 3, Article 1537080. https://doi.org/10.3389/fhpcp.2025.1537080

Image Credits: AI Generated

DOI: 10.3389/fhpcp.2025.1537080

Keywords: SWIN-based, vision, transformer, high-fidelity, high-speed, imaging, experiments, light, sources, scientific research