[Paper Review] 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 Bioinformatics17 citations

TL;DR

This paper proposes OTKE, a trainable optimal transport kernel embedding that learns fixed-size representations for variable-sized sets of features by aligning them to a learnable reference via optimal transport. It achieves state-of-the-art performance on protein fold recognition and chromatin profile detection, and shows strong results in NLP tasks using BERT features, outperforming strong baselines with minimal fine-tuning.

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.

Motivation & Objective

• To address the challenge of learning on variable-sized biological and natural language sequences with long-range dependencies and limited labeled data.
• To develop a fixed-size, differentiable embedding that enables end-to-end training for downstream tasks.
• To combine optimal transport with kernel methods for non-linear feature transformation and scalable pooling.
• To provide a simple unsupervised pretraining mechanism with low computational overhead.
• To demonstrate effectiveness on long biological sequences and transfer to NLP tasks.

Proposed method

• The method uses optimal transport to compute a transport plan between input feature sets and a trainable reference set.
• Features are embedded into a reproducing kernel Hilbert space (RKHS) before pooling using the transport plan as weights.
• A finite-dimensional embedding is obtained via Nyström approximation for scalability.
• The reference set is optimized end-to-end in a supervised or unsupervised manner.
• The approach is interpretable as a form of attention or as a surrogate for classical OT-based kernels.
• The method supports both unsupervised pretraining and fine-tuning with a simple linear classifier.

Experimental results

Research questions

• RQ1 Can optimal transport be used to create a trainable, differentiable, and scalable embedding for variable-sized feature sets?
• RQ2 How does the proposed OTKE embedding compare to attention mechanisms and classical kernel methods in terms of performance and training efficiency?
• RQ3 Can the method achieve state-of-the-art results on long biological sequences with limited labeled data?
• RQ4 Does the unsupervised pretraining of the reference set improve downstream performance on NLP tasks?
• RQ5 Can the method generalize to diverse sequence types, including protein and chromatin profiles?

Key findings

• On the SST-2 sentiment classification benchmark, the supervised OTKE embedding achieved 88.68% accuracy with 4 references and 10 supports, outperforming all baselines including Set Transformer (87.9%) and RepSet (87.1%).
• The unsupervised OTKE embedding achieved 87.2% accuracy on SST-2 with 300 supports, surpassing mean pooling (85.4%) and [CLS] (84.6%) on pre-trained BERT features.
• For protein fold recognition, the method achieved state-of-the-art results, demonstrating effectiveness on long biological sequences with few labels.
• The method showed promising results on chromatin profile detection, further validating its utility in genomics.
• Training time was efficient: under 1 hour for small models and under 3 hours for larger models on a single GPU.
• The approach outperformed a single-layer self-attention baseline (83.7%) and a dot-product variant of OTKE (86.9%), confirming the importance of optimal transport in the mechanism.