Source-linked AI summary
U-TTT: Towards Generalizable PET Image Denoising via Test-Time Training
Zhiwen Yang, Jiayin Li, Hao Lu, Hui Zhang, Zihua Wang, Yan Xu
TL;DR
PET denoising models with fixed parameters can degrade under distribution shifts, motivating adaptation to changing test conditions. U-TTT integrates spatial and frequency TTT layers into a U-shaped model for per-instance inference-time adaptation, and reports state-of-the-art denoising with superior generalization to unseen scanners and dose levels.
Problem
Fixed-parameter PET denoising models often degrade under unseen scanners and dose levels because they assume testing data shares the training distribution, limiting robust clinical deployment.
Method
U-TTT integrates TTT layers into a U-shaped PET denoising backbone, using S-TTT for spatial structural degradations and F-TTT for global noise spectra and high-frequency details.
Results
U-TTT achieves state-of-the-art PET denoising performance and superior generalization to unseen scanners and dose levels.
Takeaways & Limitations
Per-instance dual-domain adaptation enables U-TTT to learn test-image characteristics for robust PET denoising under diverse distribution shifts.
Takeaways & Limitations
Conventional TTT can overfit its auxiliary task and catastrophically forget the primary task when auxiliary and primary objectives are misaligned.
Abstract
from arXiv · showhide
Existing deep learning models for Positron Emission Tomography (PET) image denoising often suffer from severe performance degradation under distribution shifts, fundamentally restricting their robust clinical deployment. This lack of generalization stems from the conventional paradigm of fixed-parameter models that cannot adapt to variations in test data (e.g., dose levels or scanner types) after training. To overcome this limitation and achieve robust generalization, we introduce U-TTT, a novel U-shaped model that integrates Test-Time Training (TTT) layers to dynamically adjust model parameters during inference through self-supervision, thereby adapting to the specific characteristics of each test instance. Furthermore, to comprehensively capture the complex degradations of 3D PET data, U-TTT features a dual-domain adaptation mechanism comprising a Spatial Test-Time Training (S-TTT) layer and a Frequency Test-Time Training (F-TTT) layer. The S-TTT layer captures and corrects spatial structural degradations, while the F-TTT layer suppresses global noise spectra and restores delicate high-frequency details. Extensive experiments demonstrate that U-TTT achieves state-of-the-art PET denoising performance and exhibits superior generalization under challenging distribution shifts, including both unseen dose levels and unseen scanners. Our code will be available at https://github.com/Yaziwel/U-TTT.
1 Introduction
U-TTT addresses PET denoising failures under scanner and dose distribution shifts by adapting to each test image through integrated TTT layers. Its spatial- and frequency-domain adaptation targets structural degradation, global noise, and high-frequency detail recovery.
- Motivation: Fixed-parameter PET denoising models often degrade under unseen scanners and tracer dose levels because they assume testing data matches the training distribution.This limits robust deployment in real-world clinical applications.
- Motivation: Conventional TTT adapts model parameters per test instance through an auxiliary self-supervised objective, but task misalignment can cause auxiliary-task overfitting and catastrophic forgetting.The auxiliary loss does not guarantee improvement on the primary denoising task.
- U-TTT Framework: U-TTT introduces a U-shaped PET denoising backbone with TTT layers that dynamically adapt the model to each individual image during inference.The framework treats each test image as a distinct learning problem with its own generalization target.
- U-TTT Framework: S-TTT and F-TTT provide dual-domain self-supervised adaptation, correcting spatial structural degradations while suppressing global noise spectra and restoring high-frequency details.The inner-model design and optimization are adapted to scale both layers from 1D sequences to 3D vision tasks.
- Results: Extensive experiments report that U-TTT outperforms state-of-the-art PET denoising methods and generalizes to unseen scanners and dose levels.The reported conclusion covers both denoising performance and distribution-shift generalization.
2 Method
U-TTT is a four-level U-shaped PET denoising network that adapts during inference through spatial and frequency Test-Time Training blocks. Its dual-domain design combines spatial feature reconstruction with spectral modulation to process input-specific structures and global noise characteristics.
- Overall Architecture: U-TTT uses a four-level encoder–decoder whose levels contain consecutive S-TTT and F-TTT blocks for test-time parameter adaptation.Shallow features are transformed into deep features before residual reconstruction produces the restored PET image.
- Spatial Test-Time Training: S-TTT adapts an inner spatial reconstruction model by reconstructing features and applying the updated model to input-specific test features.Its spatial reconstruction proxy updates the inner model through a differentiable, unrolled optimization loop embedded in the outer denoising network.
- Spatial Test-Time Training: The spatial reconstruction model separates channel groups, using depthwise convolution for one subset and a modified gated linear unit for the remaining channels.The two branches are concatenated to form the spatial representation.
- Frequency Test-Time Training: F-TTT projects features with FFT, adapts an inner frequency reconstruction model, and applies IFFT after tailored spectral modulation.Its frequency-domain model uses a modified gated linear unit because each frequency point encodes global spatial information.
- Frequency Test-Time Training: F-TTT learns each input’s global spectral characteristics to suppress distributed noise spectra and recover fine high-frequency details.This spectral refinement complements the structural refinements provided by S-TTT.
- Training Objective: The training objective combines L1 reconstruction loss for image content with adversarial loss for fine-detail recovery.The total loss is formulated as a weighted sum of the two terms.
3 Experiments and Results
Experiments use four whole-body PET datasets to test U-TTT on in-distribution and shifted dose-level and scanner settings. U-TTT achieves strong denoising results, with dual-domain components improving performance.
- Dataset: Four whole-body PET datasets support training, validation, in-distribution testing on D1, and out-of-distribution testing on D2–D4.D2 contains unseen dose-reduction factors, while D3 and D4 come from unseen scanners.
- In-Distribution Results: 0.80 dB in PSNR and 0.0028 in SSIM are U-TTT’s average gains over VQPET, while lesion error decreases by 0.0154 on the base dataset.The comparison spans four dose-reduction factors and three metrics.
- Comparative Experiments: U-TTT outperforms five state-of-the-art methods across in-distribution and out-of-distribution evaluation settings.Compared methods include 3D-cGAN, DRMC, Spach Transformer, 3D DDPM, and VQPET.
- Out-of-Distribution Results: U-TTT achieves the best performance under unseen dose levels and scanners.The reported generalization is associated with dynamically adjusting parameters to each test instance during inference.
- Ablation Study: Combining S-TTT and F-TTT yields the best overall ablation performance, with F-TTT more effective than S-TTT individually.The modified efficient gated linear unit and depth-wise convolution further improve the inner model design.
4 Conclusion
U-TTT uses Test-Time Training for dynamic inference-time adaptation in PET denoising. Its spatial and frequency-domain layers support instance-specific restoration and noise suppression, while experiments report state-of-the-art performance and generalizability.
- Conclusion: U-TTT enables dynamic model adaptation during inference through Test-Time Training.The framework learns instance-specific characteristics from each test image.
- Conclusion: Dual-domain S-TTT and F-TTT layers learn spatial and frequency characteristics to restore structural details and suppress global noise.The two layers address complementary aspects of 3D PET degradation.
- Conclusion: Extensive experiments demonstrate state-of-the-art PET denoising performance and superior generalizability.The conclusion reports these outcomes without specifying individual metrics.