[논문 리뷰] Differentiable Top-k Operator with Optimal Transport

Yujia Xie, Hanjun Dai외 6명|arXiv (Cornell University)|2020. 02. 16.

Domain Adaptation and Few-Shot Learning참고 문헌 37인용 수 15

한 줄 요약

이 논문은 엔트로픽 최적 운반(Entropic Optimal Transport, EOT)을 사용하여 상위-k 연산자에 대한 미분 가능한 근사인 SOFT(Scalable Optimal transport-based diFferenTiable)를 제안한다. 이는 기계학습 모델에서 엔드 투 엔드 학습을 가능하게 한다. 상위-k 선택을 정규화된 최적 운반 문제로 재구성하고, 싱크본 알고리즘과 쿠른-쿤-터커(KKT) 조건을 활용함으로써 상위-k 연산을 통한 기울기 계산을 가능하게 하여 k-NN 분류, 빔 서치, 희소 어텐션 기반 메커니즘에서 향상된 성능을 달성한다.

ABSTRACT

The top-k operation, i.e., finding the k largest or smallest elements from a collection of scores, is an important model component, which is widely used in information retrieval, machine learning, and data mining. However, if the top-k operation is implemented in an algorithmic way, e.g., using bubble algorithm, the resulting model cannot be trained in an end-to-end way using prevalent gradient descent algorithms. This is because these implementations typically involve swapping indices, whose gradient cannot be computed. Moreover, the corresponding mapping from the input scores to the indicator vector of whether this element belongs to the top-k set is essentially discontinuous. To address the issue, we propose a smoothed approximation, namely the SOFT (Scalable Optimal transport-based diFferenTiable) top-k operator. Specifically, our SOFT top-k operator approximates the output of the top-k operation as the solution of an Entropic Optimal Transport (EOT) problem. The gradient of the SOFT operator can then be efficiently approximated based on the optimality conditions of EOT problem. We apply the proposed operator to the k-nearest neighbors and beam search algorithms, and demonstrate improved performance.

연구 동기 및 목표

• 표준 상위-k 연산의 비미분성 문제를 해결하여 신경망에서 엔드 투 엔드 학습이 불가능한 문제를 해결한다.
• 기울기 기반 최적화를 지원하는 부드럽고 연속적인 상위-k 연산자 근사치를 개발한다.
• 통합된 미분 가능한 상위-k 연산자를 통해 상위-k 기반 분류, 빔 서치, 희소 어텐션 메커니즘의 미분 가능성을 확보한다.
• 엔트로픽 정규화와 싱크본 알고리즘을 통해 근사 오차를 제어하면서도 계산 효율성을 유지한다.

제안 방법

• 상위-k 선택을 운반 계획으로 파arameter화한 최적 운반(OT) 문제의 해로 상위-k 연산을 재구성한다.
• OT 문제에 엔트로픽 정규화를 적용하여 미분 가능한 근사치를 생성함으로써 엔트로픽 최적 운반(EOT) 공식을 도출한다.
• 싱크본 알고리즘을 활용해 최적 운반 계획을 효율적으로 계산함으로써 확장 가능하고 미분 가능한 상위-k 선택을 가능하게 한다.
• EOT 문제의 카르누쉬-쿤-터커(KKT) 조건을 이용해 SOFT 상위-k 연산자의 기울기를 유도함으로써 역전파를 상위-k 레이어를 통해 수행할 수 있도록 한다.
• k-NN 분류기, 빔 서치, 희소 어텐션 레이어와 같은 최종 응용 모델에 SOFT 연산자를 통합한다.
• 입력 요소 수 n에 대해 O(n)의 계산 복잡도를 구현함으로써 확장성을 확보한다.

실험 결과

연구 질문

• RQ1 비미분 가능한 상위-k 연산이 선택 의미를 유지하면서도 미분 가능하게 만들 수 있는가?
• RQ2 엔트로픽 최적 운반은 제어 가능한 오차로 이산적인 상위-k 연산에 대해 부드럽고 연속적인 근사치를 제공할 수 있는가?
• RQ3 제안된 미분 가능한 상위-k 연산자가 k-NN 분류, 빔 서치, 어텐션 메커니즘에서 엔드 투 엔드 학습 성능을 향상시키는가?
• RQ4 역전파 과정에서 SOFT 연산자의 기울기는 어떻게 행동하며, 목적 함수 최적화와 일치하는가?

주요 결과

• SOFT 상위-k 연산자는 이전에 상위-k의 비미분성으로 인해 두 단계 학습이 필요했던 모델의 엔드 투 엔드 학습을 가능하게 한다.
• k-NN 분류에서 SOFT 상위-k 연산자는 소프트맥스 기반 및 사전 학습된 CNN 기반 기준 모델 대비 향상된 정확도를 달성한다.
• 신경 기계 번역을 위한 빔 서치에서 SOFT 상위-k 연산자는 기울기 흐름을 허용하는 미분 가능한 디코딩을 가능하게 하여 생성 품질을 향상시킨다.
• 희소 어텐션 메커니즘에서 상당한 성능 향상을 달성하여, 어텐션 기반 모델에서 미분 가능한 상위-k의 효과를 입증한다.
• 간단한 실험을 통해 SOFT 연산자의 기울기가 목표 레이블에 따라 점수를 상위-k 집합으로 향하거나 멀어지게 정확히 이동시킴을 검증하였다.