[論文レビュー] 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特徴量を用いた自然言語処理タスクにおいても、最小限の微調整で強力なベースラインを上回る優れた結果を示している。

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近似を用いて有限次元の埋め込みが得られる。
• 基準セットは教師ありまたは教師なしの方法でエンドツーエンドで最適化される。
• 本手法は、アテンションの一種として解釈可能であり、古典的なOTベースのカーネルの代用としても機能する。
• 教師なし事前学習と、シンプルな線形分類器を用いた微調整の両方をサポートする。

実験結果

リサーチクエスチョン

• 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％の精度を達成し、事前学習されたBERT特徴量における平均プーリング（85.4％）および[CLS]トークン（84.6％）を上回った。
• タンパク質フォールド認識タスクでは、最先端の結果を達成し、ラベルが少ない長大な生物学的シーケンスにおいても有効であることを示した。
• クロマチンプロファイル検出においても有望な結果が得られ、ゲノム分野における本手法の有効性をさらに裏付けた。
• 学習時間は効率的で、小規模モデルでは1時間未塔、大規模モデルでは3時間未塔（1GPU）で実行可能であった。
• 単層の自己アテンションベースライン（83.7％）およびOTKEのドット積バージョン（86.9％）を上回ったことから、メカニズムにおける最適輸送の重要性が確認された。