Memory-Efficient Prompt Tuning for Incremental Histopathology Classification
Authors: Yu Zhu, Kang Li, Lequan Yu, Pheng-Ann Heng
Source: The Thirty-Eighth AAAI Conference on Artificial Intelligence (AAAI-24)
Abstract
最近的研究在组织病理学分类方面取得了显着进展。 基于当前的成功，当代的工作提出通过从顺序交付的领域中增量学习，进一步将模型升级到更通用和更稳健的方向。 与之前基于参数隔离的方法在模型更新过程中通常需要大量计算资源不同，我们提出了一种内存高效的提示调整框架，以经济的内存成本培养模型泛化潜力。 对于每个传入域，我们重用初始分类模型的现有参数并附加轻量级可训练提示进入其中进行定制调整。 考虑到领域异构性，我们进行解耦提示调整，我们为每个域采用特定于域的提示来独立研究其独特的特征，并在跨域共享一个域不变的提示所有领域不断探索共同的内容嵌入。 所有特定于域的提示都将附加到提示库中，并与进一步的更改隔离，以防止忘记早期看到的域的独特特征。 而领域不变的提示将通过风格增强的提示引用来传递和迭代发展，以随着时间的推移提高模型的泛化能力。 具体来说，我们用现有的提示构建一个图，并构建一个风格增强的图注意网络来引导域不变提示探索所有交付域之间的重叠潜在嵌入，以获得更多域通用表示。 我们通过两项组织病理学任务（即乳腺癌转移分类和上皮-间质组织分类）广泛评估了我们的框架，其中我们的方法比竞争方法产生了卓越的性能和记忆效率。
#背景
组织病理学图像分类是癌症诊断中的一项基本任务，其目的是通过显微镜检查来确定可疑组织的恶性和良性。深度学习模型在该任务上取得了显著进展，但当代研究并不满足于此，而是致力于通过增量学习来升级和更新模型功能，使其更加完善。
挑战
- 数据异质性: 不同医疗中心的组织学数据由于技术人员技能和数字扫描仪品牌的差异，在苏木精和伊红（H&E）染色后呈现出不同的外观，从深蓝紫色到浅粉紫色不等。
- 领域增量学习（DIL）: 模型需要逐步适应越来越多的异质领域。
- 灾难性遗忘: 模型更新过程中需要避免遗忘先前学习的知识。
- 数据隐私: 由于数据隐私问题和存储负担，模型更新时无法访问早期数据。
- 内存消耗: 现有方法通常需要大量的计算资源和内存，限制了其在千兆像素级组织病理学图像中的应用。
贡献
本文提出了一种基于提示调优的内存高效增量学习框架，用于组织病理学图像分类。该框架能够逐步提升模型的泛化能力，同时保持较低的内存消耗。
主要贡献:
- 内存高效的提示调优框架: 该框架通过在初始模型上进行提示调优，避免了对模型参数的直接修改，从而节省了内存消耗。
- 解耦提示调优: 该框架采用域特定提示和域不变提示，有效地处理了输入域的异质性。
- 风格增强的提示细化: 该框架通过风格增强的提示细化方法，迭代地演化域不变提示，提升了模型在未见过数据上的泛化能力。
- 实验验证: 该框架在两个组织病理学图像分类任务上取得了优于其他比较方法的性能，表现出更好的泛化能力和更少的遗忘。
Methodology
领域增量学习方法论总结 (Domain Incremental Learning Methodology Summary)
目标： 在领域增量学习 (DIL) 环境中，模型需要在多个不同领域的数据流 (D1, D2, ..., DT) 上进行增量学习，并避免灾难性遗忘，同时保持对未见领域数据的泛化能力。由于隐私问题，模型无法访问过去领域的数据，也无法得知当前数据的领域标签。
方法： 本文提出了一种基于解耦提示调整 (Decoupled Prompt Tuning) 的方法。模型包含一个 Transformer 主干网络 (例如 ViT) 和两个可训练的提示：
- 领域特定提示 (Domain-Specific Prompt, DSP): 用于学习每个领域独特的特征，每个领域都有一个独立的 DSP。
- 领域不变提示 (Domain-Invariant Prompt, DIP): 用于学习所有领域共有的特征，通过风格增强提示精炼 (Style-Augmented Prompt Refining) 迭代更新。
图1：我们的内存高效提示调优框架的概述。我们提出使用两个轻量级提示对初始模型进行解耦提示调优，旨在在经济的内存成本下获得最新的领域知识。我们为每个域独立地使用一个域特殊提示符，以获得其独特的特征，如外观。学习到的域特殊提示将被存储和隔离在提示银行中，以帮助减轻对早期获得的域的遗忘。同时，我们保持一个跨域共享的域不变提示符，以便随着时间的推移逐步学习公共内容，就像形状先验一样。我们在之前的域不变提示上执行了风格增强提示重构，其中我们将其探索范围限制在所有可见域的重叠潜在嵌入中，并指导它学习域通用表示，随着时间的推移逐渐增强泛化潜力。
解耦提示调整：
- 模型初始化: 使用预训练模型初始化 Transformer 主干网络，并冻结其参数。
- 领域特定提示学习: 对于每个新领域，独立学习一个新的 DSP，使用公式 (5) 进行优化：
min p t s L c e ( f ϕ ( f b ( x ; p t ( I ) , p t s ) ) , y ) , \min_{p_t^s} L_{ce}\left(f_{\phi}\left(f_b\left(x; p_t^{(I)}, p_t^s\right)\right), y\right), ptsminLce(fϕ(fb(x;pt(I),pts)),y),
其中 L c e L_{ce} Lce 表示交叉熵损失， f ϕ f_{\phi} fϕ 表示分类层， f b f_b fb 表示 Transformer 主干网络， x x x 表示输入图像， p t ( I ) p_t^{(I)} pt(I) 表示当前 DIP， p t s p_t^s pts 表示当前 DSP， y y y 表示标签。
- 领域不变提示精炼: 使用风格增强图注意力网络 (Style-Augmented Graph Attention Network) 更新 DIP。
风格增强图注意力网络:
- 图构建: 将所有现有的提示 (包括 DIP 和所有 DSP) 作为图的节点。
- 注意力系数计算: 计算 DIP 与其他节点之间的注意力系数：
e i I S = a ( W p t ( I − 1 ) , W p i s ) , e I I = a ( W p t ( I − 1 ) , W p t ( I − 1 ) ) , e_i^{IS} = a\left(Wp_t^{(I-1)}, Wp_i^s\right), \\ e_{II} = a\left(Wp_t^{(I-1)}, Wp_t^{(I-1)}\right), eiIS=a(Wpt(I−1),Wpis),eII=a(Wpt(I−1),Wpt(I−1)),
其中 a a a 表示单层神经网络， W W W 表示线性变换矩阵。
图2：生成样式增强数据的说明。
- 知识聚合: 使用将其他节点的知识聚合到 DIP：
p t ( I ) = f G A T ( P t ) = ∑ i = 1 t α i I S W p i s + α I I W p t ( I − 1 ) , p_t^{(I)} = f_{GAT}(P_t) = \sum_{i=1}^t \alpha_i^{IS}Wp_i^s + \alpha_{II}Wp_t^{(I-1)}, pt(I)=fGAT(Pt)=i=1∑tαiISWpis+αIIWpt(I−1),
其中 f G A T f_{GAT} fGAT 表示图注意力网络， α i I S \alpha_i^{IS} αiIS 和 α I I \alpha_{II} αII 表示归一化后的注意力系数。
- 风格增强训练: 对输入图像进行风格增强，并将增强后的图像用于训练图注意力网络。
推理: 对于测试图像，首先选择最匹配的 DSP，然后使用 DIP 和选择的 DSP 进行预测。
总结： 该方法通过解耦提示调整和风格增强图注意力网络，有效地解决了领域增量学习中的灾难性遗忘问题，并提升了模型对未见领域数据的泛化能力。
实验部分:
- 与现有方法的比较:
- 本文提出的方法在乳腺癌转移分类任务的9个评价指标中的8个上取得了最佳效果,在上皮间质组织分类任务的8个评价指标中的7个上取得了最佳效果。
- 与现有的参数隔离方法相比,本文方法在缓解灾难性遗忘和提高模型泛化能力方面表现出色。
- 关键组件的分析:
- 在模型性能和内存效率的权衡方面,本文方法在使用有限的额外内存空间的情况下,相比现有的基于提示的方法,在平均准确率、向后传递性能和前向传递性能方面都有显著提升。
- 解耦提示调优的可视化结果表明,学习到的特定领域提示可以有效捕捉领域特征,模型可以很好地区分正常组织和肿瘤组织。
- 消融实验表明,通过GAT对通用不变提示进行精炼,并使用风格增强的训练数据,可以进一步提高模型性能。
总的来说,实验结果全面展示了本文提出方法的有效性。在多个医学图像分析任务上,该方法在缓解灾难性遗忘和提高模型泛化能力方面优于现有的增量学习方法,同时保持了较高的内存效率。关键组件的分析验证了方法设计的合理性。
Reference
[1] Zhu, Y., Li, K., Yu, L., & Heng, P. A. (2024, March). Memory-Efficient Prompt Tuning for Incremental Histopathology Classification. In Proceedings of the AAAI Conference on Artificial Intelligence (Vol. 38, No. 7, pp. 7802-7810).