【论文阅读|cryoET】Adaptive differentiable grids for cryo-electron tomography reconstruction and denoising

期刊：Bioinformatics Advances, 2023
源码（c++）：https://github.com/yuanhaowang1213/adaptivediffgrid_ex

概述

提出了一种cryoET重建框架，使用基于八叉树结构的自适应密度网格联合重建断层图像并对其进行去噪。在神经表示的优化过程中，组成八叉树结构的节点将被划分成新节点、合并或者如果为空则禁用。八叉树更新旨在通过更加重视具有更多特征的区域来加速计算。每个八叉树节点都具有可微密度网格，我们根据捕获的噪声投影对其进行优化。

本文工作借鉴了机器学习的最新进展，通过使用可微图像形成模型结合反向传播。实验表明，经典分层数据结构和可微分图像形成和误差反向传播的结合优于最近的cryoET神经表示。

为了解决冷冻电子断层成像中存在的高噪声问题，本文在密度更新期间将多尺度八叉树更新策略和TV先验和跨节点非局部约束（CNLC）相结合，还使用边界一致性先验来确保相邻八叉树节点之间的连续性。

本文提出的框架如图1，主要基于对多尺度分层方式实现的扫描对象的自适应网格表示。该结构经过更新和优化以表示连续的3D密度场。然后通过对感兴趣区域（ROI）中的体积进行均匀采样并查询这些样本位置处的密度来获得重建。训练步骤中使用的损失函数由来自断层扫描形成模型的数据保真度项和三个不同的正则化项组成，以提高重建质量。

方法

框架如图1，主要基于以多尺度分层方式实现的扫描对象的自适应网格表示。该结构经过更新和优化以表示连续的3D密度场。然后通过对感兴趣区域（ROI）中的体积进行均匀采样并查询这些样本位置处的密度来获得重建。训练中使用的损失函数由来自断层扫描形成模型的数据保真度项和三个不同的正则化项组成，以提高重建质量。

下面，首先介绍图像形成模型，以及损失函数中的不同约束：TV loss，non-local constrait， boundary consistency。然后提出了体积的自适应可微网格表示、采样策略和模型优化。

2.2 图像形成模型

对于断层扫描重建，传感器捕获的投影图像对应于沿源和不同探测器之间的射线的密度对数空间中的积分。对于给定的射线，离散化后可以使用以下方程表示：
$${\mathbf{b}}_i={\mathbf{A}}_i{\mathbf{x}}_i+{\mathbf{n}}_i$$
其中${\mathbf{b}}_i$和${\mathbf{n}}_i$分别对应于探测器i测量的强度及其相关噪声。${\mathbf{A}}_i$代表沿射线i的Radon变换操作。${\mathbf{x}}_i$是沿射线采样的3D密度值向量。对于cryoET序列，在原始捕获数据中噪声通常被假设遵循泊松噪声模型。然而，经过一些预处理操作，例如运动校正和强度校正，投影中的噪声变得更加复杂，为了简单起见可以将其视为高斯噪声。通过将所有光线重新组合在一起并应用mask来禁用与fiducial marker相交的光线，我们定义了以下数据保真度损失（data-fidelity loss）：

其中M是一个binary mask，用于将优化过程限制为不与基准标记相交的光线。

mask生成方法：
对于每个投影，设计一个mask。使用IMOD进行对齐的预处理，会输出包含基准标记位置的“*.fid”文件，然后用IMOD的 $model2point$ 命令将其转成“.txt”文件。
在此阶段，我们会仔细检查是否检测到感兴趣区域内的所有标记。对于某些数据集，输入投影已经对齐。这种情况下，我们使用IMOD中的 $imodfindbeads$ 命令来检测。
此外，还会在mask中排除填充区域，只保留感兴趣区域。如果我们知道了对齐矩阵和原始投影，则可以自动执行此步骤。然后通过合并fiducial marker上的mask和填充区域上的mask来获得最终的mask。

2.3 基于坐标的表示

基于坐标的网络在tomography问题中的应用包括将ROI区域内的3D空间坐标映射到密度场。该映射可以表示为：

在传统神经场中，$f_{\theta }$是一个全连接的MLP（图2a）。这种表示由于重建场景的大小，很快会受到限制。所以后面一些工作例如ACORN和NeAT提出使用基于八叉树的多尺度结构。这些方法中，每个八叉树节点都有一个较小的MLP或者解码器网络来表示local density。

本文框架（图2b）使用基于自适应八叉树结构的类似表示。然而，每个八叉树节点存储一个可优化的3D密度网格，以局部离散化方式表示映射函数$f_{\theta }$。对于给定的3D点$\mathbf{p}$它的密度是通过包含$\mathbf{p}$的密度网格块的8个顶点的存储密度通过三线性插值获得的，如图2b中红色虚线矩形所示。

由于插值计算是可微分的，因此可以优化所提出的3D密度表示。所以 data-fidelity loss可以写成：

N是训练步骤每次迭代时使用的采样点数量（batch size）。为了简化，本文所有密度网格定义成相同尺寸（$N_x=N_y=N_z$）。

本文表达方式的优势：

• 在大多数cryoET数据集中，扫描样本的某些区域具有详细的特征，而其他区域则是uniform的。通过使用多尺度八叉树结构去表达这些不同的区域，我们的重建能更好地保护样本中的细节，同时对均匀区域进行去噪。
• 使用插值去查询密度比使用MLP或解码器网络快得多。
• 密度网格表示允许更灵活地在损失函数中添加正则项，以处理cryoET数据中的高水平噪声。

2.4 正则化

为了降低重建中的噪声水平并确保神经场在网格边缘连续，我们在损失函数中引入了三个先验。

Total variation

TV loss 通常在断层扫描重建中用作空间正则化器来平滑重建。在神经场中不好直接使用它，因为需要在每次迭代时查询完整的体积来评估 TV loss。Zang(2021) 的工作中只利用神经场取估计缺失楔投影，然后在具有完整正弦图的经典重建中引入了 TV loss。Ruckert(2022) 的工作中提出在特征空间中计算此损失，在使用解码器获取密度之前进行。在本文实现中，将TV先验应用于每个启用的八叉树节点的密度网格，然后对所有节点进行平均，表示如下：

其中| |中的代表八叉树结构每个网格中的密度梯度，N是启用的节点数。

跨节点非局部约束

Non-local means （非局部均值）是一种经典的过滤方法，对图像中所有像素进行平均，并根据它们和目标像素的相似度进行加权。它比局部滤波方法具有更好的去噪效果。

最近，Wang（2018）的工作中引入了非局部神经网络，以改进3D视频分类任务中的特征表示。
Zhang(2019)的工作中采用残差非局部注意网络来执行多个图像恢复任务，包括图像去噪。
所提出的Non-local操作的目的是通过使用自注意力的概念总结特征网络中的相关采样特征来探索图像内部的自相似性。定义如下：

$${\mathbf{y}}_{\mathbf{i}}=softmax(f_{\varphi }({\mathbf{p}}_i{)}^T{f}_{\varphi }({\mathbf{p}}_j))f_{\varphi }({\mathbf{p}}_j{)}^T$$

其中的$f_{\varphi }$都是在特征网格中${\mathbf{p}}_i$和${\mathbf{p}}_j$位置采样的特征, ${\mathbf{y}}_i$是在位置${\mathbf{p}}_i$进行non-local 操作后的结果特征。

在本文实现中，直接将non-local操作用于大型三维体积上是不可行的，所以提出了在每次迭代时在ROI内随机选择两个启用节点。然后我们采样其密度网格内部的密度以计算非局部操作。Non-local loss 定义如下：

其中$f_{\varphi }({\mathbf{p}}_i^m)$和$f_{\varphi }({\mathbf{p}}_j^n)$分别是在节点m和节点n的密度网格内的位置$p_i$、$p_j$处采样的密度。
softmax运算将其和密度网格n中的密度向量的转置相乘，估计和密度网格m中的密度向量可能存在的相似性。

边界一致性

本文框架里的八叉树结构会在重建密度场中引入不连续伪影，因为每个节点都有自己的密度网格，并且不同八分圆之间的优化执行不一致。所以这里引入边界一致性损失（类似于NeAT中的）。

对于不同网格之间的所有边界，最小化使用每个相关密度网格计算的密度之间的差异。边界一致性损失定义如下：

Ob指所有相邻八叉树节点对，∩n,m是节点n和m之间边界表面上的采样点集合。f是对应的边界上的密度，用m和n的网格计算。

总损失函数

2.6 自适应密度网格优化

八叉树更新

第一步包括定义octree。首先八叉树从ROI初始化，屏蔽掉外面的节点（图3）。

![在这里插入图片描述](https://img-blog.csdnimg.cn/direct/6d16571c3358421ba9de5e70d494cd0a.png =700
然后均匀地对每个八分之一进行采样，并计算每个节点（密度网格）内部密度的SD，以定义更新损失。当给定节点中的SD很重要时，这意味着该节点可能包含详细的特征。所以这个节点将更有可能被分成8个子节点。

本文采用和ACORN、NeAT相同的八叉树更新约束。这包括求解一个混合整数程序，其中会根据更新损失，在每次迭代中拆分、合并或保持相同的octants，以保持总节点数低于固定限制。本文的更新损失和它们不同，在ACORN中GT密度值在图像、体积中的每个位置都是可用的，但这在断层重建中并非如此。
NeAT中使用重投影误差来定义八叉树更新损失，但是这对于cryoET不合适，因为相应的预测中存在高噪声。本文工作中发现将每个节点的SD作为八叉树更新损失是cryoET数据集的更好选择。

为了解决高噪声问题，本文进一步应用coarse-to-fine策略来加速更新。八叉树最初使用下采样投影来进行更新。再次步骤中，同时会优化每个节点内的密度网格的密度。低分辨率投影可以减少八叉树更新期间高噪声的影响。经过一些轮次后，修复了八叉树的结构，并使用原始投影只对每个节点的密度网格进行优化。在此步骤中，网格是从其下采样版本初始化的。

射线采样

在优化步骤中，对每条光线进行采样，以定义用于强度积分和损失评估的3D位置列表。这种采样并不均匀，但考虑了当前的八叉树结构。对于射线穿过的每个八叉树节点，我们沿射线选择$N_{oc}$个3D位置，分层随机采样（图4）。其中：

$N_{max}$是对应于每个节点的最大样本数的超参数，$l_{oc}$是八叉树节点$oc$内部的射线长度，$d_{oc}$是$oc$的对角线长度。

实验和结果

本文在模拟数据集和真实数据集上设计了几个实验来证明框架的有效性。和三种主要的最先进的cryoET重建技术进行了对比：

1. SART+TV。这种迭代优化的方法在缺失楔形采集方面比WFBP提供更好的结果；
2. Kniesel 提出的方法（2022）
3. NeAT(2022)
   还比较了本文方法不加跨节点非局部约束的结果。
   为了公平比较，所有重建密度进行归一化。我们进一步定性和定量地评估去噪的有效性。
   此外，提出了一个实验来评估本文方法在保留扫描样本的详细特征方面的性能。还使用FSC度量估计重建分辨率。最后，评估了方法对噪声的鲁棒性，并提供了消融实验说明框架组件的重要性。

3.1 Datasets

• EMPIAR 10643：（-60°，60°，3°），重建了40和51两个单独的序列。
• EMPIAR 10453：（-60°，60°，3°），重建一个序列。
• EMPIAR 11462：（-56°，56°，1°），重建一个序列。

3.2 数据预处理

三个数据集嵌入了基准标记来执行对齐。先对投影进行预处理，使用IMOD完成对齐，然后使用等于投影均值的灰度值来填充由于对齐而丢失的区域。对于重建，应用mask来禁用缺失的区域和基准标记。

输入投影有 4K(4096×4096)或8K(8192×8192)的分辨率。
为方便计算，EMPIAR 10643和 EMPIAR 10453的投影下采样5倍，EMPIAR 11462 的投影下采样8倍，以获得接近 1024×1024 的分辨率。还对其进行了裁剪，以便更好的关注病毒副本，同时保持所有数据集的大小相同。

3.3 去噪效果评估

图5中展示了使用不同方法重建数据集的结果比较。整体来说，除了Kniesel的方法，其他方法在重建体积上产生类似的对比度。这可以用过学习该方法的数据集和比较中使用的真实数据集之间的差异来解释。下面，重点比较这些方法对数据去噪的有效性以及在重建中保存有效特征的能力。

图中可以看到本文方法在抑制均匀区域中的噪声方面是最有效的。即使用了TV约束的SART和NeAT，均匀区域中仍包含噪声。如果增加这个约束的权重来提高去噪能力，那重建中的主要特征会丢失。和NeAT相比，本文的去噪效果可以用CNLC和八叉树更新策略来解释。

此外，表1展示了三个统计指标：

• Contrast-to-Noise Ratio (CNR)
• Equivalent Number of Look (ENL)
• Signal-to-noise Ratio (SNR)

这些指标通常用于评估去噪质量。
CNR评估降噪策略如何增加ROI和均匀背景之间对比度，ENL衡量均匀区域的平滑度，SNR指标测量重建中的信噪比。

表中结果显示本文方法具有更高的指标。增加了具有特征的区域和背景之间的对比度，并产生更平滑的均匀区域。从表中可以注意到CNLC在提高本文方法的去噪效果方面做出了巨大贡献。表中还显示了SART+TV的结果噪声最大。

3.4 详细特征分析

图5显示SART+TV擅长保留详细特征。但是当放大这些特征时，似乎淹没在背景噪声中。很难用这种方法来分析特征的结构。

NeAT重建技术在EMPIAR 10453重建中相对较好地保留了病毒的结构，即使重建中存在不均匀的对比度水平。然而，在EMPIAR 10643数据集中，受噪声影响很大，特征很难和背景区分开。

最后，该图说明了使用本文方法在保留两个数据集中的详细特征方面具有更好的性能。在不使用非局部约束的情况下，结果包含数据集的主要特征，也包含一些残余噪声。当使用非局部先验时，这种残余噪声就会减少，因此病毒的结构很容易从统一的背景中区分出来。

为了对每种方法在恢复详细特征方面的性能进行数值评估，我们分析了 EMPIAR 10643-40 数据集中病毒周围的线的重建轮廓（图6a）。该线是手动选择和采样的，以突出每种方法在保留病毒尖峰周期性结构方面的性能。

图6a显示了第二三列的方法不能将尖峰和背景分离，所以重点比较了其余方法。图6b显示了三种比较方法的概况，SART+TV方法的轮廓不如我们的两种方法。可以看到重建中残留的噪声引起的一些中间峰和谷，也很难使用此轮廓来准确定位尖峰的位置。另一方面，在有或没有非局部约束的情况下使用我们的方法会产生更规则的轮廓，尖峰和背景可以用简单的阈值过滤分开。从这个剖面分析来看，非局部约束的影响似乎很小，事实上，该先验更适合于减少均匀区域中的噪声，其中穿过体积的多个区域具有相同的统计数据。对于病毒区域只有少数区域有类似的统计数据。我们对这一先验的实施减少了其对这些地区的影响。由于我们每次选择一对随机的密度网格，因此可能会得到均匀的密度网格或者病毒的不同部分。

3.5 分辨率验证

为了评估重建精度，FSC在单颗粒分析中被广泛应用，其根据同一结构的两个不同重建计算得到：

$r_i$是傅里叶空间中半径为r的体素元素，$F_1$、$F_2$是两个体积的complex structure factors。在cryoET倾斜序列中，我们只有一份重建体积。所以为了计算其FSC，Diebolder(2015)的工作中建议将偶数和奇数投影分来以重建同一样本的两个不同体积。本文采用了这个方法并评估了 EMPIAR 10643-40 和 EMPIAR 11462 数据集的FSC。

重建分辨率定义为 FSC 保持高于给定阈值（通常选择0.5和0.143）的最高频率。这里使用 FSC=0.5，因为0.143无法区分某些重建方法的分辨率。结果表明，我们的方法用于重建时，高空间频率的幅度在从奇数和偶数投影获得的两次重构中更一致。

表2中是不同方法获得的定量分辨率估计值。尽管本文方法提供了最佳分辨率估计，但不要过度解读FSC值，因为基于进一步减少的投影数量的重建中存在高残余噪声水平。请注意FSC将在两次重建中产生一致但是错误的结构的方法排名很高。和SPA不同，cryoET的FSC度量数值不能给出绝对分辨率的可靠估计。考虑到这一点，结果似乎缺失证实了可以从图5结果得出的定性分析。

结论

本文提出了一种新的cryoET数据的自适应重建方法，使用八叉树结构，每个八分圆存储一个3D密度网格。优化时先使用投影的下采样来更新结构和3D密度网格，再冻结八叉树并使用原始投影微调密度网格。损失函数包括图像形成模型、TV损失、跨节点非局部约束和边界一致性先验。

通过使用可微网格，我们可以客服由缺失楔采集引起的伪影。另外，八叉树结构允许在合理的时间内更好地处理大量数据。同时，三个约束项大大降低了噪声水平。另外参数的调整对于避免过度降噪导致的信号丢失至关重要。

当前框架仅限于1K分辨率的投影，进一步减少内存和时间上的计算是未来研究趋势，从而重建更高分辨率的体积。此外，还会考虑联合对齐和重建的方法。

---

Supplementary

实现细节

在密度网格的优化过程中采用RMSProp优化器，学习率0.01。25个epoch后收敛，每个epoch会遍历所有可用投影。前15epoch使用下采样投影进行更新。剩下的epoch中，八叉树结构固定，用原始投影微调密度网格表示。

八叉树最大节点数设置为320。网格大小取决于要重建的体积大小。这里对真实场景设置成42，模拟数据设置成24。

模拟数据实验

用的和Kniesel相同的数据集。(-70°,69.5°,1.5°)。比较本文方法在不同网格大小和不同噪声水平下的性能。

grid size

网格尺寸越小训练速度更快，重建体积更平滑更干净。但是不能学习详细特征。网格尺寸较大时可以恢复更多细节，但重建容易过拟合而且再现噪声。需要找到折中，来更好地恢复细节特征的同时消除尽可能多的噪声。

还对不同网格尺寸的重建体积的峰值信噪比PSNR和结构相似性指标测量SSIM进行了定量评估。

noise level

真实数据集上的消融实验

TV loss

BC prior

没有BC之前有明显的边界不连续性。使用BC后伪影被消除。

CNLC prior

可以清楚看到非局部先验的去噪效果。

CNR和ENL区域的选择

手动选择。使用3D度量，将深度10添加到选定的矩形平面，如图8。CNR区域对（一个具有特征，一个具有背景噪声）红色表示；ENL所选均匀区域用黄色表示。

参数量和训练时间

参数量：

• Kniesel 等人使用的基于MLP表示，需要约2.7M参数
• NeAT：168M
• Ours：28M。虽然参数比基于MLP的方法多10倍，但是具有更快的收敛速度并且对GPU来说仍是可管理的。

表2比较了每种方法执行训练/优化和重建所需的总执行时间。由于SART+TV是基于OpenMP实现和加速并在CPU运行，这里没有比较，在作者的工作站使用64个线程大约需要1小时。
非局部约束约占方法执行时间的66%。然而即便使用此约束仍比NeAT稍快。

---

$$Author:Chier$$