U-TTT: Towards Generalizable PET Image Denoising via Test-Time Training

Abstract: 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.

• The paper introduces a dual-domain test-time training method using specialized spatial (S-TTT) and frequency (F-TTT) layers integrated within a U-Net for enhanced PET denoising.
• It leverages meta-learned inner-loop adaptation with a U-shaped encoder-decoder, improving PSNR, SSIM, and lesion accuracy across both in-distribution and out-of-distribution datasets.
• The results demonstrate that aligning self-supervised feature reconstruction with denoising objectives yields robust performance in clinically variable environments.

U-TTT: Test-Time Adaptation for Generalizable PET Image Denoising

Introduction and Motivation

The performance of deep learning-based PET image denoising has been consistently challenged by real-world distribution shifts, including variations in scanner hardware and radiotracer dose levels. Conventional fixed-parameter CNNs, Transformers, or diffusion models exhibit marked degradation outside their training distribution, fundamentally limiting their utility in heterogeneous clinical environments. Recent advances in Test-Time Training (TTT) frameworks attempt to address this by introducing self-supervised adaptation at inference, but naive implementations introduce issues of primary/auxiliary task misalignment and lack task/domain specificity required for PET restoration. This paper introduces U-TTT, a U-shaped model for PET image denoising that addresses these deficits by embedding specialized TTT layers in both the spatial and frequency domains.

Figure 1: Overview of the proposed U-TTT architecture, illustrating dual-domain TTT integration within a U-Net backbone.

Methodology

Architecture Overview

U-TTT adopts a U-shaped encoder-decoder backbone tailored for 3D PET. Feature extraction is achieved via a 4-level convolutional hierarchy, enriched at every stage with two distinctly constructed TTT layers: S-TTT (spatial) and F-TTT (frequency). The output is formed by aggregating a learned residual onto the low-dose input, facilitating robust denoising.

The central innovation lies in dual-domain TTT adaptation:

• S-TTT: Learns and adapts spatially via a self-supervised feature reconstruction task. The inner model, a hybrid of depthwise conv and gated linear units, performs local spatial refinement, meta-learned to benefit denoising.
• F-TTT: Adapts in the frequency domain using FFT/IFFT projections. Local convolutions are omitted due to the global nature of the spectral basis. Instead, parametrized spectral modulation via gated units addresses globally distributed noise, crucial for high-frequency detail recovery.

Both TTT layers optimize inner-loop proxies (feature reconstruction) by differentiable online updates, ensuring their auxiliary learning is aligned with the denoising objective through end-to-end meta-learning.

Loss and Optimization

Following prevailing practice, reconstruction quality is enforced via an $\mathcal{L}_1$ loss, and structural detail is encouraged with a generative adversarial penalty, balanced by a fixed weight. Inner learning rates and architectural parameters are meta-learned.

Experimental Analysis

Benchmarking and Settings

U-TTT's generalizability is assessed using a rigorous multi-domain dataset, with both in-distribution (ID) and strict out-of-distribution (OOD) splits—OOD covering both unseen dose levels and entirely novel PET scanners. Competing methods include adversarial, Transformer, diffusion, and vector quantized approaches, all trained on the ID domain.

Quantitative and Qualitative Results

• In-Distribution: U-TTT improves over VQPET (the strongest non-TTT baseline) by +0.80 dB PSNR, +0.0028 SSIM, and reduces lesion MAE by 0.0154. Its computational footprint remains moderate (10M params; 43.52 GFLOPs), much smaller than comparable Transformers or Diffusion models.
• Out-of-Distribution: On OOD-DRF and OOD-Scanner settings, U-TTT maintains absolute performance margins of 0.67–1.0 dB in PSNR and 0.006–0.009 in SSIM over all baselines, demonstrating robust adaptation and detail recovery.

Qualitative analysis (Figure 2) indicates U-TTT's ability to restore small lesion contrast and fine anatomical structure where competitors over-smooth or hallucinate.

Figure 2: Visualization comparison at DRF=12; U-TTT recovers sharp detail and lesion contrast unachievable by previous methods.

Ablation and Component Analysis

• Isolated S-TTT (spatial) and F-TTT (frequency) blocks each confer substantial improvements; combined, they yield the maximum PSNR/SSIM and lesion accuracy.
• The custom inner model outperforms linear/MLP proxies (as used in TTT-RNNs), especially when depthwise convolutions are employed in S-TTT.
• F-TTT is particularly critical under severe spectral noise (e.g., extreme dose reduction).

Implications and Future Directions

U-TTT's contribution is twofold. The explicit demonstration that per-instance, self-supervised adaptation in both spatial and frequency domains can resolve the generalization bottleneck in 3D PET restoration is significant. By meta-learning both the adaptation proxies and feature extractors, U-TTT achieves robust alignment between self-supervised tasks and clinical objectives—addressing a key limitation of prior TTT systems.

Practically, U-TTT provides a viable pathway for robust PET denoising in diverse, real-world clinical environments where hardware and protocol drift is unavoidable. Theoretically, it motivates further work on:

• Unified dual-/multi-domain TTT layers for other volumetric imaging modalities
• Advanced meta-optimization schemes for inner loop learning rates and adaptation policies
• Extension to broader image restoration tasks under variable, non-stationary corruptions

The integration of test-time self-supervision as a modular architectural motif rather than an external patch suggests a new generation of adaptive, robust neural networks for real-world health AI deployments.

Conclusion

U-TTT sets a new state-of-the-art for generalizable PET image denoising by integrating spatial and frequency domain TTT layers for differentiable, dynamic adaptation at inference. The architecture achieves strong empirical gains over competitive baselines and is validated under severe distributional variability, offering a robust, meta-learned solution for clinical PET restoration.