Attentionpit: Soft Permutation Invariant Training for Audio Source Separation with Attention Mechanism

Permutation invariant training (PIT) has recently attracted attention as a framework to achieve end-to-end time-domain audio source separation. Its goal is to train a separation network that takes a mixture signal as input and produces the J underlying source signals. Since the order of the output signals is arbitrary, the idea of PIT is to first find the best output-target assignment and then update the network parameters based on the error given by that assignment at each iteration. However, there are two known problems with PIT: One is that it has a time complexity of $\mathcal{O}\left( {J!} \right)$, which makes it infeasible as J increases, and the other is that it is prone to getting stuck in bad local optimal solutions due to the hard output-target assignment process. To overcome these problems simultaneously, in this paper, we propose AttentionPIT, which uses an attention mechanism to find soft output-target assignments for separation network training, and can be run in polynomial time in J, as with the recently proposed fast PIT variants such as SinkPIT and HungarianPIT. The training loss of AttentionPIT is fully differentiable, allowing us to simultaneously perform processes corresponding to soft output-target assignment and network parameter update through backpropagation. Experiments on the LibriMix corpus revealed that while AttentionPIT works reasonably well on its own, it works even better when combined with SinkPIT and HungarianPIT so that AttentionPIT is run only in the early stages of training.