[论文解读] Neuromorphologicaly-preserving Volumetric data encoding using VQ-VAE

Petru-Daniel Tudosiu, Thomas Varsavsky+6 more|arXiv (Cornell University)|Feb 13, 2020

Medical Image Segmentation Techniques参考文献 33被引用 12

一句话总结

本文提出一种3D VQ-VAE架构，可将全分辨率脑部MRI体积压缩至原始大小的0.825%，同时保持图像保真度和神经解剖结构。通过整合子像素反卷积、FixUp模块以及3D自适应损失，该模型在重建质量方面达到最先进水平，并在健康及阿尔茨海默病患者脑部分别展现出更优的形态保持能力。

ABSTRACT

The increasing efficiency and compactness of deep learning architectures, together with hardware improvements, have enabled the complex and high-dimensional modelling of medical volumetric data at higher resolutions. Recently, Vector-Quantised Variational Autoencoders (VQ-VAE) have been proposed as an efficient generative unsupervised learning approach that can encode images to a small percentage of their initial size, while preserving their decoded fidelity. Here, we show a VQ-VAE inspired network can efficiently encode a full-resolution 3D brain volume, compressing the data to $0.825\%$ of the original size while maintaining image fidelity, and significantly outperforming the previous state-of-the-art. We then demonstrate that VQ-VAE decoded images preserve the morphological characteristics of the original data through voxel-based morphology and segmentation experiments. Lastly, we show that such models can be pre-trained and then fine-tuned on different datasets without the introduction of bias.

研究动机与目标

• 通过实现高效无监督压缩，应对高维、低数据量的3D医学影像挑战。
• 克服GAN（模式崩溃、训练不稳定）和VAE（重建保真度低）在3D医学数据中的局限性。
• 确保压缩后的重建结果保留对临床和研究应用至关重要的神经解剖形态结构。
• 通过在一种数据集上预训练并在另一数据集上微调，实现模型的可迁移性，且不引入偏差。
• 证明生成模型在体素神经影像中可兼具紧凑性与生物学合理性。

提出的方法

• 通过将VAE的高斯后验替换为大小为K×D的可学习向量量化嵌入空间，将VQ-VAE框架适配至3D MRI。
• 采用子像素反卷积以减轻3D体素重建中的网格状伪影。
• 使用FixUp残差模块以在小批量条件下稳定训练，尤其避免使用批归一化。
• 引入Barron（2019）自适应损失的3D改进版本，并与Baur等人（2019）启发的损失进行比较，以提升感知质量。
• 通过组合重建损失与感知正则化，端到端训练模型，以增强结构保真度。
• 利用PixelCNN进行自回归解码，以提升从离散潜在码生成图像的质量。

实验结果

研究问题

• RQ1 3D VQ-VAE能否在保持神经解剖结构的前提下，实现全分辨率脑部MRI的高保真、低压缩比重建？
• RQ2 与最先进的基于GAN的方法相比，所提出的VQ-VAE在图像保真度和形态保持方面表现如何？
• RQ3 VQ-VAE重建结果在灰质、白质和脑脊液分割中的准确性保持程度如何？
• RQ4 在压缩后，模型是否仍能保留病理人群（如阿尔茨海默病患者与健康对照组）之间的群体水平形态差异？
• RQ5 该模型是否可在新数据集上有效微调，而不会引入偏差或性能下降？

主要发现

• 所提出的VQ-VAE模型将3D脑部MRI体积压缩至原始大小的0.825%，同时保持高重建保真度。
• 与α-WGAN基线相比，该模型在MS-SSIM指标上取得统计显著提升，且MMD更低（p < 0.01，Wilcoxon符号秩检验）。
• 重建体积中灰质、白质和脑脊液的Dice分数显著高于α-WGAN的重建结果。
• VBM分析显示，与原始数据相比，该方法在阿尔茨海默病队列中诱导的形态差异显著更少。
• 在VBM T检验图中，采用自适应损失的VQ-VAE优于Baur损失变体，表明其对群体水平形态模式的保持更优。
• 该模型展现出可迁移性，预训练权重可在新数据集上微调，且未引入偏差或性能下降。