Diffusion Transformer Meets Random Masks: An Advanced PET Reconstruction Framework

2025-03-11

AI Summary

Traditional PET reconstruction methods lack effective joint masking mechanisms in both sinogram and latent spaces, limiting reconstruction fidelity and computational efficiency. To address this, we propose DREAM—a novel diffusion-based framework that introduces, for the first time, randomized masking in both domains (sinogram and latent space). DREAM employs high-dimensional stacked fusion to integrate multi-scale features, a hierarchical masking strategy to progressively guide learning from local details to global structure, and a mask-driven compact prior to accelerate the diffusion process. Built upon a diffusion Transformer architecture, it incorporates hierarchical masking scheduling and latent-space regularization. Experiments demonstrate that DREAM significantly improves signal-to-noise ratio and structural fidelity, better preserves clinically critical lesions, and achieves faster inference than existing diffusion-based PET reconstruction methods. The source code is publicly available.

Technology Category

• Computer Vision: Diffusion Models for Vision
• Machine Learning: Feature Construction/Reformulation
• Search and Optimization: Distributed Search

Application Category

• Responsible Web: Data and user privacy-enhancing technologies for the Web
• Graph Algorithms and Modeling for the Web: Representation, reconstruction, and subgraph or motif discovery in Web-related graphs
• Search and Retrieval-Augmented AI: Vertical and domain-specific search

Abstract

Deep learning has significantly advanced PET image re-construction, achieving remarkable improvements in image quality through direct training on sinogram or image data. Traditional methods often utilize masks for inpainting tasks, but their incorporation into PET reconstruction frameworks introduces transformative potential. In this study, we pro-pose an advanced PET reconstruction framework called Diffusion tRansformer mEets rAndom Masks (DREAM). To the best of our knowledge, this is the first work to integrate mask mechanisms into both the sinogram domain and the latent space, pioneering their role in PET reconstruction and demonstrating their ability to enhance reconstruction fidelity and efficiency. The framework employs a high-dimensional stacking approach, transforming masked data from two to three dimensions to expand the solution space and enable the model to capture richer spatial rela-tionships. Additionally, a mask-driven latent space is de-signed to accelerate the diffusion process by leveraging sinogram-driven and mask-driven compact priors, which reduce computational complexity while preserving essen-tial data characteristics. A hierarchical masking strategy is also introduced, guiding the model from focusing on fi-ne-grained local details in the early stages to capturing broader global patterns over time. This progressive ap-proach ensures a balance between detailed feature preservation and comprehensive context understanding. Experimental results demonstrate that DREAM not only improves the overall quality of reconstructed PET images but also preserves critical clinical details, highlighting its potential to advance PET imaging technology. By inte-grating compact priors and hierarchical masking, DREAM offers a promising and efficient avenue for future research and application in PET imaging. The open-source code is available at: https://github.com/yqx7150/DREAM.

Problem

• Integrates mask mechanisms in PET reconstruction for enhanced fidelity.
• Uses high-dimensional stacking to capture richer spatial relationships.
• Introduces hierarchical masking for detailed and global pattern preservation.

Innovation

• Integrates mask mechanisms in sinogram and latent space
• Uses high-dimensional stacking for richer spatial relationships
• Implements hierarchical masking for detailed and global feature capture