[论文解读] A Trainable Optimal Transport Embedding for Feature Aggregation and its Relationship to Attention

Grégoire Mialon, Dexiong Chen+2 more|arXiv (Cornell University)|Jun 22, 2020

Machine Learning in BioinformaticsBiochemistry, Genetics and Molecular Biology被引用 17

一句话总结

本文提出 OTKE，一种可训练的最优传输核嵌入方法，通过最优传输将可变大小的特征集对齐到可学习的参考集，从而学习固定大小的表示。该方法在蛋白质折叠识别和染色质谱检测任务上达到最先进性能，并在使用 BERT 特征的 NLP 任务中表现出色，仅需极少微调即可超越强基线模型。

ABSTRACT

We address the problem of learning on sets of features, motivated by the need of performing pooling operations in long biological sequences of varying sizes, with long-range dependencies, and possibly few labeled data. To address this challenging task, we introduce a parametrized representation of fixed size, which embeds and then aggregates elements from a given input set according to the optimal transport plan between the set and a trainable reference. Our approach scales to large datasets and allows end-to-end training of the reference, while also providing a simple unsupervised learning mechanism with small computational cost. Our aggregation technique admits two useful interpretations: it may be seen as a mechanism related to attention layers in neural networks, or it may be seen as a scalable surrogate of a classical optimal transport-based kernel. We experimentally demonstrate the effectiveness of our approach on biological sequences, achieving state-of-the-art results for protein fold recognition and detection of chromatin profiles tasks, and, as a proof of concept, we show promising results for processing natural language sequences. We provide an open-source implementation of our embedding that can be used alone or as a module in larger learning models at https://github.com/claying/OTK.

研究动机与目标

• 解决在长程依赖关系强且标注数据有限的可变大小生物序列和自然语言序列上进行学习的挑战。
• 开发一种固定大小、可微分的嵌入方法，以支持下游任务的端到端训练。
• 将最优传输与核方法结合，实现非线性特征变换和可扩展的池化操作。
• 提供一种计算开销低的简单无监督预训练机制。
• 在长生物序列上验证方法的有效性，并实现向 NLP 任务的迁移。

提出的方法

• 该方法使用最优传输计算输入特征集与可学习参考集之间的传输计划。
• 在使用传输计划作为权重进行池化之前，先将特征嵌入到再生核希尔伯特空间（RKHS）。
• 通过 Nyström 近似获得有限维嵌入以实现可扩展性。
• 参考集以端到端方式在有监督或无监督设置下进行优化。
• 该方法具有可解释性，可视为一种注意力机制，或作为经典基于最优传输的核函数的代理。
• 该方法支持无监督预训练和使用简单线性分类器的微调。

实验结果

研究问题

• RQ1 能否利用最优传输构建一种可训练、可微分且可扩展的嵌入方法，用于可变大小的特征集？
• RQ2 与注意力机制和经典核方法相比，所提出的 OTKE 嵌入在性能和训练效率方面表现如何？
• RQ3 该方法能否在标注数据有限的长生物序列上实现最先进性能？
• RQ4 参考集的无监督预训练是否能提升在 NLP 任务中的下游性能？
• RQ5 该方法能否泛化到多种序列类型，包括蛋白质和染色质谱？

主要发现

• 在 SST-2 情感分类基准上，监督式 OTKE 嵌入在使用 4 个参考点和 10 个支持点时达到 88.68% 的准确率，优于所有基线模型，包括 Set Transformer（87.9%）和 RepSet（87.1%）。
• 无监督 OTKE 嵌入在 SST-2 上使用 300 个支持点时达到 87.2% 的准确率，超过平均池化（85.4%）和 [CLS]（84.6%）在预训练 BERT 特征上的表现。
• 在蛋白质折叠识别任务中，该方法取得了最先进结果，证明了其在少量标签的长生物序列上的有效性。
• 该方法在染色质谱检测任务中也表现出色，进一步验证了其在基因组学中的实用性。
• 训练时间高效：小型模型训练时间不足 1 小时，大型模型训练时间不足 3 小时（单张 GPU）。
• 该方法优于单层自注意力基线模型（83.7%）和 OTKE 的点积变体（86.9%），证实了最优传输机制在该方法中的关键作用。