对比学习作为一种无监督的表征学习方法，具有训练简单，泛化性强，可用于多种下游任务以及可扩展到多模态数据等优点，已在AI领域被广泛应用，如OpenAI的CLIP模型[1]和DeepLabCut的CEBRA模型[2]。由于神经元Sorting任务本质上是一种基于神经元波形的聚类任务，因此可以基于大量已收集到的波形数据采用对比学习方法预训练一个编码器，以更好地提取波形特征，再用于下游的聚类任务。本文将从对比学习基本原理出发，介绍对比学习方法在神经元Sorting任务中的应用以及前沿研究。
1
对比学习简介
对比学习作为一种无监督的训练方法，由于摆脱了监督学习的标签限制，因此可以基于大数据集做预训练，从而提高网络模型的泛化性。本部分将介绍对比学习的基本思想、对比学习模型训练的两大要素和应用最多的对比学习框架——SimCLR[3]。
01
对比学习的基本思想
图1：监督模型在识别一张豹子图片后的概率输出
受到监督学习模型启发，在一个已训练好的用于图片识别任务的监督学习模型中，如果令其识别一张豹子的图片，会发现模型输出概率比较高的都是与豹子有关的类别，如图1所示[4]。因此，正是数据本身存在某种“相似性”，才会使得有监督学习模型给与相似类别更高概率。
图2：对比学习的基本流程
对比学习方法期望摆脱监督学习标签的限制，将每一个样本都看作单独的类别，通过数据增强定义正负样本对，对数据本身进行表征学习。最终通过优化对比损失函数，使得“相似的”数据（正样本对）在样本空间中的距离尽可能近，而“不相似的”数据（负样本对）在样本空间中尽可能远，如图2所示[5]。
02
对比学习的两大要素 1.2.1 定义正负样本
要通过对比学习方法训练一个可表征数据本身特征的网络模型，首先需要通过数据增强定义正负样本对。如以图2左侧为例，一个样本与由它增强而来的样本之间互为正样本对，而与其他样本及其他样本的增强之间互为负样本对。三张鸟的图片,, 经过数据增强后，得到了, , 。如此，和为一对正样本，而和除外的其他所有样本之间互为负样本对。
1.2.2 定义损失函数
原始数据在经过网络前向过程变为低维特征向量后，在样本特征空间中需要使得正样本之间的距离彼此靠近，负样本之间则反之，如图2右侧。通过定义对比损失，可以优化这一过程。在有监督学习过程中，一般采用交叉熵作为损失函数，然而由于对比学习将每一个样本都看作单独的类别，因此softmax操作由于计算复杂度太高实际上是不可行的。因此，一些专用的对比损失函数如NCE loss (noise contrastive estimation loss)[6]和InfoNCE loss[7]等被提出，此处仅介绍常用的InfoNCE：
其中，τ是调节分布形状的超参数，为用于计算相似性的点积运算，为一个样本经编码器编码后输出的特征，为此样本的正样本输出的特征，为包括和的其他负样本的所有样本输出的特征。
03
一种简单有效的对比学习框架 —— SimCLR
图3：SimCLR框架
SimCLR 框架[3]因其简单有效而被广泛应用。训练过程如图3所示，样本x在经过数据增强变为和，后经过编码器编码为特征向量和，最后再经过一个全连接层得到和。训练目标是最大化和之间的一致性，在SimCLR框架中，对比损失采用normalized temperature-scaled cross entropy loss，为InfoNCE的改进。
模型训练完成后，只有编码器被用于新数据集以得到特征向量用于下游任务。
2
基于SimCLR架构的Sorting模型
Sorting任务本质上是对神经元的波形进行聚类分析，以得到每个神经元的放电时间戳。对比学习方法可以基于大量已经提取到的无标签的波形数据进行表征学习，预训练一个编码器，对波形数据进行表征，用于下游聚类任务。本部分将基于已发表在NeurIPS 2023的一项工作[8]简要介绍对比学习方法在神经元sorting任务中的应用。
图4：基于对比学习方法的sorting
如图4所示，基于SimCLR框架，对提取到的波形数据进行数据增强可以得到正负样本对，之后便可优化对比损失调整网络参数。研究人员将此方法命名为CEED (Contrastive Embeddings for Extracellular Data)。
此研究采用的原始数据集来自IBL（International Brain Laboratory）数据库[9]，为Neuropixels 1.0电极记录到的小鼠神经元波形数据，ground truth由Kilosort 2.5 sorting而来。在此基础上，为了对比不同方法之间的性能，研究人员设计了3种对比数据集，分别为①10个神经元的训练集和同分布的10个神经元测试集，②400个神经元的训练集和同分布的10个神经元测试集，③390个神经元的训练集和训练集分布之外的10个神经元测试集。
表1：CEED方法评估（编码器采用MLP架构）
为了评估CEED性能，研究人员以主成分分析（PCA）和自编码器（Autoencoder）给出的特征为基准进行对比（输出特征的维度都为5）。下游采用高斯混合模型（GMM）对编码器输出的特征进行聚类分析。最终，进行50次GMM聚类后的ARI（Adjusted Rand Index）均值和标准差如表1所示。结果显示，在下游采用相同聚类算法的前提下，基于对比学习方法输出的神经元波形特征在3种数据集上效果最优，聚类精度最高。
References
[1] Radford A, Kim JW, Hallacy C, Ramesh A, Goh G, Agarwal S, Sastry G, Askell A, Mishkin P, Clark J, Krueger G. Learning transferable visual models from natural language supervision. International conference on machine learning 2021 Jul 1 (pp. 8748-8763). PMLR.
[2] Schneider S, Lee JH, Mathis MW. Learnable latent embeddings for joint behavioural and neural analysis. Nature. 2023 May 11;617(7960):360-8.
[3] Chen T, Kornblith S, Norouzi M, Hinton G. A simple framework for contrastive learning of visual representations. International conference on machine learning 2020 Nov 21 (pp. 1597-1607). PMLR.
[4] Wu Z, Xiong Y, Yu SX, Lin D. Unsupervised feature learning via non-parametric instance discrimination. Proceedings of the IEEE conference on computer vision and pattern recognition 2018 (pp. 3733-3742).
[5] Ye M, Zhang X, Yuen PC, Chang SF. Unsupervised embedding learning via invariant and spreading instance feature. Proceedings of the IEEE/CVF conference on computer vision and pattern recognition 2019 (pp. 6210-6219).
[6] Gutmann, M. and Hyvärinen, A., 2010, March. Noise-contrastive estimation: A new estimation principle for unnormalized statistical models. Proceedings of the thirteenth international conference on artificial intelligence and statistics (pp. 297-304). JMLR Workshop and Conference Proceedings.
[7] Oord AV, Li Y, Vinyals O. Representation learning with contrastive predictive coding. arXiv preprint arXiv:1807.03748. 2018 Jul 10..
[8] Vishnubhotla A, Loh C, Srivastava A, Paninski L, Hurwitz C. Towards robust and generalizable representations of extracellular data using contrastive learning. Advances in Neural Information Processing Systems. 2024 Feb 13;36.
[9] https://viz.internationalbrainlab.org/app
本文作者：王兴琪 ｜华东师范大学邮箱：hsingchi.wang@gmail.com
欢 迎 沟 通 交 流 ！
发表回复