← 返回 PaperDaily
视觉与图像
告别显存焦虑！PPDM让3D扩散模型在普通显卡上跑出SOTA
3D扩散模型做医学图像翻译效果虽好，但训练时显存随分辨率三次方暴增，一张H100都勉强。PPDM (Pixel Puzzling Diffusion Model) 用像素拼图妙招，把空间尺寸“平移”成通道数，显存从49GB直降到6.94GB，效果还更稳。不是炫技，是真正解决落地痛点的好工作。
龙哥读论文
阅读 3
查看原文
🐉 龙哥读论文知识星球来了！公众号每日8篇拆解不够看？星球无上限更AI领域论文、资讯、招聘、招博、开源代码，一站式干货，每日2分钟刷完即赚！
👇扫码加入「龙哥读论文」知识星球，前沿干货、实用资源一站式拿捏～
龙哥推荐理由：
3D扩散模型做医学图像翻译效果虽好，但训练时显存随分辨率三次方暴增，一张H100都勉强。PPDM (Pixel Puzzling Diffusion Model) 用像素拼图妙招，把空间尺寸“平移”成通道数，显存从49GB直降到6.94GB，效果还更稳。不是炫技，是真正解决落地痛点的好工作。
原论文信息如下：
3D医学图像翻译的瓶颈
在医学影像领域，图像到图像翻译（Image-to-Image Translation, I2I）是一项核心任务，比如把低剂量的PET图像翻译成全剂量、把CT图像合成MRI，或者把T2加权的MRI变成PD加权。这些翻译能帮助医生在减少辐射、缩短扫描时间的同时，获得更丰富的诊断信息。近年来，扩散模型（Diffusion Model）凭借其稳定的训练过程和生成的高保真度，在这一领域表现惊艳，甚至超越了传统的U-Net、条件生成对抗网络（cGAN）等方法。
但一个致命的限制是：3D扩散模型的计算量和显存需求随着空间分辨率近似成立方级增长。一个128×128×64的3D体积，用全3D扩散模型训练，动辄需要49GB甚至59GB的GPU显存（比如3D BBDM），这连一张H100（80GB）都吃得够呛，更不用说主流实验室的24GB或32GB显卡了。现有的记忆高效方案，比如基于补丁的方法PatchDDM（Patch-based Diffusion Model）、基于小波的WDM（Wavelet-based Diffusion Model）、以及潜在扩散模型ALDM（Adaptive Latent Diffusion Model），要么牺牲全局一致性产生块状伪影，要么丢失高频细节导致过度平滑。总之，成了高保真3D医学图像翻译的硬瓶颈。
那么有没有一种方法，既能保持全3D扩散模型的效果，又能把显存砍到七分之一、甚至十分之一？西北大学和耶鲁大学的研究团队就在刚刚发表的论文 PPDM: Pixel Puzzling Diffusion Model for Speed and Memory Efficient Volumetric Medical Image Translation 中给出了一个非常巧妙的答案：像素拼图（Pixel Puzzling）。
像素拼图：巧妙的降维策略
PPDM的核心思想用一个词概括就是：“空间换通道”。这听起来有点像卷积神经网络里常用的Patchify操作？但PPDM的像素拼图更柔性，它保留了全局结构信息，仅仅把体素的空间位置重新排列到通道维度上。
图1：PPDM整体架构图。输入3D体积首先经过像素解拼（Unpuzzle）操作，将空间分辨率降低r倍，通道数增加r³倍，然后在低分辨率空间进行直接桥接扩散，最后通过像素拼图（Puzzle）操作恢复原分辨率。
具体来说，PPDM定义了一对可逆操作：像素解拼（Pixel Unpuzzle）和像素拼图（Pixel Puzzle）。给定一个输入张量 X ∈ RC×D×H×W，解拼操作按空间下采样因子 r 将每个r×r×r的体素块内的所有体素重新排列到通道维度上。结果得到的张量 XL 的空间维数变为 D/r × H/r × W/r，而通道数变为 C × r³。这相当于把相邻的空间信息“打包”成了通道特征，让后面的3D卷积在更低的空间分辨率上运行，同时不丢失任何原始信息。反向的拼图操作则把通道再展开回空间，无损恢复原分辨率。