TTT-KD: Test-Time Training for 3D Semantic Segmentation through Knowledge Distillation from Foundation Models

📄 arXiv: 2403.11691v1 📥 PDF

作者: Lisa Weijler, Muhammad Jehanzeb Mirza, Leon Sick, Can Ekkazan, Pedro Hermosilla

分类: cs.CV

发布日期: 2024-03-18

💡 一句话要点

提出TTT-KD以解决3D语义分割中的数据分布适应问题

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 3D语义分割 知识蒸馏 测试时训练 自监督学习 数据分布适应

📋 核心要点

  1. 现有的3D语义分割方法在面对测试时数据分布变化时，缺乏有效的适应机制，导致性能下降。
  2. 本文提出TTT-KD方法，通过知识蒸馏自监督目标，实时更新3D分割网络以适应测试数据分布。
  3. 在多个基准测试中，TTT-KD在分布相似的情况下提升了最多13%的mIoU，在OOD样本中提升了最多45%。

📝 摘要（中文）

测试时训练（TTT）旨在实时适应预训练网络以应对变化的数据分布。本文提出了首个针对3D语义分割的TTT方法TTT-KD，该方法将知识蒸馏（KD）从基础模型（如DINOv2）建模为自监督目标，以适应测试时的数据分布变化。在拥有配对的图像-点云（2D-3D）数据的情况下，首先优化3D分割主干网络以进行语义分割任务，同时利用现成的2D预训练基础模型进行2D到3D的知识蒸馏。在测试时，TTT-KD通过知识蒸馏的自监督任务更新3D分割主干网络，以便在最终预测之前对每个测试样本进行适应。大量在多个室内和室外3D分割基准上的评估表明，TTT-KD在分布内（ID）和分布外（OOD）测试数据集上均提升了性能。

🔬 方法详解

问题定义：本文旨在解决3D语义分割在测试时面对数据分布变化的适应性不足问题。现有方法通常依赖于固定的训练数据分布，无法有效应对测试时的变化，导致性能下降。

核心思路：TTT-KD通过引入知识蒸馏作为自监督学习目标，使得3D分割网络能够在测试时根据输入数据的特征进行动态调整，从而提高模型的适应能力。

技术框架：TTT-KD的整体架构包括两个主要阶段：首先，在训练阶段优化3D分割主干网络和2D到3D的知识蒸馏；其次，在测试阶段，针对每个测试样本进行自监督知识蒸馏以更新分割网络。

关键创新：TTT-KD的创新之处在于将知识蒸馏与测试时训练结合，形成了一种新的自监督学习机制，使得模型能够在测试时实时适应数据分布的变化，这在现有文献中尚属首次。

关键设计：在实现中，TTT-KD采用了配对的图像-点云数据进行训练，使用了标准的损失函数来优化分割性能，并结合了现成的2D预训练模型进行知识蒸馏，确保了模型在不同数据分布下的鲁棒性。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

TTT-KD在多个3D分割基准测试中表现出色，在分布相似的情况下，mIoU提升了最多13%，而在分布外样本中，提升幅度更是达到45%。这些结果表明，TTT-KD在应对数据分布变化方面具有显著的优势。

🎯 应用场景

TTT-KD方法在3D语义分割领域具有广泛的应用潜力，尤其是在自动驾驶、机器人导航和增强现实等场景中，能够有效应对实时数据变化，提高系统的智能化水平。未来，该方法还可以扩展到其他领域，如医学影像分析和环境监测等。

📄 摘要（原文）

Test-Time Training (TTT) proposes to adapt a pre-trained network to changing data distributions on-the-fly. In this work, we propose the first TTT method for 3D semantic segmentation, TTT-KD, which models Knowledge Distillation (KD) from foundation models (e.g. DINOv2) as a self-supervised objective for adaptation to distribution shifts at test-time. Given access to paired image-pointcloud (2D-3D) data, we first optimize a 3D segmentation backbone for the main task of semantic segmentation using the pointclouds and the task of 2D $\to$ 3D KD by using an off-the-shelf 2D pre-trained foundation model. At test-time, our TTT-KD updates the 3D segmentation backbone for each test sample, by using the self-supervised task of knowledge distillation, before performing the final prediction. Extensive evaluations on multiple indoor and outdoor 3D segmentation benchmarks show the utility of TTT-KD, as it improves performance for both in-distribution (ID) and out-of-distribution (ODO) test datasets. We achieve a gain of up to 13% mIoU (7% on average) when the train and test distributions are similar and up to 45% (20% on average) when adapting to OOD test samples.