前几篇博客中已经介绍了ViT、MAE两个模型，最近看到一篇和笔者研究方向相关的基础模型的论文——基于ViT使用OCT图像进行层分割的基础模型，这个任务可以简单理解为从OCT的B扫描图像中分割出不同的层结构（在医学中有明确的定义），层结构的分割结果可以为疾病诊断提供见解（一些眼底疾病会导致层结构出现明显的变形），而笔者认为任务难点主要有：
- 朴素的语义分割方法难以获得明确的层次结构，在不确定性大的区域会出现层断裂等结果，导致无法获得连续的层结构；
- 面对疾病样本的分割效果很差，可能是病变少样本的原因；
而最近看到一篇2024年发表在DAGM German Conference on Pattern Recognition上的论文《Foundation Models Permit Retinal Layer Segmentation Across OCT Devices》，使用OCT数据进行视网膜层分割的基础模型，作者是来自波恩大学的团队，下面就开始正文。
摘要
在眼科疾病诊断过程中，从OCT（Optical Coherence tomography，光学相干断层扫描）图像中分割出视网膜层结构是非常关键的一步。当前CNN-based模型在同种OCT设备采集的图像上表现很好，但是如果遇到其他设备采集的图像，性能会大幅降低。而作者构建了首个基于ViT的、用于OCT层分割任务的基础模型。通过试验验证了，和目前的SOTA CNN模型相比，模型在推理阶段对未见分布域数据的泛化性得到显著提升，表现了基础模型进行更加稳健的医学图像分割的潜力。同时作者分析了不同基础模型的效果，值得注意的一点是：在通用任务上训练的基础模型比只使用OCT图像训练得到的基础模型效果更好。
引言
OCT图像由于非侵入性、高分辨率，在可视化眼底结构时有重要价值，利用OCT图像的关键挑战在于对其中视网膜层的准确分割，其对于疾病诊断和病后追踪至关重要。
在层分割和其他任务中最主要的问题在于：当测试数据和训练数据处于不同的分布时，模型性能会大幅降低。并且这种域偏移问题会因为多种条件而产生：不同的成像条件、患者的自身情况、模型实际部署的环境等。
最近，ViT（Vision Transformer）已经成为了视觉领域非常有效的工具，和传统的CNN不同，ViT使用自监督机制直接建模图像中的长距离依赖关系，因此为图像分析提供了更加有效、可扩展的模式。ViT的可捕获全局上下文的能力使其特别适用于需要空间细节理解的任务，例如从OCT中分割出层结构。另一方面，从头开始训练这样的模型非常贵，例如最初的ViT的 batch size为4096。
因此为了利用ViT优势，最常用的方法是使用预训练模型。这些基于ViT的基础模型推动了稳健性、泛化性更好的医学图像分析模型的构建，因为其已经在大规模数据集上得到训练，捕获了多种视觉模式和特征。广泛的预训练使得基础模型可以更好地泛化到未见数据域中，有效解决域偏移问题。利用预训练阶段学到的特征表征，这些模型可以用小数据集微调来应用于特定任务，从而推动在不同临床场景中的应用。
作者的主要贡献在于构建了OCT图像中视网膜层分割的基础模型，通过对预训练的encoder进行低秩自适应（Low-rank Adaping)微调、使用两个Transformer块和处理特定任务的网络头构成的解码器，来实现基础模型的构建。对于Encoder，有非常多的预训练好的ViT可以使用，因此可以使用多种规模的模型，但是也提出了一个问题：哪种模型会得到最佳结果。
作者的第二个贡献是对比了在不同基础模型上构建的层分割模型的性能，在同域和异域数据中都进行了广泛的实验，对多种模型和数据集都给出了定性分析，在传统方法的基线模型中选择了一个专为OCT层分割设计的CNN网络）——FeatureRefinementNet（FRN）。
总结
作者提出了一个用于OCT层分割的ViT-based基础模型，验证了和传统CNN-based方法相比泛化性得到明显提升，特别是存在数据域偏移的情况下。
ViT-based的基础模型在稳健的、可泛化的医学图像分析中展现出了潜力，但是仍然存在问题，例如遇到训练数据中不存在的病理数据。同时不同基础模型之间的对比，展示了基础模型的选择非常重要，特别地出现了意料之外的情况——在通用任务中预训练的模型比任务专业的预训练模型效果更好。
方法
数据集和预处理
所有模型都在Duke数据集上训练，包含正常和AMD（年龄性黄斑病变）样本（115：269），带有3个层的标注（ILM、IBPRE、BM），所有的数据在Bioptigen设备上采集。数据集的划分和之前的一致（164：20：200），删除某些标注不足50%的B扫描图像后，训练数据共包含8928张B扫描图像。在域内验证中使用Duke测试集，域偏移验证使用两个公开OCT数据集定量对比。（OCT5k和AROI，笔者不做过多说明，只需要注意数据集的成像设备不同）另外在RETOUCH数据集中进行定性对比。
B扫描图像被预处理为224*224分辨率的图像，进行通道归一化、数据增强等操作。
模型结构
Baseline选择CNN-based模型FeatureRefinementNet(FRN)，使用CNN学习全局上下文、通过Layer Head实现层分割。Layer Head的设计思路为：层上方区域标记为1、下方标记为0，那么最终可以通过在列方向求和得到层的位置。
作者提出的模型中，使用预训练好的ViT框架作为特征编码器，两个或更多的transformer块作为解码器。所有的transformer块都使用1024的 embedding dim、16个注意力头，并处理16*16大小的patch。解码器的输出会投影到固定长度，然后通过Layer Head恢复到原始图像，类似自监督预训练的思路。
作者选择的编码器有多种：
- RETFound：ViT-large编码器在OCT数据集上进行预训练，原模型处理分类任务，但通过masked auto-encoder的思路训练，可以提取图像的详细和局部信息。
- ViT-large模型（和RETFound同规模），但在ImageNet上预训练。
- SAM（Segment Anything Model)，探索了两种规模：Large和Base。
- MedSAM：使用SAM模型的权重初始化，在大规模、多分布的医学图像上训练。（虽然其中使用的OCT图像很少，但是作者认为医学图像中的灰度特性等比较相近，可能带来性能上的提升）
SAM的设计中包含为自注意力设计的额外添加的相对位置编码，由于其他模型中没有这些参数，因此在其他模型中随机初始化相对位置编码并从头训练这些参数。除此之外，原始的编码器参数在训练中冻结。另外会给编码器的线性层添加低秩自适应（LoRA），来减少编码器微调成本。（具体的LoRA参数设计这里略去，未来有机会单独看一下这篇工作）
另外以MAE encoder为例说一下模型的计算过程，首先对图像进行patch、添加位置编码、随机mask，经过encoder处理，接着encoded feature和 mask tokens（可学习参数）输入到解码器，最后经过Layer Head得到层分割结果。这里需要注意的是，和ViT的推理阶段不同（不会进行mask操作），这里的mask是为了保持分布一致性，避免输入全部patch tokens导致分布差异过大，影响解码器的学习。那么这里模型可学习的参数包括：
- 相对位置编码；
- LoRA矩阵；
- mask tokens；
- 解码器的参数；
具体的实验参数不作详细介绍，下面来看一下实验结果。
实验结果
列出几点比较重要的结论：
- 在同域测试集上，CNN-based模型性能最好，但和vit-based模型的差异不大；
- 异域数据集上，vit-based模型性能显著优于CNN-based，说明作者提出的基础模型泛化性显著提升；
- 不同的encoder选择中，SAM-L效果最好，即在自然图像上预训练得到的encoder好于只使用OCT数据预训练的；
讨论
当测试数据是同域时，所有的方法都取得了较小的中位数误差和95%的分位数，基于ViT和CNN的模型在分割精度上差异很小，表明在训练数据足够多、没有域偏移的情况下，OCT层分割不需要大量的预训练。
而在域偏移情况下，作者提出的基础模型优势明显。在两个域外数据集中，ViT-based模型都显著优于CNN-based（FRN），后者会出现显著的性能退化。FRN的95%误差分位数数值很大，表明其在扩展数据情况下完全失败，而ViT-based模型可以实现成功泛化。
在ViT-based模型中，在大的、多样的自然图像数据上预训练的SAM-L和MAE-L编码器都可以取得最佳泛化性。这些模型在自然图像上预训练，而自然图像和OCT图像非常不同，因此作者任务使用在大规模OCT数据上预训练的RETFound（这是另一个分别为彩色眼底图像和OCT图像设计两个不同基础模型的方法，其下游任务是疾病分类、病后预测和系统疾病预测）可能更适合处理层分割任务，但实验却没有达到预期的效果，考虑到数据有限无法进行更多的实验，只能对背后的原因进行推测。
作者提到其实验结果可能反映了RETFound使用的训练数据的偏差，尽管其使用了来自不同设备的OCT图像，但其中某些成像类型占主导地位（RETFound原论文中确实写到大部分数据均来自同一个数据集中）。这可能可以解释为什么通用数据中训练的模型可以取得最佳的泛化结果。（但实际上，我觉得除了数据分布的原因外，可能和数据量也非常相关，ViT原论文中就展示了只有在数据规模非常大的情况下才会取得超越CNN的结果）作者特地选择了没有在RETFound和MedSAM中使用的数据集以便进行对比实验，这可能也是RETFound不如其他基础模型的原因。
除了训练数据的差异外，几种基础模型的训练机制也不同，例如：
- RETFound的
batch_size为1792，学习率为1e-3 - MAE使用
batch_size为4096，学习率为1e-4
在对比实验中发现，更大的SAM-L encoder比更小的SAM-B效果更好，特别是在域偏移问题下。同时对比了这两个模型使用和不使用低秩自适应（LoRA）情况下的性能，发现两个模型之间的差异始终存在。
SAM-B和MedSAM-B的结果很相似，表现出MedSAM-B在医学图像上的额外预训练没有为本任务带来另外的帮助，可能和MedSAM-B的训练数据中OCT图像较少有关。
最后虽然作者提出的基础模型的泛化性得到提升，但是域偏移情况下的性能仍然不如域内的表现，精度有明显下降，未来仍需要进一步探索。
后记
笔者认为这篇论文的特色在于针对OCT层分割任务尝试了多种编码器的选择，进行广泛的实验（模型结构没有太多改进），通过微调技术得到该任务下的基础模型，可以提升域偏移问题下的泛化性。同时在对比实验中发现自然图像上预训练得到的模型反而优于只使用OCT预训练的模型和其他医学图像预训练模型，这一点作者也做出了比较合理的猜测。总之，这是笔者看到的第一篇和自己工作非常相近的基础模型的工作，未来有机会（也有资源）的条件下，或许可以尝试一下这个方向^_^