ICASSP 2022 - 2022 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP) | 978-1-6654-0540-9/22/$31.00 ©2022 IEEE | DOI: 10.1109/ICASSP43922.2022.9746439

ATTENTIONPIT: SOFT PERMUTATION INVARIANT TRAINING FOR AUDIO SOURCESEPARATION WITH ATTENTION MECHANISM
Hirokazu Kameoka, Shogo Seki, Li Li, and Chihiro Watanabe
NTT Communication Science Laboratories, Nippon Telegraph and Telephone Corporation
an input mixture signal [1–3]. The time-domain approach, on
the other hand, aims to perform separation in an end-to-end
fashion by designing and training a neural network model thattakes a mixture signal as input and outputs a set of J signals(hereafter referred to as a separation network). One problem
in training a separation network is that the order of the output signals depends on the architecture and parameters of the
network and it is usually difﬁcult to predict in advance whichsource each output signal corresponds to. Therefore, simply
using the error between the output signals and the target signals arranged in a predetermined order as the training loss
will not provide satisfactory performance to the network. To
tackle this problem, a framework called permutation invariant
training (PIT) [4] has been proposed. The idea is to ﬁrst ﬁnd
the best output-target assignment for each training sample,
and then update the network parameters based on the error
given by that assignment at each iteration. This technique, or
its utterance-level extension called uPIT [5], has been shown
to be effective for the above problem and is currently being
employed to train many types of separation networks, including the time-domain audio separation network (TasNet) [6],convolutional TasNet (ConvTasNet) [7], dual-path recurrent
neural network (DPRNN) [8], Wavesplit [9], as well as for the
successive downsampling and resampling of multi-resolution
features (SuDoRM-RF) [10].

ABSTRACTPermutation 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 networkthat takes a mixture signal as input and produces the J underlying source signals. Since the order of the output signals isarbitrary, the idea of PIT is to ﬁrst ﬁnd 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 atime complexity of O(J!), 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 ﬁnd soft output-target assignments for separation network training, and can be run in polynomial timein 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 andnetwork 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.

However, since the time complexity of the exhaustive
search for the best output-target assignment is O(J!), PIT
becomes more challenging as the number of sources increases (say, to more than ten). For example, when J = 20,
J! = 2.4 × 1018 , which makes training infeasible. Somemethods have recently been proposed to overcome this obstacle. One of them is called SinkPIT [11], which exploits thefact that permutation matrices are a special case of doubly
stochastic matrices, and that the Sinkhorn-Knopp algorithm
can converge an arbitrary positive matrix to a doubly stochastic matrix. The authors showed that the Frobenius innerproduct between the pairwise error matrix (namely, a matrix
with each element being the error between a different pair of
the output and target signals) and the doubly stochastic matrix
converged from a certain matrix expressed using the pairwise
error matrix is equivalent to the PIT loss. Using this loss leads
to a time complexity of O(J 2 ). Another example is HungarianPIT [12], which uses the Hungarian algorithm [13–15], aIndex Terms— End-to-end audio source separation, permutation invariant training (PIT), attention
1. INTRODUCTION
The task of separating and extracting the signal of each soundsource from a monaural mixture of multiple sources is called
monaural source separation. Since the emergence of powerful
methods based on deep neural networks (DNNs), the level ofsound source separation performance has been dramatically
improved. The methods proposed so far can be broadly classiﬁed into two approaches: spectral-based and time-domainbased ones. The spectral-domain approach aims to predict binary, soft, or complex-valued time-frequency masks that extract the components corresponding to the source signals from
This work was supported by JST CREST Grant JPMJCR19A3, Japan.

978-1-6654-0540-9/22/$31.00 ©2022 IEEE

706

ICASSP 2022

Authorized licensed use limited to: NTT. Downloaded on November 07,2023 at 09:25:51 UTC from IEEE Xplore. Restrictions apply.

method that can solve the assignment problem in polynomial
time, to ﬁnd the best output-target assigment instead of an
exhaustive search. The time complexity of HungarianPIT isO(J 3 ), which is still at a level sufﬁcient to handle realistic
scenarios (e.g., up to twenty sources).
Another problem that PIT poses is convergence to a bad
local optimal solution. Since PIT alternates between outputtarget assignment and network parameter update, it is inherently prone to getting stuck in bad local solutions, comparedto the optimization in ordinary regression model training.
This is in some sense analogous to the k-means algorithm, awell-known clustering method that alternates between cluster
assignment and cluster centroid update. Especially in theearly stages of training, when the separation network tends
to still be immature and produces only seemingly noisy signals, it would be better to leave the chance for every output
signal to correspond to any source, rather than assigning each
output signal to a different source exclusively. For the same
reason that the expectation-maximization algorithm, or soft
k-means algorithm, is less prone to local solutions than the
k-means algorithm, we believe that soft assignment is the key
to avoiding bad local solutions in PIT. SinkPIT and another
PIT variant called ProbPIT1 [16] are examples of methods
designed to ﬁnd soft output-target assignments in the PIT
framework, and the method proposed in this paper takes asimilar approach in that regard. In particular, motivated by
the fact that the attention mechanism [17,18] has been provento be very good at ﬁnding soft assignments between the items
in a sequence or a pair of sequences, in this paper, we propose
AttentionPIT, which uses an attention mechanism for ﬁnding
soft output-target assignment in separation network training.
As will be explained later, the complexity of AttentionPIT isO(J 2 ) or O(J 3 ) depending on the choice of regularization
loss. We will also show that while AttentionPIT works reasonably well on its own, it works even better when combinedwith HungarianPIT or SinkPIT so that AttentionPIT is run
only in the early stages of training.

Fig. 1. Illustration of AttentionPIT
The DNN-based time-domain approach tackles this problem by designing a separation network fθ (·)
Ŝ = fθ (y)with some architecture parametrized by θ, where Ŝ =[ŝ1 ; . . . ; ŝJ ] ∈ RJ×N , and by training θ so that Ŝ matches
the target S = [s1 ; . . . ; sJ ] ∈ RJ×N . Note that here we haveused ; to denote vertical concatenation of vectors or matrices.
However, which source each output signal ŝj corresponds todepends on the architecture and parameters of fθ , and cannot
be easily predicted prior to training. The idea of PIT [4] is toﬁnd the best output-target assignment at each iteration before
updating θ. We can write this in the form of a training loss as[]
L(θ) = ES,y min D(Pfθ (y), S) ,(3)
P∈P

where ES,y [·] means the sample mean over all training examples, and P ∈ P denotes a permutation matrix that corresponds to a particular output-target assignment. D represents
a measure of the errors between the output and target signals.
One widely used measure is the mean of the negative scaleinvariant signal-to-distortion ratios (SI-SDRs) [19].
HungarianPIT [12] provides a way to search for the best
P in a O(J 3 ) running time under ﬁxed θ. Hence, local optimality can be efﬁciently achieved by alternately updating P
and θ. However, for the reason mentioned earlier, it may notnecessarily be the best strategy for ﬁnding the jointly optimal
solution of P and θ.

2. METHOD
2.1. Notation and Problem Formulation
By using sj (n) to denote the signal of source j, the mixturey(n) of J source signals, s1 (n), . . . , sJ (n), is given by
y(n) =∑sj (n) (n = 1, . . . , N ),(2)

2.2. AttentionPIT(1)Here, we describe AttentionPIT, which provides yet anotherway to ﬁnd soft output-target assignments during model training, different from SinkPIT [11] and ProbPIT [16]. In AttentionPIT, an additional encoder network gϕ is introduced andtrained along with the separation network. The encoder is
conﬁgured to take Ŝ or S as input and to output

jwhere n and N denote the time index and the number ofsample points, respectively. The goal of monaural source
separation is to restore sj = [sj (1), . . . , sj (N )] ∈ R1×N
(j = 1, . . . , J ) solely from y = [y(1), . . . , y(N )] ∈ R1×N .
1 The complexity of ProbPIT is the same as PIT, namely, O(J!).

K = gϕ (Ŝ),

707
Authorized licensed use limited to: NTT. Downloaded on November 07,2023 at 09:25:51 UTC from IEEE Xplore. Restrictions apply.

(4)

′

O(J 3 ). Note that as an alternative, we also tried a sparsityinducing loss [20][ ∑ (∑ |a ′ |)/√∑ |a ′ |2 −1 ]
jj
jj
j′
j′
√
, (11)
R(θ, ϕ) = λES,y J1 j
J−1

where softmax denotes a softmax function applied to eachcolumn in a matrix so that the elements in each column sum
to 1. Here, it is important to note that the way our attention
matrix is computed differs from the regular attention mechanism employed in sequence-to-sequence models in that A
represents the correspondence between channels rather thanbetween time points: In the context of sequence-to-sequencemodels, the scaled√ dot-product attention is given as A =softmax(KT Q/ J), when S and Ŝ are interpreted as Jdimensional vector sequences. In our A, element ajj ′ represents how likely the target signal j ′ is to correspond to the
ouput signal j, which takes a value close to 1 if the output jand the target j ′ are similar, and 0 otherwise, andwhich takes a value of zero if each row of A contains only one
non-zero element and requires a time complexity of O(J 2 ) tocompute, but there was no signiﬁcant difference in the actual
training time up to J = 20 compared to (10).
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. Here, it is interesting tocompare how SinkPIT and AttentionPIT differ in the way theyﬁnd soft output-target assignments during training. The difference between the two is the point in the loss computation
process at which the weighted mean is taken: SinkPIT uses
as the loss the weighted mean of the SI-SDRs of all possible output-target pairs, whereas AttentionPIT uses as the loss
the SI-SDR between each target signal and a weighted mean
(convex combination) of the output signals. Of course, thisdifference alone is not conclusive enough to argue which issuperior, so they need to be compared experimentally.
While the time complexities of SinkPIT, HungarianPIT,
and AttentionPIT in terms of J are as decribed above, the actual training time is inﬂuenced not only by J but also by otherfactors such as the hyperparameter settings and input signal
lengths, unless we assume a considerably large J. Therefore,
the superiority of each method should be discussed after thesource separation accuracy and actual training time have been
experimentally compared.

Q = gϕ (S).

(5)

Each row of K, Q ∈ RJ×N is expected to be a sequence
consisting of the features of the corresponding input signal,
where N ′ denotes the length of the output sequences. Using
K and Q, we can compute an “attention” matrix A with eachelement ajj ′ representing the similarity between the output
signal j and the target signal j ′ . If we use the expression of
the scaled dot-product attention [18], A is given by√
A = softmax(KQT / N ′ ),(6)S̃ = AT Ŝ

(7)becomes a matrix such that each row is a convex combinationof the output signals (the rows of Ŝ) that approximates one of
the target signals. Our goal is to train θ and ϕ so that S̃ ﬁts S.
If we put (6) and (7) together and denote them as
S̃ = Attn(Q, K, Ŝ),(8)

2.3. Architecture
The encoder network gϕ was designed to consist of four 1D
strided convolution layers with the kernel size of 8, stride sizeof 2, and J output channels, each followed by an instancenormalization layer and a sigmoid linear unit (SiLU) function
except for the last layer. Hence, the length of the output sequence becomes 1/16th of the input signal (i.e., N ′ = N/16).
For the separation network fθ , we chose to use the ConvTasNet architecture [7].

the training objective to be minimized becomes[]L(θ, ϕ) = ES,y D(S̃, S)(9)[]= ES,y D(Attn(gϕ (S), gϕ (fθ (y)), fθ (y)), S) .
The time complexity for computing this loss is O(J 2 ).
There should be a one-to-one correspondence between the
output and target signals. This corresponds to the case where
A satisﬁes the conditions for a permutation matrix. Sinceusing the above loss alone does not necessarily converge A
to a permutation matrix, we need a regularization term thatencourages A to become one. To this end, we introduce[]
R(θ, ϕ) = λES,y J12 ∥AAT − I∥1 ,(10)

3. EXPERIMENT
3.1. Datasets
For the experiment, we used the speech signals in the LibriSpeech corpus [21] and the script2 made available by theauthors of [12, 22] to create training, validation, and evaluation datasets for the clean mixtures (without noise) of 5, 10,15, and 20 speakers. All the signals were sampled at 8 KHz,
and all the mixtures were created in the min mode [23]. These
datasets are referred to as LibriJMix.

as the regularization loss, where ∥ · ∥1 denotes the sum of
the absolute values of all the elements in a matrix and λ denotes the regularization weight. In the following experiment,
we gradually increased λ from 0 to 50 according to λ =min(1.05e − 1, 50) during training, where e denotes the number of epochs. The time complexity for computing this loss is

2 https://github.com/ShakedDovrat/LibriMix

708
Authorized licensed use limited to: NTT. Downloaded on November 07,2023 at 09:25:51 UTC from IEEE Xplore. Restrictions apply.

3.2. Baselines and Experimental Setup

Table 1. Average number of iterations per second.

SinkPIT, HungarianPIT, and AttentionPIT (hereafter abbreviated as Sink, Hun, and Attn, respectively) were compared interms of computation time and separation accuracy. Note that
if the computation time is unlimited, Hun is equilvalent to PIT
in terms of separation accuracy. In addition to these, versions
in which Attn was run up to 20 epochs and Sink or Hun was
run thereafter (abbreviated as Attn→Sink and Attn→Hun, respectively) were also considered for comparison. For the implementations of Sink and Hun, the scripts3 provided in the
Asteroid platform [23] were used as they are. Attn was implemented on the same platform to make the condition as consistent as possible. All the methods were run using the Adam
optimizer [24] with the batch size of 8 and learning rate of
0.001. For all the methods, the learning rate was reduced bya factor of 0.5 if the loss was not improved for ﬁve consecutive epochs. The β scheduler for SinkPIT was conﬁgured in
the same manner as in [11]. The negative mean SI-SDR wasused as D. The mean SI-SDR under the best output-target
assignment (found using the Hungarian algorithm [13–15])
was used for the validation and evaluation metrics. For eachmethod, the model that obtained the best validation metricwithin 200 epochs was used for evaluation. The average number of iterations per second was used as the computational efﬁciency metric. All the methods were run on a single Tesla
V100 SXM2 GPU with a 32.0 GB memory and an Intel(R)Xeon(R) Platinum 8160 24-core CPU @ 2.10GHz.
In a previous work [12], Hun and Sink were only compared under different separation networks. In this experiment,
the separation network was ﬁxed to ConvTasNet in order tomake a pure comparison of the training methods.

DatasetLibri5Mix
Libri10Mix
Libri15Mix
Libri20Mix

Sink
2.91
2.77
2.62
2.45

Hun
3.06
2.92
2.77
2.37

Attn
2.95
2.77
2.62
2.66

Table 2. SI-SDR improvement (dB).
DatasetLibri5Mix
Libri10Mix
Libri15Mix
Libri20Mix

Sink
8.16
4.97
4.29
3.91

Hun
8.24
4.91
4.06
3.66

Attn
8.22
5.20
4.07
3.51

Att→Sin Att→Hun
8.30
8.43
5.29
5.17
4.36
4.22
3.92
3.89

The relatively lower performance of Hun compared to Sink
as J increases may be due to the fact that Hun is more prone
to bad local optima, as mentioned earlier. In Attn, there weremany cases where the attention matrices did not converge topermutation matrices as J became larger, probably due to an
insufﬁcient regularization effect. This is thought to be thecause of the relative performance degradation as J becomeslarger. Of particular note is that Attn→Sink and Attn→Hun
performed better than Sink and Hun, respectively, which implies the effect of Attn in avoiding bad local optima. It may
be possible to further improve the performance of Attn alone
by improving the λ scheduler or the loss deﬁnition for regularization, but given the above results, it seems reasonable touse Attn in combination with Sink or Hun. The above results
were obtained for the case where the separation network was
ConvTasNet, but a similar trend was observed for the case of
DPRNN [8].

3.3. Results
Table 1 shows the average number of iterations run per second
for each method. The results show that Sink and Attn wereequally efﬁcient and Hun was slightly more efﬁcient than theother two when J ≤ 15, but none of them had any practical problems in terms of training efﬁciency, at least up to
J = 20. Note that Attn was slightly faster when J = 20than when J = 15 because the length of each mixture signal
was shorter on average due to the min mode, and probably
because the computational efﬁciency of Attn depends slightly
more on the length of each input signal than the other two:In the min mode, each mixture signal stops with the shortest source signal [22]. Table 2 shows the SI-SDR improvement obtained with each method in each condition. Interestingly, the method that showed the best performance differeddepending on J: Hun tended to show relatively higher performance with smaller J, while Sink tended to show relatively
higher performance with larger J. On the other hand, Attn
showed the best performance when J = 10, but as with Hun,the performance tended to become lower as J became larger.

4. CONCLUSION
In this paper, we proposed AttentionPIT, a new PIT variant
for end-to-end audio source separation network training thatcan be run in polynomial time with respect to the number ofsources. The attention mechanism in AttentionPIT allows fora seamless process of updating the network parameters whileﬁnding soft output-target assignments through backpropagation. While the original PIT is inherently prone to getting
stuck in bad local optimal network parameters due to its hard
output-target assignment process, AttentionPIT is expected tobe effective in avoiding them. In experiments on the LibriMix
corpus, we found that while AttentionPIT works reasonably
well on its own, it works even better when combined with
SinkPIT or HungarianPIT so that AttentionPIT is run only inthe early stages of training. In the current experiment, we provided only the results of AttentionPIT applied to ConvTasNet.
Although we preliminarily tested its behavior when applied to
DPRNN as well, we plan to more thoroughly verify its effectiveness on models other than ConvTasNet in the future.

3 https://github.com/asteroid-team/asteroid

709
Authorized licensed use limited to: NTT. Downloaded on November 07,2023 at 09:25:51 UTC from IEEE Xplore. Restrictions apply.

[13] H. W. Kuhn, “The Hungarian method for the assignment
problem,” Naval research logistics quarterly, vol. 2,
no. 1-2, pp. 83–97, 1955.

5. REFERENCES[1] J. R. Hershey, Z. Chen, J. Le Roux, and S. Watanabe,“Deep clustering: Discriminative embeddings for segmentation and separation,” in Proc. ICASSP, pp. 31–35,
2016.

[14] J. Munkres, “Algorithms for the assignment and transportation problems,” Journal of the society for industrial
and applied mathematics, vol. 5, no. 1, pp. 32–38, 1957.

[2] Z. Chen, Y. Luo, and N. Mesgarani, “Deep attractor
network for single-microphone speaker separation,” inProc. ICASSP, pp. 246–250, 2017.

[15] N. Tomizawa, “On some techniques useful for solutionof transportation network problems,” Networks, vol. 1,
no. 2, pp. 173–194, 1971.

[3] Z.-Q. Wang, J. Le Roux, D. L. Wang, and J. R. Hershey,“End-to-end speech separation with unfolded iterative
phase reconstruction,” in Proc. Interspeech, 2018.

[16] M. Youseﬁ, S. Khorram, and J. H. Hansen, “Probabilistic permutation invariant training for speech separation,”in Proc. Interspeech, pp. 4604–4608, 2019.

[4] D. Yu, M. Kolbæk, Z.-H. Tan, and J. Jensen, “Permutation invariant training of deep models for speakerindependent multi-talker speech separation,” in Proc.
ICASSP, pp. 241–245, 2017.

[17] M.-T. Luong, H. Pham, and C. D. Manning, “Effective approaches to attention-based neural machine translation,” in Proc. EMNLP, 2015.
[18] A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit,
L. Jones, A. N. Gomez, L. Kaiser, and I. Polosukhin,“Attention is all you need,” in Adv. NeurIPS, 2017.

[5] M. Kolbæk, D. Yu, Z.-H. Tan, and J. Jensen, “Multitalker speech separation with utterance-level permutation invariant training of deep recurrent neural networks,” IEEE/ACM Trans. ASLP, vol. 25, no. 10,
pp. 1901–1913, 2017.

[19] J. Le Roux, S. Wisdom, H. Erdogan, and J. R. Hershey,“SDR –half-baked or well done?,” in Proc. ICASSP,
pp. 626–630, 2019.

[6] Y. Luo and N. Mesgarani, “TasNet: Time-domain audio
separation network for real-time, single-channel speech
separation,” in Proc. ICASSP, pp. 696–700, 2018.

[20] P. O. Hoyer, “Non-negative matrix factorization withsparse constraints,” J. MLR, vol. 5, pp. 1457–1469,
2004.

[7] Y. Luo and N. Mesgarani, “Conv-TasNet: Surpassing ideal time-frequency magnitude masking for speech
separation,” IEEE/ACM Trans. ASLP, vol. 27, no. 8,
pp. 1256–1266, 2019.

[21] V. Panayotov, G. Chen, D. Povey, and S. Khudanpur,“Librispeech: An ASR corpus based on public domain
audio books,” in Proc. ICASSP, 2015.
[22] J. Cosentino, M. Pariente, S. Cornell, A. Deleforge,and E. Vincent, “LibriMix: An open-source dataset
for generalizable speech separation,” arXiv:2005.11262[eess.AS], 2020.

[8] Y. Luo, Z. Chen, and T. Yoshioka, “Dual-Path RNN: Efﬁcient long sequence modeling for time-domain singlechannel speech separation,” in Proc. ICASSP, pp. 46–
50, 2020.

[23] M. Pariente, S. Cornell, J. Cosentino, S. Sivasankaran,
E. Tzinis, J. Heitkaemper, M. Olvera, F.-R. Stoter,
M. Hu, J. M. Martin-Donas, D. Ditter, A. Frank,
A. Deleforge, and E. Vincent, “Asteroid: The PyTorchbased audio source separation toolkit for researchers,”
in Proc. Interspeech, 2020.

[9] N. Zeghidour and D. Grangier, “Wavesplit: Endto-end speech separation by speaker clustering,”
arXiv:2002.08933 [eess.AS], 2020.
[10] E. Tzinis, Z. Wang, and P. Smaragdis, “Sudo rm-rf: Efﬁcient networks for universal audio source separation,”in Proc. MLSP, pp. 1–6, 2020.

[24] D. Kingma and J. Ba, “Adam: A method for stochastic
optimization,” in Proc. ICLR, 2015.

[11] H. Tachibana, “Towards listening to 10 people simultaneously: An efﬁcient permutation invariant training ofaudio source separation using sinkhorn’s algorithm,” inProc. ICASSP, pp. 491–495, 2021.
[12] S. Dovrat, E. Nachmani, and L. Wolf, “Many-speakers
single channel speech separation with optimal permutation training,” arXiv:2104.08955 [cs.SD], 2021.

710
Authorized licensed use limited to: NTT. Downloaded on November 07,2023 at 09:25:51 UTC from IEEE Xplore. Restrictions apply.