TranSort: transformer for differentiable sorting

Abstract

Differentiable Sorting Algorithm is used in end-to-end differentiable frameworks, enabling gradient-based optimization of models that involve sorting operations. The Differentiable Sorting Network, the most recent state-of-the-art Differentiable Sorting Algorithm, necessitates an equal gap between input scalars for accurate sorting. We consider the sorting operation as a seq2seq generation task, where the input sequence consists of unsorted scalars, and the output sequence represents the argsort result of the unsorted scalars. From that perspective, we present TranSort, a transformer architecture proposed as an alternative to Differentiable Sorting Algorithm. TranSort demonstrates stable sorting performance on various distribution of input scalars, distinguishing itself from Differentiable Sorting Network. Moreover, we present empirical evidence highlighting the enhanced performance of end-to-end learning tasks when utilizing TranSort compared to previous Differentiable Sorting Algorithms.