Title: Annealed Multiple Choice Learning: Overcoming limitations of Winner-takes-all with annealing
Venue: NeurIPS 2024
Accepted: yes
Keywords: multiple choice learning

winner-takes-all

deterministic annealing

uncertainty quantification
Abstract: We introduce Annealed Multiple Choice Learning (aMCL) which combines simulated annealing with MCL. MCL is a learning framework handling ambiguous tasks by predicting a small set of plausible hypotheses. These hypotheses are trained using the Winner-takes-all (WTA) scheme, which promotes the diversity of the predictions. However, this scheme may converge toward an arbitrarily suboptimal local minimum, due to the greedy nature of WTA. We overcome this limitation using annealing, which enhances the exploration of the hypothesis space during training. We leverage insights from statistical physics and information theory to provide a detailed description of the model training trajectory. Additionally, we validate our algorithm by extensive experiments on synthetic datasets, on the standard UCI benchmark, and on speech separation.

Annealed Multiple Choice Learning: Overcoming limitations of Winner-takes-all with annealing

Abstract

We introduce Annealed Multiple Choice Learning (aMCL) which combines simulated annealing with MCL. MCL is a learning framework handling ambiguous tasks by predicting a small set of plausible hypotheses. These hypotheses are trained using the Winner-takes-all (WTA) scheme, which promotes the diversity of the predictions. However, this scheme may converge toward an arbitrarily suboptimal local minimum, due to the greedy nature of WTA. We overcome this limitation using annealing, which enhances the exploration of the hypothesis space during training. We leverage insights from statistical physics and information theory to provide a detailed description of the model training trajectory. Additionally, we validate our algorithm by extensive experiments on synthetic datasets, on the standard UCI benchmark, and on speech separation.

1 Introduction

Ambiguous prediction tasks arise in deep learning when the target y is ill-defined from the input x . Predicting y directly from x can be challenging due to the partial predictability of y from the information contained in x . Multiple Choice Learning (MCL) [25, 41] addresses these challenges by providing a small set of plausible hypotheses , each representing a different possible outcome given the input. MCL learns these hypotheses using a competitive training scheme that promotes the specialization of the hypotheses in distinct regions of the output space Y . The framework iteratively partitions Y into a Voronoi tesselation and guides each hypothesis toward the barycenter of its respective Voronoi cell [63]. This mechanism makes MCL akin to a gradient-descent-based and conditional variant of the popular K-means algorithm [45]. Like K-means, MCL is sensitive to initialization, subject to hypothesis collapse [49], and more generally may converge toward arbitrarily suboptimal hypothesis configurations [63]. While there is a substantial body of literature addressing the limitations of K-means [3, 20, 58], relatively little has been done to address these challenges in the context of MCL [63, 49, 54]. The core issue of MCL lies with its greedy gradient-based update of the hypotheses. This greediness precludes the exploration of the hypothesis space, preventing MCL from optimally capturing the ambiguity of y . We propose to incorporate annealing into this gradient descent update in order to improve the robustness of MCL.

Simulated annealing, inspired by the gradual cooling of metals, was originally introduced for statistical mechanics applications [28, 52] and was later applied to combinatorial problems [35]. It is

a random exploration process concurrent to the popular stochastic gradient descent, with a significant difference: gradient descent always tries to improve performance while annealing also accepts to temporarily degrade it for the sake of exploration. The range of this exploration is controlled by a temperature parameter: with a high temperature, annealing explores wide regions of the search space; when the temperature decreases, the exploration becomes narrow, and the system is able to refine its performance. This strategy has been shown to converge to an optimal state, provided that the cooling is sufficiently slow [26].

Deterministic annealing [62] is a variant of simulated annealing. In simulated annealing, exploration relies on a sequence of random moves across the search space, whereas deterministic annealing seeks greater efficiency by replacing this random process with the exact minimization of a deterministic functional, namely the free energy of the system. It has been shown that deterministic annealing can be efficiently applied to clustering [61, 62]. In this article, we show that it can be further adapted to the conditional and gradient-based setting of MCL. The resulting algorithm, which we name aMCL for annealed MCL , addresses the main issues of MCL and achieves strong performance in practical settings while being straightforward to implement and amenable to analysis. Specifically, we make the following contributions.

We introduce Annealed Multiple Choice Learning (aMCL) , a novel algorithm that incorporates annealing into the multiple choice learning framework (Section 3).

We propose a theoretical analysis of aMCL , to understand its advantages in comparison to vanilla MCL. We characterize the training trajectory of the model, by establishing an analogy with statistical physics and information theory (Section 4).

We provide extensive experimental validation, by applying this method i) to illustrative synthetic examples; ii) to a standard distribution estimation benchmark (UCI datasets); and also iii) to the challenging audio task of speech separation (Section 5). The accompanying code is made available. 1

2 Related Work

Multiple choice learning. MCL has been successfully applied to various machine learning tasks, typically using multi-head neural networks, with each head providing a prediction [40, 22, 42]. Several works observed the phenomenon of hypothesis collapse [8, 19, 33, 40, 67, 63], where some hypotheses are left unused during training. Various solutions have been proposed to tackle collapse [63, 49, 54]. Notably, [63] introduces Relaxed-WTA, which updates non-winning hypotheses with a gradient scaled by a small constant ε . However, this small gradient biases the hypotheses toward the global barycenter of the target distribution, which can be shown to be suboptimal [14].

Simulated and deterministic annealing. Deterministic annealing is a variant of simulated annealing [28, 52, 35]. Rose et al. extensively investigated its properties, particularly in relation to statistical physics and clustering [61, 62, 59, 60]. We are, to the best of our knowledge, the first to combine this technique with Winner-takes-all training in a conditional setting.

Information theory and quantization. Quantization [64] consists of discretizing continuous variables over a finite set of symbols. The rate-distortion theory studies the minimal number of bits necessary to encode information at a given level of quantization error [5, 2, 4]. Recently, a relation has been established between optimal quantization of conditional distributions [14] and multiple choice learning [63, 42, 43]. In this paper, we propose to integrate annealing for conditional quantization.

3 Annealed Multiple Choice Learning

3.1 Winner-takes-all loss and its limitations

Let X and Y denote subsets of Euclidean vector spaces. We are interested in so-called ambiguous tasks , i.e. , for any given input x ∈ X , there may be several plausible outputs y ∈ Y . Formally, let p ( x, y ) denote a joint distribution on X ×Y . Multiple Choice Learning (MCL) [25, 41] was proposed to train neural networks in this setting, and has proven its effectiveness in a wide range of machine vision [40], natural language [22] and signal processing tasks [42].

MCL consists in training several predictors ( f 1 , . . . , f n ) ∈ F ( X , Y n ) , typically a multi-head neural network derived from a common backbone, such that for each input x ∈ X , the predictions ( f 1 ( x ) , . . . , f n ( x )) provide an efficient quantization of the conditional distribution p ( y | x ) [63, 42, 43]. This goal is achieved by minimizing the distortion

where ℓ : Y 2 → R is an underlying loss function, for instance, the squared Euclidean distance ℓ (ˆ y, y ) = ∥ y -ˆ y ∥ 2 . Eq. (1) can be seen as a generalization of the conditional distortion [56].

More specifically, MCL training is an iterative procedure optimizing (1) by alternating the two following steps.

Assign each y to the closest hypothesis f k ( x ) to build the Voronoi cells:

Minimize the distortion within each cell by taking a gradient step on the WTA loss:

The prediction models can be paired with scoring models ( γ 1 , . . . , γ n ) ∈ F ( X , [0 , 1] n ) , which are trained to estimate the Voronoi regions' probability mass using the scoring loss [42]

where BCE( p, q ) ≜ -p log( q ) -(1 -p ) log(1 -q ) . In practice, the two losses (3) and (4) are optimized in a compound objective L = L WTA + L scoring . A probabilistic interpretation of such trained predictors has been developed [63, 42]. It shows that the predictions and scores can be interpreted as a mixture model approximating the conditional density p ( y | x ) by ∑ n k =1 γ k ( x ) δ f k ( x ) ( y ) .

It has been shown that WTA is sensitive to initialization [54], and often leads to suboptimal hypothesis positions, similarly to K-means. Indeed, WTA is a greedy procedure that updates only the best hypotheses: if one hypothesis falls outside the support of the data density p ( · | x ) , it may be isolated from its competitors at initialization, and remain so across training (the collapse issue).

Our method improves the WTA training scheme by addressing the inherent greediness of gradient descent and introducing variability in the exploration of the hypothesis space through deterministic annealing. Figure 1 illustrates the limitations of the aforementioned algorithms and the comparative advantage of aMCL.

3.2 Combining deterministic annealing with Multiple Choice Learning

We introduce aMCL, which combines MCL and annealing. Let t ↦→ T ( t ) denote a temperature schedule decreasing with the training step t , and vanishing at the end of the training. Similarly to MCL, aMCL alternates between an assignation and a minimization step, as follows:

Softly assign each y to all f k ( x ) using the softmin operator (or Boltzman distribution q T ( t ) ):

Minimize the distortion within each soft cell by taking a gradient step on the aWTA loss:

where q T ( t ) is kept constant ( i.e., the stop _ gradient operator is applied).

MCL

Relaxed WTA

Annealed MCL

Therefore, at the lowest level, aMCL simply consists of replacing the min operator from (2) by softmin . aMCL introduces the temperature schedule as an additional hyperparameter. As highlighted by the literature on simulated annealing [26], it is crucial to ensure that the temperature decreases slowly enough to benefit from the advantages of annealing. In practice, we experimented with both linear and exponential schedulers (see also Section 5).

On a higher level, we can interpret the objective of aMCL as a smoothed version of the MCL objective. Smoothing with high temperature simplifies the optimization problem (6), making the loss landscape easier to navigate: we can conjecture from this analysis that aMCL will find a global minimum at high temperature, and we can expect it to stay optimal as long as the temperature decreases slowly enough [9]. We can also see aMCL as an input-dependent version of deterministic annealing [61, 62]. In this view, a high temperature encourages the exploration of the hypothesis space and mitigates the greediness of the gradient descent update (3). Moreover, following [26], we can posit that there exists an optimal temperature schedule striking a balance between exploration and optimization. Yet another interpretation is that aMCL constitutes an adaptative extension of Relaxed-MCL [63], as q T ( t ) ( f k | x, y ) depends both on the distance between the hypothesis f k ( x ) and the target y , and the training step t . These interpretations shed light on the inner workings of aMCL. However, the complete training dynamic of the algorithm appears when we analyze aMCL through the lens of information theory and statistical physics, which is the purpose of the next section.

4 Theoretical analysis

In this Section, we theoretically investigate the properties of our algorithm. Specifically, we detail its training dynamic in Section 4.1. In Section 4.2, we explore how this dynamic relates to the rate-distortion curve. This relationship allows us to study in Section 4.3 the phenomenon of phase transition , where hypotheses merge and split into sub-groups depending on the temperature. Throughout this Section, we will focus on the squared Euclidean distance ℓ (ˆ y, y ) = ∥ y -ˆ y ∥ 2 .

4.1 Soft assignation and entropy constraints

Minimizing (1) is NP-hard [1, 10, 48]. Unsurprisingly, MCL can get trapped in local minima during training. In this Section, we discuss why the aMCL training scheme in Section 3.2 is more resilient to this pitfall. The first step toward our analysis is to observe that the stop _ gradient operator used in (6) of the aMCL update effectively turns the algorithm into an alternating optimization of the soft distortion

where the variables q and f are treated as independent, a procedure similar to Expectation Maximization [12]. This observation is captured by Proposition 1, where F = F ( X , Y n ) denotes the set of functions from X to Y n , ∆ n the set of all distributions on n items conditioned by points on X × Y , H ( π ) = -∑ n k =1 π k log π k the entropy of a discrete distribution π , H T = H ( q T ) the entropy of the Boltzmann distribution at temperature T , and λ t the learning rate of the gradient descent at step t .

Proposition 1 (Entropy-constrained alternated minimization) . The assignation (5) and optimization (6) steps of aMCL correspond to an entropy-constrained block coordinate descent on the soft distortion.

This is a corollary of Proposition 2, which provides additional insights on the training dynamics of aMCL.

Proposition 2 (aMCL training dynamic) . See Proposition 5 in Appendix. The following statements are true for all T > 0 , f ∈ F , strictly positive q ∈ ∆ n , x ∈ X and y ∈ Y .

Part ( i ) states that the softmin operator is the solution of the entropy-constrained minimization of the soft distortion. Part ( ii ) states that a necessary condition for minimizing the soft distortion is that each f k is a soft barycenter of the assignation distribution for each temperature T . Part ( iii ) states that each gradient update moves f k toward this soft barycenter f ⋆ , and that the update speed depends on the probability mass γ ⋆ k of the points softly assigned to f k . Together, they describe the training dynamics of aMCL. Note that as T → 0 , q T converges to a one-hot vector, and the soft barycenter in ( ii ) becomes a hard barycenter. This is consistent with the necessary optimal condition for MCL, f ⋆ k ( x ) = E Y ∼ p ( y | x ) [ Y | Y ∈ Y k ( x )] , proved by Rupperecht et al. [63].

4.2 Rate-distortion curve

We have established in Section 4.1 that the aMCL training scheme is equivalent to an entropyconstrained alternating optimization of the soft-distortion (7), with each hypothesis f k moving toward a soft barycenter. In this Section, we describe the impact of temperature cooling on this training dynamic.

First, observe that when the temperature is high, the Boltzmann distribution q T becomes uniform. Therefore, the soft Voronoi cells merge into a single cell Y , the hypotheses f k converge toward the barycenter of Y , and they all fuse into a single hypothesis:

Remarkably, this phenomenon occurs even at finite temperatures (see Appendix, Proposition 9). As the temperature decreases, a phenomenon of bifurcation occurs [61, 59]. During this process, the hypotheses iteratively split into sub-groups, as shown in Figure 4. The virtual number of hypotheses [60] for each x at a given distortion level is captured by the conditional rate-distortion function

In Eq. (10), Y ∼ p ( y | x ) follows the target distribution, the hypothesis position ˆ Y ∼ q ( f k | x ) follows a distribution over Y with q ( f k | x ) = ∫ Y q ( f k | x, y ) p ( y | x )d x , I x ( ˆ Y ; Y ) is their mutual information, and D x ( q, f ) = ∫ Y ∑ k ℓ ( f k ( x ) , y ) q ( f k | x, y ) p ( y | x )d y is the distortion for input x .

The rate-distortion function R x ( D ⋆ ) has the following key properties.

Rate-Distortion Curve

ShannonLower Bound

Rate R

(b)

aMCL

Distortion D

10

-

2

10

-

1

10

0

Temperature

10

-

2

10

-

1

Distortion

Distortion during training

T c

0

(upper bound)

2

×

σ

2

Proposition 3 (Rate-distortion properties) . See Proposition 7 in Appendix.

For each x ∈ X , let D max x (31) denote the optimal conditional distortion when using a single hypothesis, we have the following results.

(i) For each T > 0 , minimizing the free energy

over all hypotheses positions f ∈ F ( X , Y n ) comes down to solving the optimization problem that defines (10) for each x ∈ X .

(ii) R x is a non-increasing, convex and continuous function of D ⋆ , R x ( D ⋆ ) = 0 for D ⋆ ≥ D max x , and for each x the slope can be interpreted as R ′ x ( D ⋆ ) = -1 T when it is differentiable.

(iii) For each x , R x ( D ⋆ ) is bounded below by the Shannon Lower Bound (SLB)

where Y ∼ p ( y | x ) and H ( D ⋆ ) is the entropy of a Gaussian with variance D ⋆ .

Part ( i ) establishes that the rate-distortion function is tightly linked with our problem. Provided that the hypotheses f k perfectly optimize the soft distortion (7) at any temperature level, the hypothesis configuration f will follow the optimal parametric curve ( D ⋆ , R x ( D ⋆ )) for each x ∈ X . Part ( ii ) describes the shape of the parametric curve ( D ⋆ , R x ( D ⋆ )) . The Shannon Lower Bound described in part ( iii ) is a lower bound on the virtual number of hypotheses reached by aMCL.

The rate-distortion curve effectively describes the training trajectory of our algorithm, with the deterministic annealing procedure consisting of ascending along this curve [61, 59]. Interestingly, there is a set of critical temperatures at which the hypotheses suddenly split, increasing the number of sub-groups they form. By analogy with statistical physics, the behavior at these points has been coined phase transitions [23]. An illustration of the trajectory of MCL and aMCL in the rate-distortion space is shown in Figure 2. We see that MCL evolves along a constant rate R = log 2 ( n ) (in bits) following ( a ) . In contrast, aMCL initially reaches the critical state at D ⋆ = D max x following ( b ) . After the transition, the trajectory of aMCL returns to the maximal rate following ( c ) . We expect the optimization at a lower rate to be simpler and this training trajectory to provide a better initialization for the set of hypotheses compared to the vanilla MCL.

4.3 Phase transitions

During training, as the temperature decreases, the hypotheses f undergo a sequence of phase transitions at specific critical temperatures. The right tool to exhibit these critical temperatures is the Hessian of the free energy f ↦→ F ( f, q T ) . Transitions occur when the minimum of F is no longer stable, and it can be shown [61] that this relates to the eigenvalues of the following block-diagonal covariance matrix:

where q ( y | f k , x ) denotes the posterior probability of assigning a point y to the hypothesis k , calculated using Bayes's rule [60]. At high temperatures, all hypotheses merge into the conditional barycenter of the distribution (9). In this setting, all the matrices C k,k are equal to the data covariance matrix C ( x ) ≜ C ov ( X,Y ) ∼ p ( x,y ) [ Y | X = x ] , F has a unique minimizer, and the stability of this global optimum is conditioned on the strict positivity of the Hessian of F . The first critical temperature T c 0 is defined as the first temperature for which the Hessian of F is no longer positive definite.

This temperature is connected to D max x , the vanishing point of the rate-distortion function R x ( D ⋆ ) , which also corresponds to the optimal 1-hypothesis distortion [4]. Indeed, R x ( D ⋆ ) measures the virtual number of hypotheses, so that the first splitting of the hypotheses from a single point to several groups will coincide with the moment when R x ( D ⋆ ) > 0 . We summarize these observations in the following Proposition 4, which is illustrated in Figure 3.

Proposition 4 (First critical temperature) . See Definition 1, Propositions 6 ( ii ) and 8 in Appendix. Let λ max ( · ) denote the maximum eigenvalue of a matrix, D ∗ x ( T ) ≜ inf f ∈F D x ( q T , f ) , D ∗ ( T ) ≜ inf f ∈F D ( q T , f ) , and D max ≜ ∫ x ∈X D max x p ( x )d x . Then the two following properties hold.

(i) D ∗ x and D ∗ are non-decreasing functions of T and admit generalized inverses (with the convention g -1 ( β ) = inf { α | g ( α ) ≥ β } for the generalized inverse of a function g ).

(ii) The conditional (resp. non-conditional) first critical temperature T c 0 ( x ) (resp. T c 0 ) satisfy

T(t1)

T(t2)

-1

T(t3)

-1

T(t4)

1

c=0.1

=0.6

c=0.9

-1

c=0.1

=0.6

c=0.9

After the first transition, an interesting phenomenon occurs for some distributions p ( x, y ) . Instead of splitting in all directions, the hypotheses f k continue to form a small number of subgroups. The hypotheses f k may undergo many phase transitions before they all split apart from each other: this is illustrated in Figure 4. Generally, this recursive splitting reaches an end: under mild conditions [36], there is a critical temperature below which the hypotheses all separate from each other. This defines the last critical temperature T c ∞ . Remarkably, T c ∞ is associated with the Shannon Lower Bound: the temperature at which R x ( D ⋆ ) hits the lower bound SLB( D ⋆ ) corresponds to the moment when the hypotheses f k completely separate from each other (see Theorem 1 and 3 in [59]).

5 Experimental validation

In this Section, we experimentally investigate the advantage of our algorithm in practical settings. Specifically, we evaluate it on the standard UCI benchmark in Section 5.1, and we apply it to the challenging task of speech separation in Section 5.2.

5.1 UCI datasets

5.1.1 Setup

General setup. We followed the experimental protocol described by [29] for the UCI benchmark [15]. Specifically, we used the official train-test splits, with 20 folds except for the Protein dataset, which is split into 5 folds, and the Year dataset, which uses a single fold.

Baselines. In our result tables, we also include data from three baseline methods detailed in Table 1 of Lakshminarayanan et al. 's paper [38] ('Deep Ensembles'), which serves as a reference. These baselines include Probabilistic Back Propagation [29] (denoted 'PBP'), and Monte Carlo Dropout [21] (denoted 'MC-dropout'). As additional baselines, we include the standard score-based MCL ( e.g. , [42]). We also include the Relaxed-MCL variant [63] with ε = 0 . 1 . The impact of ε is discussed in Appendix D.2. Our method (aMCL), was trained with an exponential scheduler of the form T ( t ) = T 0 ρ t , with ρ = 0 . 95 and T 0 = 0 . 5 . Comparison with a linear scheduler is also provided in Appendix D.1. Both aMCL and Relaxed-MCL were trained for 1,000 epochs. Each MCL system was trained with n = 5 hypotheses.

Metrics. We computed the following metrics on the UCI datasets. Let d denote the squared Euclidean distance: d (ˆ y i , y i ) = ∥ y i -ˆ y i ∥ 2 , and N the number of samples in each dataset. The RMSE ( ↓ ) is defined as RMSE = √ 1 N ∑ i d (ˆ y i , y i ) , where ˆ y i denotes the estimated conditional mean. The latter was computed with ∑ n k =1 γ k ( x i ) f k ( x i ) for the MCL variants. The results of this experiment are summed up in Table 1 (Distortion), and also in Table 2 (RMSE). Rows are ordered by dataset size N , with the intensity of the grey color proportional to N (excluding the Year dataset).

5.1.2 Results

±

±

±

±

±

±

±

±

±

Comparison of aMCL and MCL. We can observe that aMCL performs comparably or outperforms vanilla MCL in most settings, especially for large dataset sizes, both in terms of distortion and RMSE. This outcome supports the claims made in the paper and is especially promising, given that the temperature scheduler was not specifically optimized for each dataset.

Comparison of aMCL and standard UCI baselines. We observe that aMCL performs on par with, and in some cases exceeds, standard baselines on the RMSE metric. This is noteworthy, as aMCL is not explicitly optimized for RMSE during training-its primary focus is on quantization. While perfect quantization would naturally result in optimal RMSE, achieving low RMSE is not the main objective of aMCL. For example, the RMSE performance of MCL is slightly worse across most datasets in those experiments.

Comparison with Relaxed-MCL. Finally, we compare aMCL and Relaxed-MCL, since aMCL can be interpreted as an adaptative version of Relaxed-MCL. Our results indicate that aMCL generally outperforms Relaxed-MCL in terms of distortion across nearly all datasets. However, we also observe that Relaxed-MCL demonstrates strong performance in terms of RMSE. We attribute these findings to the bias of Relaxed-MCL toward the conditional barycenter of the target distribution, which seems to improve RMSE at the expense of increased distortion. This trade-off arises because RMSE evaluates the accuracy of the barycenter estimation, while distortion measures the quantization performance. The trade-off between distortion and RMSE can be adjusted by tuning the value of ε . Further analysis of this parameter is presented in Appendix D.1, where we show that aMCL strikes a good balance between these two metrics.

These results strongly support the use of aMCL as a quantization algorithm in practical settings. To further evaluate its effectiveness, we also apply aMCL to a more challenging task, namely speech separation.

5.2 Application to speech separation

Speech separation consists of isolating each speaker's signal from a mixture in which they are simultaneously active. This task is of major interest for automatic speech processing applications [50, 44, 71]. In these experiments, we explore the application of MCL and the proposed aMCL to the task of speech separation. An extensive description of the experiments is proposed in Appendix E.2.

5.2.1 Experimental setting

General purpose . Speech separation consists in obtaining the source signals y 1 , . . . , y m ∈ R l from a mixture x = ∑ m s =1 y s . Hence, the task is to provide estimates ˆ y 1 , . . . , ˆ y m of the isolated speech tracks from the mixture x .

Dataset . Source separation experiments are conducted on the Wall Street Journal dataset [30] (WSJ0-mix), a standard benchmark for speech separation. We focus on the 2- and 3- speaker mixture scenarios, with each scenario including 20,000 training, 5,000 validation, and 3,000 testing mixtures.

Model architecture . The source separation task is solved using the Dual-Path Recurrent Neural Network (DPRNN) architecture [46]. DPRNN has been extensively used in speech separation, as it strikes a good balance between performance and number of trainable parameters [73, 65].

Separation metrics . We use the Scale-Invariant Signal-to-Distortion Ratio (SI-SDR) to measure the separation quality [70, 39]. There is an ambiguity in finding the best assignment between predicted and active sources. The PIT SI-SDR loss [37] initially addresses this issue. We propose to use MCL and our new variant aMCL to perform this matching (See Appendix E).

5.2.2 Results

Comparing PIT, MCL, Relaxed-WTA and aMCL. Table 3 presents the source separation results in the 2- and 3-speaker scenarios. First, aMCL demonstrates performance equivalent to or better than MCL. Both methods can be used for the separation task. However, we observed that MCL is subject to hypothesis collapse for some training seeds, while aMCL is more robust to initialization. This translates into a lower interseed standard deviation for aMCL. Second, we observe the advantage of aMCL over RelaxedWTA, which is consistent with our previous analysis of the barycenter bias of this method. Third, aMCL performs equivalently to PIT in both sce-

narios. Note that by using MCL or aMCL, the number of predictions n could exceed the number of sources m . This could be leveraged to improve the separation metrics (cf. Appendix E.3.2). Finally, aMCL improves the algorithmic complexity of PIT from O ( m 3 ) to O ( mn ) (cf. Appendix E.3.1). These results make aMCL stand as a good alternative to PIT.

Observing phase transition. When the metric is the Euclidean distance, the theoretical analysis of Section 4.3 and the synthetic experiments (see Figure 3) have highlighted a phenomenon of phase transition for aMCL. Here, we analyze the validation loss trajectory as a function of the temperature for different initial temperatures, and using an exponential scheduler. The curves with the two higher initial temperatures in Figure 8 exhibit a plateau until a given temperature. After this critical point, the loss decreases. Although the SI-SDR is non-Euclidean, this behavior resembles that observed for the Euclidean metric. This is detailed in Section E.3.3 of the Appendix.

6 Conclusion

This article introduces aMCL, a novel training method that combines deterministic annealing and the Winner-Takes-all training scheme to address two key issues of MCL: hypothesis collapse and convergence toward a suboptimal local minimum of its quantization objective. We provide a detailed analysis of aMCL's training dynamics. Moreover, drawing on statistical physics and information theory, we provide insights into the trajectory of the aMCL predictions during training. In particular, we exhibit a phase transition phenomenon and establish its connection to the rate-distortion curve. We validate our analysis with experiments on synthetic data, on the UCI datasets, and on a real-world speech separation task. This demonstrates that aMCL is a theoretically grounded alternative to MCL in diverse settings. Future work includes a detailed analysis of the temperature schedule's impact, the derivation of performance bounds, a thorough examination of our algorithm's convergence, particularly at finite temperature, as well as an evaluation of its generalization capabilities on out-ofdistribution samples.

Limitations. First, aMCL introduces a temperature schedule: this is a challenging hyperparameter tightly linked to the optimizer and its learning rate. The derivation of optimal schedules is left to future work. Second, annealing requires a slow temperature schedule to maintain model performance. This potentially leads to longer training times.

Broaden impact. This paper introduces research aimed at progressing the field of Machine Learning. While our work has numerous potential societal implications, we believe there are no specific consequences that need to be emphasized in this paper.

7 Acknowledgments

This work was funded by the French Association for Technological Research (ANRT CIFRE contract 2022-1854) and Hi! PARIS through their PhD in AI funding programs, and was performed using HPC resources from GENCI-IDRIS (Grant 2021-AD011013406R1). We are grateful to the reviewers for their insightful comments.

References

Daniel Aloise, Amit Deshpande, Pierre Hansen, and Preyas Popat. Np-hardness of euclidean sum-of-squares clustering. Machine learning , 2009. 4, 15

Suguru Arimoto. An algorithm for computing the capacity of arbitrary discrete memoryless channels. IEEE Transactions on Information Theory , 1972. 2

David Arthur, Sergei Vassilvitskii, et al. k-means++: The advantages of careful seeding. In Annual ACM-SIAM Symposium on Discrete Algorithms, SODA , 2007. 1

Toby Berger. Rate-distortion theory. Wiley Encyclopedia of Telecommunications , 2003. 2, 7, 20

Richard Blahut. Computation of channel capacity and rate-distortion functions. IEEE Transactions on Information Theory , 1972. 2

Richard E Blahut. Principles and practice of information theory . Addison-Wesley Longman Publishing Co., Inc., 1987. 20

David P Bourne and Steven M Roper. Centroidal power diagrams, lloyd's algorithm, and applications to optimal location problems. SIAM Journal on Numerical Analysis , 2015. 15

Mike Brodie, Chris Tensmeyer, Wes Ackerman, and Tony Martinez. Alpha model domination in multiple choice learning. In IEEE International Conference on Machine Learning and Applications, ICMLA , 2018. 2

Imre Csiszár. On the computation of rate-distortion functions (corresp.). IEEE Transactions on Information Theory , 1974. 4

Sanjoy Dasgupta. The hardness of k-means clustering. 2008. 4, 15

L Davisson. Rate distortion theory: A mathematical basis for data compression. IEEE Transactions on Communications , 1972. 20

Arthur P Dempster, Nan M Laird, and Donald B Rubin. Maximum likelihood from incomplete data via the em algorithm. Journal of the Royal Statistical Society , 1977. 5

Qiang Du, Maria Emelianenko, and Lili Ju. Convergence of the lloyd algorithm for computing centroidal voronoi tessellations. SIAM Journal on Numerical Analysis , 2006. 15

Qiang Du, Vance Faber, and Max Gunzburger. Centroidal voronoi tessellations: Applications and algorithms. SIAM Review , 1999. 2

Dheeru Dua, Casey Graff, et al. Uci machine learning repository. 2017. 8, 24

Jack Edmonds and Richard M Karp. Theoretical improvements in algorithmic efficiency for network flow problems. Journal of the ACM, JACM , 1972. 29

Paul Embrechts and Marius Hofert. A note on generalized inverses. Mathematical Methods of Operations Research , 2013. 20

Maria Emelianenko, Lili Ju, and Alexander Rand. Nondegeneracy and weak global convergence of the lloyd algorithm in R d . SIAM Journal on Numerical Analysis , 2008. 15

Michael Firman, Neill DF Campbell, Lourdes Agapito, and Gabriel J Brostow. Diversenet: When one right answer is not enough. In IEEE/CVF Conference on Computer Vision and Pattern Recognition, CVPR , 2018. 2

Pasi Fränti and Sami Sieranoja. How much can k-means be improved by using better initialization and repeats? Pattern Recognition , 2019. 1

Yarin Gal and Zoubin Ghahramani. Dropout as a bayesian international workshop on approximation algorithms for combinatorial optimization, approximation: Representing model uncertainty in deep learning. In International Conference on Machine Learning, ICML , 2016. 8, 24

Nuno Cruz Garcia, Sarah Adel Bargal, Vitaly Ablavsky, Pietro Morerio, Vittorio Murino, and Stan Sclaroff. Distillation multiple choice learning for multimodal action recognition. In IEEE/CVF Winter Conference on Applications of Computer Vision, WACV , 2021. 2

Thore Graepel, Matthias Burger, and Klaus Obermayer. Phase transitions in stochastic selforganizing maps. Physical Review E , 1997. 6

Robert M Gray. Source coding theory . Springer Science & Business Media, 1989. 20

Abner Guzman-Rivera, Dhruv Batra, and Pushmeet Kohli. Multiple choice learning: Learning to produce multiple structured outputs. Advances in Neural Information Processing Systems, NeurIPS , 2012. 1, 2

Bruce Hajek. Cooling schedules for optimal annealing. Mathematics of Operations Research , 1988. 2, 4, 22, 23

Xizewen Han, Huangjie Zheng, and Mingyuan Zhou. Card: Classification and regression diffusion models. Advances in Neural Information Processing Systems, NeurIPS , 2022. 24

WKeith Hastings. Monte carlo sampling methods using markov chains and their applications. 1970. 1, 2, 22, 23

José Miguel Hernández-Lobato and Ryan Adams. Probabilistic backpropagation for scalable learning of bayesian neural networks. In International Conference on Machine Learning, ICML , 2015. 8, 24

John R Hershey, Zhuo Chen, Jonathan Le Roux, and Shinji Watanabe. Deep clustering: Discriminative embeddings for segmentation and separation. In IEEE International Conference on Acoustics, Speech and Signal Processing, ICASSP , 2016. 9, 27

Po-Sen Huang, Scott Deeann Chen, Paris Smaragdis, and Mark Hasegawa-Johnson. Singingvoice separation from monaural recordings using robust principal component analysis. In IEEE International Conference on Acoustics, Speech and Signal Processing, ICASSP , 2012. 26

Mikaela Iacobelli. Asymptotic quantization for probability measures on riemannian manifolds. ESAIM: Control, Optimisation and Calculus of Variations , 2016. 20

Eddy Ilg, Ozgun Cicek, Silvio Galesso, Aaron Klein, Osama Makansi, Frank Hutter, and Thomas Brox. Uncertainty estimates and multi-hypotheses networks for optical flow. In European Conference on Computer Vision, ECCV , 2018. 2

Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In International Conference on Learning Representations, ICLR , 2015. 24

Scott Kirkpatrick, C Daniel Gelatt Jr, and Mario P Vecchi. Optimization by simulated annealing. Science , 1983. 1, 2

Tobias Koch. The shannon lower bound is asymptotically tight. IEEE Transactions on Information Theory , 2016. 8

Morten Kolbæk, Dong Yu, Zheng-Hua Tan, and Jesper Jensen. Multitalker speech separation with utterance-level permutation invariant training of deep recurrent neural networks. IEEE/ACM Transactions on Audio, Speech, and Language Processing , 2017. 10

Balaji Lakshminarayanan, Alexander Pritzel, and Charles Blundell. Simple and scalable predictive uncertainty estimation using deep ensembles. Advances in Neural Information Processing Systems, NeurIPS , 2017. 8, 9, 24

Jonathan Le Roux, Scott Wisdom, Hakan Erdogan, and John R Hershey. Sdr-half-baked or well done? In IEEE International Conference on Acoustics, Speech and Signal Processing, ICASSP , 2019. 10, 26

Kimin Lee, Changho Hwang, KyoungSoo Park, and Jinwoo Shin. Confident multiple choice learning. In International Conference on Machine Learning, ICML , 2017. 2

Stefan Lee, Senthil Purushwalkam Shiva Prakash, Michael Cogswell, Viresh Ranjan, David Crandall, and Dhruv Batra. Stochastic multiple choice learning for training diverse deep ensembles. Advances in Neural Information Processing Systems, NeurIPS , 2016. 1, 2, 4

Victor Letzelter, Mathieu Fontaine, Mickaël Chen, Patrick Pérez, Slim Essid, and Gael Richard. Resilient multiple choice learning: A learned scoring scheme with application to audio scene analysis. Advances in Neural Information Processing Systems, NeurIPS , 2023. 2, 3, 4, 8, 15, 16, 23

Victor Letzelter, David Perera, Cédric Rommel, Mathieu Fontaine, Slim Essid, Gaël Richard, and Patrick Perez. Winner-takes-all learners are geometry-aware conditional density estimators. In International Conference on Machine Learning, ICML , 2024. 2, 3, 15, 23

Chenda Li, Jing Shi, Wangyou Zhang, Aswin Shanmugam Subramanian, Xuankai Chang, Naoyuki Kamo, Moto Hira, Tomoki Hayashi, Christoph Boeddeker, Zhuo Chen, et al. Espnetse: End-to-end speech enhancement and separation toolkit designed for asr integration. In IEEE Spoken Language Technology Workshop, SLT , 2021. 9, 26

Stuart Lloyd. Least squares quantization in pcm. IEEE Transactions on Information Theory , 1982. 1

Yi Luo, Zhuo Chen, and Takuya Yoshioka. Dual-path rnn: efficient long sequence modeling for time-domain single-channel speech separation. In IEEE International Conference on Acoustics, Speech and Signal Processing, ICASSP , 2020. 9, 27, 28

Yi Luo and Nima Mesgarani. Conv-tasnet: Surpassing ideal time-frequency magnitude masking for speech separation. IEEE/ACM Transactions on Audio, Speech, and Language Processing , 2019. 27, 28

Meena Mahajan, Prajakta Nimbhorkar, and Kasturi Varadarajan. The planar k-means problem is np-hard. Theoretical Computer Science , 2012. 4, 15

Osama Makansi, Eddy Ilg, Ozgun Cicek, and Thomas Brox. Overcoming limitations of mixture density networks: A sampling and fitting framework for multimodal future prediction. In IEEE/CVF Conference on Computer Vision and Pattern Recognition, CVPR , 2019. 1, 2, 29

Amparo Marti, Maximo Cobos, and Jose J Lopez. Automatic speech recognition in cocktailparty situations: A specific training for separated speech. The Journal of the Acoustical Society of America, JASA , 2012. 9, 26

Neri Merhav. Rate-distortion function via minimum mean square error estimation. IEEE Transactions on Information Theory , 2011. 20

Nicholas Metropolis, Arianna W Rosenbluth, Marshall N Rosenbluth, Augusta H Teller, and Edward Teller. Equation of state calculations by fast computing machines. The Journal of Chemical Physics , 1953. 1, 2, 22, 23

Fionn Murtagh and Pedro Contreras. Algorithms for hierarchical clustering: an overview. Wiley Interdisciplinary Reviews: Data Mining and Knowledge Discovery , 2012. 22

Sriram Narayanan, Ramin Moslemi, Francesco Pittaluga, Buyu Liu, and Manmohan Chandraker. Divide-and-conquer for lane-aware diverse trajectory prediction. In IEEE/CVF Conference on Computer Vision and Pattern Recognition, CVPR , 2021. 1, 2, 3

Elias Nehme, Rotem Mulayoff, and Tomer Michaeli. Hierarchical uncertainty exploration via feedforward posterior trees. Advances in Neural Information Processing Systems, NeurIPS , 2024. 22, 23

Gilles Pagès and Jacques Printems. Optimal quadratic quantization for numerics: the gaussian case. Monte Carlo Methods and Applications , 2003. 3, 15

Gilles Pages and Jun Yu. Pointwise convergence of the lloyd i algorithm in higher dimension. SIAM Journal on Control and Optimization , 2016. 15

José M Pena, Jose Antonio Lozano, and Pedro Larranaga. An empirical comparison of four initialization methods for the k-means algorithm. Pattern Recognition Letters , 1999. 1

Kenneth Rose. A mapping approach to rate-distortion computation and analysis. IEEE Transactions on Information Theory , 1994. 2, 5, 6, 8, 20, 22

Kenneth Rose. Deterministic annealing for clustering, compression, classification, regression, and related optimization problems. Proceedings of the IEEE , 1998. 2, 5, 7, 21

Kenneth Rose, Eitan Gurewitz, and Geoffrey C Fox. Statistical mechanics and phase transitions in clustering. Physical Review Letters , 1990. 2, 4, 5, 6, 7, 15, 21

Kenneth Rose, Eitan Gurewitz, and Geoffrey C Fox. Vector quantization by deterministic annealing. IEEE Transactions on Information Theory , 1992. 2, 4, 15

Christian Rupprecht, Iro Laina, Robert DiPietro, Maximilian Baust, Federico Tombari, Nassir Navab, and Gregory D Hager. Learning in an uncertain world: Representing ambiguity through multiple hypotheses. In International Conference on Computer Vision, ICCV , 2017. 1, 2, 3, 4, 5, 8, 15, 23, 29

Claude E Shannon et al. Coding theorems for a discrete source with a fidelity criterion. International Convention Record, IRE , 1959. 2, 20

Cem Subakan, Mirco Ravanelli, Samuele Cornell, Mirko Bronzi, and Jianyuan Zhong. Attention is all you need in speech separation. In IEEE International Conference on Acoustics, Speech and Signal Processing, ICASSP , 2021. 9

Hideyuki Tachibana. Towards listening to 10 people simultaneously: An efficient permutation invariant training of audio source separation using sinkhorn's algorithm. In IEEE International Conference on Acoustics, Speech and Signal Processing, ICASSP , 2021. 29

Kai Tian, Yi Xu, Shuigeng Zhou, and Jihong Guan. Versatile multiple choice learning and its application to vision computing. In IEEE/CVF Conference on Computer Vision and Pattern Recognition, CVPR , 2019. 2

Efthymios Tzinis, Zhepei Wang, and Paris Smaragdis. Sudo rm-rf: Efficient networks for universal audio source separation. In International Workshop on Machine Learning for Signal Processing, MLSP , 2020. 28

Andrea Vattani. K-means requires exponentially many iterations even in the plane. In Annual symposium on Computational geometry, SoCG , 2009. 15

Emmanuel Vincent, Rémi Gribonval, and Cédric Févotte. Performance measurement in blind audio source separation. IEEE/ACM Transactions on Audio, Speech, and Language Processing , 2006. 10, 26

Thilo Von Neumann, Keisuke Kinoshita, Marc Delcroix, Shoko Araki, Tomohiro Nakatani, and Reinhold Haeb-Umbach. All-neural online source separation, counting, and diarization for meeting analysis. In IEEE International Conference on Acoustics, Speech and Signal Processing, ICASSP , 2019. 9, 26

DeLiang Wang and Jitong Chen. Supervised speech separation based on deep learning: An overview. IEEE/ACM Transactions on Audio, Speech, and Language Processing , 2018. 26

Zhong-Qiu Wang, Samuele Cornell, Shukjae Choi, Younglo Lee, Byeong-Yeol Kim, and Shinji Watanabe. Tf-gridnet: Making time-frequency domain models great again for monaural speaker separation. In IEEE International Conference on Acoustics, Speech and Signal Processing, ICASSP , 2023. 9

Kenneth Zeger, Jacques Vaisey, Allen Gersho, et al. Globally optimal vector quantizer design by stochastic relaxation. IEEE Transactions on Signal Processing , 1992. 23

Organisation of the Appendix

The Appendix is organized as follows. Appendix A presents the theoretical analysis of our algorithm. It begins with the introduction of the notations and the definition of the training scheme in Appendices A.1 and A.2 respectively. This is followed by an interpretation of the algorithm in terms of entropyconstrained alternate optimization in Appendix A.3. Appendix A.4 provides an analysis of the training dynamics of the algorithm in relation to rate-distortion theory, and Appendix A.5 discusses phase transitions. Connection with the literature and additional discussions are provided in Appendix B. Additional details and results from experiments on synthetic data and UCI datasets are provided in Appendices C and D. Finally, Appendix E offers an extensive description of the Source Separation experiment, including the impact of the number of hypotheses in Appendix E.3.2, and the analysis of phase transitions for this task in Appendix E.3.3.

A Theoretical analysis

A.1 Notations and motivations

Following the main paper notations, let X ⊂ R d 1 and Y ⊂ R d 2 denote the input and target spaces respectively. We will assume that X and Y are bounded. We note p ( x, y ) the continuous density of a joint data distribution on X × Y . Let f = ( f 1 , . . . , f n ) ∈ F Θ ⊂ F ( X , Y n ) denote the hypothesis models, which are families of n neural networks with parameters in Θ . Likewise, let γ = ( γ 1 , . . . , γ n ) ∈ F Θ ′ ⊂ F ( X , [0 , 1] n ) denote the score models, with parameters in Θ ′ . We will sometimes write F = F ( X , Y n ) .

Let t ∈ N denote a training step, and t epoch ∈ N ∪ {∞} denote the number of training epochs. Let T : t ↦→ T ( t ) ≥ 0 be a temperature schedule that decreases with the training step and verifies lim t → t epoch T ( t ) = 0 . Note that we consider the temperature to be constant across an epoch, with linear or exponential decrease between epochs. Let ℓ : Y 2 → R + denote an underlying loss function. We will restrict our analysis to the Euclidean squared distance ℓ (ˆ y, y ) = ∥ y -ˆ y ∥ 2 unless otherwise stated. For any vector z , we note z t as its vector transposition. Gradient descent involves a learning rate schedule, that we note λ t .

Wedenote by ∆ n the set of discrete distributions on n items conditioned by points on X×Y (typically representing a soft assignment of a target to the hypotheses). We denote by H ( q ) the entropy of a distribution q . If q ∈ ∆ n , we write H ( q ) = -∫ X×Y ∑ n k =1 q ( f k | x, y ) log q ( f k | x, y ) p ( x, y )d x d y . If Z is a random variable with density p ( z ) and Y is another random variable, we define the differential entropy of Z as H ( Z ) = -E Z [log( p ( Z ))] and its conditional entropy as H ( Z | Y ) = -E ( Z,Y ) [log( p ( Z | Y ))] . Mutual information I ( Z, Y ) between two random variables Z and Y is defined as I ( Z ; Y ) = H ( Z ) -H ( Z | Y ) .

Multiple Choice Learning (MCL) aims at training the hypotheses ( f 1 , . . . , f n ) such that, for each input x ∈ X , the hypothesis predictions ( f 1 ( x ) , . . . , f n ( x )) provide an efficient quantization of the conditional distribution p ( y | x ) [63, 42, 43]. The quality of the quantization can be evaluated using the distortion :

which can be seen as a generalization of the conditional distortion [56]. The two definitions coincide when p ( y | x ) is independent of x . Minimizing (16) is NP-hard [1, 10, 48]. The K-means algorithm, which is the standard approach for this task, can take an exponential number of steps before converging [69], and is only guaranteed to find a local minimum in general settings [13, 7]. The same applies to WTA: due to its greedy nature, it can be trapped, at any point of its training trajectory, into a local minimum that it will never manage to escape. One common issue is the hypothesis collapse [57, 18], a situation in which some of the hypotheses are left unused, therefore leading to a suboptimal configuration (see Figure 1).

To mitigate these issues, we introduce an annealed version of MCL. It can be seen as an inputdependent version of deterministic annealing [61, 62]. Our training scheme is described in Appendix A.2.

A.2 Training scheme of aMCL

Let ( x, y ) be a training example sampled from p at step t . Our proposed annealed MCL training process is defined through the following subsequent steps:

Perform a forward pass through f 1 ( x ) , . . . , f n ( x ) .

Compute the Boltzmann distribution q T ( t ) defined for each f k by

and detach it from the computational graph.

Perform a gradient step on the loss

Note that we used the same scoring loss L scoring (4) as in rMCL, the score-based method in [42].

In this paper, we extend the distortion D ( f ) to account for a distribution q on the hypotheses, and define the resulting soft distortion as

When q = q T ( t ) , the soft distortion D ( q T ( t ) , f ) corresponds to the expectation of the aWTA loss (18).

We will now introduce a decoupling assumption, which states that the family of models F Θ is expressive enough to encompass the global minimizer of the soft distortion. Noting D x ( q, f ) the integrand of (19) over the X integral, we can formalize this assumption as follows.

Assumption 1 (Decoupling) . We assume that the model family is perfectly expressive, i.e. F Θ = F .

Note that under this assumption, the global minimizer of the soft distortion D ( q, f ) , minimizes the integrand of (19) for all x ∈ X :

In the subsequent Sections, our goal is to analyze the aMCL training scheme and justify its design. Specifically, we show that annealing simplifies the optimization problem defined by (16). In Appendix A.3, we first focus on studying the properties of the soft distortion D ( q, f ) .

In the following, if x ∈ X is fixed and when the context is clear, we will omit the dependency on the input and denote the hypotheses position ( f 1 , . . . , f n ) ∈ Y n .

A.3 Soft assignation and entropy constraint

In this Section, we rewrite the assignation and optimization steps of aMCL as an entropy-constrained block optimization of the soft distortion D ( q, f ) . Noting H T = H [ q T ] , we formalize this result with the following Theorem.

Theorem 2 (Entropy constrained distortion minimization) . The assignation (17) and optimization (18) steps of aMCL correspond to a block coordinate descent on the soft distortion. For all k ∈ [ [1 , n ] ] ,

Theorem 2 is a corollary of the following Proposition, which provides additional intuition into the dynamic of aMCL.

Proposition 5 (aMCL training dynamic) . For all T > 0 , f ∈ F , strictly positive q ∈ ∆ n , x ∈ X and y ∈ Y , the following statements are true.

Part ( i ) states that the softmin operator is the solution of the entropy-constrained minimization of the soft distortion. Therefore, gradient-based optimization of D ( q, f ) along axis q is unnecessary, and we can directly use the closed-form expression in ( i ) for each temperature level T ( t ) , or equivalenty each corresponding entropy level H T ( t ) . Consequently, the optimization of D ( q, f ) reduces to the gradient-based minimization D ( q T ( t ) , f ) along axis f for each temperature level T ( t ) .

Part ( ii ) states that a necessary condition for the minimization of the soft distortion is that each f k is a soft barycenter of the assignation distribution q T ( t ) for each temperature T . Part ( iii ) states that each gradient update moves f k towards this soft barycenter f ⋆ k , and that the update speed depends on the probability mass γ ⋆ k of the points softly assigned to f k . Together, these results describe the training dynamics of aMCL. Let us prove these results.

Proof. ( i ) First observe that under the decoupling assumption, minimizing D ( q, f ) amounts to minimizing the integrand D x ( q, f ) = ∫ Y ∑ n k =1 ℓ ( f k ( x ) , y )) q ( f k | x, y ) p ( y | x )d y for each x ∈ X (see (20)). When x is fixed, the distribution q becomes a tuple of n scalars ( q 1 , . . . , q n ) , and the Lagragian of (21) writes

Using the first-order necessary optimization conditions, we find that

From this it immediately follows that λ = 1 + 1 / ∑ n k =1 exp ( -ℓ ( f k ,y ) T ) is a normalization factor, q k = exp ( -ℓ ( f k ,y ) T ) / ∑ n s =1 exp ( -ℓ ( f s ,y ) T ) is the Boltzmann distribution, and H T = H [ q T ] correspond to its entropy.

( ii ) Using the same reasoning than in ( i ) we set x ∈ X . Then each f k ∈ Y becomes a simple vector. Observing that ∇ f k D x ( q, f ) vanishes at the minimum of D x ( q, f ) , a necessary condition to minimize D ( q, f ) is

(25)

Then, as we assumed q > 0 , we have ∫ Y q ( f k | x, y ) p ( y | x )d y > 0 , and

( iii ) It follows immediately from ( ii ) that

Before further analyzing the training dynamics of aMCL, let us pause to examine a few properties of the soft distortion D ( q, f ) . First, observe that the hard distortion D ( f ) given by (16) is a particular case of soft distortion D ( q, f ) where q ( f k | x, y ) = 1 [ k ∈ argmin s ∈ [ [1 ,n ] ] ℓ ( f s ( x ) , y )] . Second, inf f ∈F D x ( q T , f ) and inf f ∈F D ( q T , f ) are non-decreasing functions of T . Third, entropyconstrained soft distortion minimization (21) is equivalent to the minimization of the free energy

a quantity defined in statistical physics, that will be analyzed in more detail in the following Section. Moreover, optimal free energy has a closed-form formula in our case. We group these observations in the following Proposition.

Proposition 6 (Additional properties of the distortion and the free energy) . For all T > 0 , f ∈ F , q ∈ ∆ n , x ∈ X and y ∈ Y , the following statements are true.

( i ) q ( f | x, y ) = ( 1 [ k ∈ argmin s ∈ [ [1 ,n ] ] ℓ ( f s ( x ) , y )] ) k ∈ [ [1 ,n ] ] ⇒∀ q ′ ∈ ∆ n , D ( q, f ) ≤ D ( q ′ , f ) .

Proof. ( i ) Notice that argmin s ∈ [ [1 ,n ] ] ℓ ( f s ( x ) , y ) selects the hypothesis f s with lowest loss ℓ ( f s ( x ) , y ) . Therefore, for all q ′ ∈ ∆ n , omitting the x dependency in the notations,

We conclude by multiplying by p ( x, y ) and integrating over X × Y . ( ii ) Let us first show that for any f ∈ F , T ↦→ D x ( q T , f ) is non-drecreasing. For that, let us define φ ( T ) ≜ E k ∼ q T ( f k | x,y ) [ ℓ ( f k ( x ) , y )] = ∑ n k =1 q T ( f k | x, y ) ℓ ( f k ( x ) , y ) for each T > 0 , x, y ∈ X ×Y , and f ∈ F . φ is a differentiable function of T , with derivative

Therefore, φ is non-decreasing. By the growth of the integral, we deduce that T ↦→ D x ( q T , f ) is non-decreasing.

Let T ′ ≥ T > 0 . Then D x ( q T ′ , f ) ≥ D x ( q T , f ) ≥ inf f ′ ∈F D x ( q T , f ′ ) for all f ∈ F , and by definition of the infimum, we have inf f ∈F D x ( q T ′ , f ) ≥ inf f ∈F D x ( q T , f ) . So D ∗ x is a nondecreasing function of T , and the same reasoning applies to D ∗ .

( iii ) There are two equalities to prove. First observe that

We deduce

which establishes the first equality. Then we conclude by substituting the definition of q T into F ( q T , f ) and simplifying the expression. Again, we drop the x dependency in the notations for readability.

Note that the last equality can be written as:

where we define the conditional free energy F x ( q T , f ) as

A.4 Rate distortion curve

We have established in the previous Section that aMCL moves the hypotheses f k towards the soft barycenter of soft Voronoi cells. We now describe the impact of temperature cooling on the position of these soft barycenters.

At high temperature, we observe that the soft barycenters (hence the hypotheses) converge towards the conditional barycenter E [ Y | X = x ] and fuze. When temperature decreases, they iteratively split into sub-groups. To capture the virtual number of hypotheses, we introduce the conditional rate-distortion function

where the target Y ∼ p ( y | x ) , the hypothesis position ˆ Y ∼ q ( f k | x ) follows a distribution over Y with q ( f k | x ) = ∫ Y q ( f k | x, y ) p ( y | x )d x , and I x ( ˆ Y ; Y ) is their mutual information.

The following Proposition shows that (30) corresponds in fact to the same optimization problem as minimizing the conditional free energy (29), and therefore as minimizing the entropy constrained distortion (21) conditionally on x . It then describes the shape of the parametric curve ( D ⋆ , R x ( D ⋆ )) , and provides a lower bound of the rate distortion. For this purpose, we introduce for each x ∈ X

the optimal conditional distortion when using a single hypothesis.

Proposition 7 (Rate-distortion properties) . For each x ∈ X , we have the following results.

(i) For each , minimizing the free energy

over all hypotheses positions f ∈ F ( X , Y n ) comes down to solving the optimization problem that defines (10) for each x ∈ X under decoupling Assumption 1.

(ii) R x is a non-increasing, convex and continuous function of D ⋆ , R x ( D ⋆ ) = 0 for D ⋆ ≥ D max x , and for each x the slope can be interpreted as R ′ x ( D ⋆ ) = -1 T when it is differentiable.

(iii) For each x , R x ( D ⋆ ) is bounded below by the Shannon Lower Bound (SLB)

where Y ∼ p ( y | x ) and H ( D ⋆ ) is the entropy of a Gaussian with variance D ⋆ .

These results are well-known properties of the rate-distortion function [11, 4, 24, 51, 6].

Proof. ( i ) Under the decoupling assumption 1, optimizing F comes down to optimizing F x for each x ∈ X . Consequently, we can use the known relationship between the rate-distortion curve and free energy minimization [4, 24, 59].

( ii ) See [4, 24].

( iii )

A key characteristic of the SLB is its asymptotic tightness in the low distortion regime for a broad range of probability distributions [32]. For some distributions, such as the Gaussian case, the SLB and the rate-distortion curve align perfectly for D ≤ D max x .

The Rate-Distortion interpretation of aMCL applies for each x , but defining a notion of global rate that is independent of the input is more challenging. This is because the distortion optimal distortion at a given temperature level T , inf f ∈F D x ( q T , f ) , can vary across inputs x . This is left for future work.

A.5 First phase transition

In the previous Section, we mentioned that aMCL predictions fuze at high temperature into the conditional barycenter E [ Y | X = x ] . Moreover, these fuzed predictions iteratively split as the temperature (or equivalently the distortion) decreases during training. With each split, the number of sub-groups formed by the hypotheses increases. The number of sub-groups, measured in bits, is captured by the rate-distortion R x ( D ⋆ ) . In this Section, we focus on the critical temperature corresponding to the first of these splits.

The predictions fuze into a single point at high temperatures so that the number of sub-groups is 1 and R x ( D ⋆ ) = 0 in this regime. The first splitting therefore occurs when R x ( D ⋆ ) > 0 . Correspondingly, we define the first critical temperature in Definition 1.

Definition 1. The first critical temperature T c 0 ( x ) for each x and the defined global variant T c 0 are defined as:

in the conditional and the non-conditional setting, where D max ≜ ∫ X p ( x ) D max x d x and D max x is defined in (7) .

Note that we have also equality in the definition T c 0 ( x ) = inf { T | inf f ∈F D x ( q T , f ) = D max x } , since inf f ∈F D x ( q T , f ) ≤ D max x . Moreover, recalling that D ∗ x : T ↦→ inf f ∈F D x ( q T , f ) is a nondecreasing function (Proposition 6 ( ii ) ), we can equivalently define T c 0 ( x ) as the generalized inverse T c 0 ( x ) = ( D ∗ x ) -1 ( D max x ) ( e.g. , Definition 1 in [17]). The same observations hold for T c 0 .

We also know that aMCL's predictions implicitly minimize the free energy (27) (see Proposition 7, part ( i ) ). Therefore, the conditional barycenter E [ Y | X = x ] is stable as long as the Hessian of the free energy is positive definite. This observation can alternatively be used to define the first critical temperature.

Interestingly, all these definitions are equivalent, as shown by the next Proposition. For this purpose, let us introduce the following covariance matrices:

where q ( y | f k , x ) denotes the posterior probability of assigning the point y to the hypothesis k , calculated using Bayes's rule [60]. At high temperatures, all hypotheses merge into the conditional barycenter of the distribution (9), and the matrices C k,k are equal to the data covariance matrix C ( x ) ≜ C ov ( X,Y ) ∼ p ( x,y ) [ Y | X = x ] .

Proposition 8 (First critical temperature) . We have the two following results.

(i) We know from [61] that for all x ∈ X , T c 0 ( x ) = 2 λ max ( C ( x )) .

(ii) Under decoupling assumption 1, T c 0 ≤ 2 sup x ∈X λ max ( C ( x )) .

Proof. ( i ) See [61].

( ii ) Let T = sup x ∈X T c 0 ( x ) . By definition of T 0 ( x ) and due to Proposition 6 ( ii ) , inf f ∈F D x ( q T , f ) ≥ inf f ∈F D x ( q T c 0 ( x ) , f ) ≥ D max x for all x ∈ X . Then, we deduce from (20):

Using ( i ) , we can write:

We now analyze why the predictions converge towards the conditional barycenter E [ Y | X = x ] . In the next Proposition, we first focus on the asymptotic regime, when the temperature T increases to infinity, and then extend this analysis at finite temperature.

Proposition 9 (Training Dynamics with Temperature) . For all x ∈ X and all k ∈ [ [1 , n ] ] , we have the following properties.

(i) If T = ∞ and f ∈ argmin f ∈F D ( q T , f ) , then f k ( x ) = E [ Y | X = x ] .

(ii) The gradient of the soft distortion can be re-written for each T > 0 as

Therefore, the first order approximation of ∇ f k D x in T writes:

When T → ∞ , aMCL reduces to standard risk minimization, so that a necessary condition of optimization is that all the hypotheses f ( x ) are located at the conditional barycenter E [ Y | X = x ] (Proposition 9, part ( i ) ). This conditional barycenter is stable and corresponds to the global minimizer of the aMCL training objective when T → ∞ . We further analyze the force that pushes all the hypotheses toward this conditional barycenter, by looking at an expansion (36) of ∇ f k D , which provides the direction of the hypotheses updates for each T > 0 (Proposition 9, part ( ii ) ). As T →∞ , we see that a global driving force, to which all the data points contribute, is pushing the hypotheses toward the barycenter in (37). As T decreases, local interactions corresponding to the higher-order terms appear and increase in amplitude. They are responsible for the phase transitions.

Proof. ( i ) By definition of the Boltzmann distribution

Therefore D ( q T , f ) - - - - → T →∞ 1 n ∑ n k =1 ∫ Y ℓ ( f k ( x ) , y ) p ( y | x )d y , and all the hypotheses solve the same quantization problem. A necessary condition to minimize D ( q T , f ) is then

( ii ) We note Z x,y ( T ) = ∑ n s =1 exp ( -∥ f s ( x ) -y ∥ 2 /T ) . Then, accounting for the stop _ gradient operator in (6), the gradient of the soft distortion writes

Using the series expansion of the exponential, we have for each T > 0 ,

We can rewrite the gradient of the soft distortion:

Keeping only the first term, and observing Z x,y ( T ) - - - - → T →∞ n we get:

This Section focuses on the first phase transition. However, multiple phase transitions may occur during training, as shown in Figure 4. It is interesting to note that, for deterministic annealing, the final phase transition occurs when the rate-distortion curve hits the Shannon Lower Bound defined in (33) (see [59]).

B Connection with the literature and discussion

B.1 Annealing at inference time

A promising direction for future research consists in using annealing at inference time. With this scheme, aMCL could be used to perform an input-dependent hierarchical clustering [53] at test time, similarly to the idea proposed in [55]. More precisely, we can store the model's parameters at different times of the training schedule ( i.e. , at several temperature levels). During inference, this allows replaying the temperature cooling for new test samples, by performing forward passes through each of the trained models.

Replaying this trajectory may have several advantages. Indeed, the hypotheses trajectory follows the rate-distortion curve and consequently explores recursively the modes of the distribution as the temperature decays. Crucially, at each critical temperature, when the hypotheses are about to split, they are exactly located at the barycenter of these modes. If we can track these splitting moments, for instance by counting the number of distinct virtual hypotheses at each step of the cooling schedule, we can perform a hierarchical clustering that iteratively uncovers the modes of the distribution.

B.2 Discussion on Stochastic simulated annealing

Non-deterministic simulated annealing [28, 52] is a promising research direction due to its strong convergence properties (see Hajek theorem [26]). It requires to define the state f of the system and

the optimization objective D ( f ) . At each step, the state f is updated to a neighbor state ˜ f based on a stochastic exploration criterion. The probability of accepting a neighbor state ˜ f depends on the objective variation D ( ˜ f ) -D ( f ) and the temperature T .

Our objective D ( f ) corresponds to the Distortion (1). However, the state of the system can be defined in various ways. In the non-conditional setting, [74] defines the state as the hypothesis positions ( f k ) k ∈ [ [1 ,n ] ] (similarly to the present work), while [26, 28, 52] defines it as the cluster assignation of each dataset sample.

In both cases, we expect that storing and updating this state using neural networks would be costly. Moreover, evaluating D ( ˜ f ) requires going through a validation set, which is time-consuming. Further investigation in this direction is left for future research.

B.3 Comparison with an additional baseline: Relaxed WTA with annealed ε

Relaxed WTA [63] suffers from a bias toward the barycenter. As suggested in [55], a natural idea to tackle this issue would be to anneal ε during training. In Figure 1, we have shown the results of this additional baseline on a Gaussian synthetic dataset, using a linearly decreasing epsilon during training ε ( t ) = ε 0 ( 1 -t 1000 ) with ε 0 = 0 . 98 .

In Figure 5, we show its training trajectory when using the same setup as in Figure 1. All hypotheses initially converge to the barycenter, then the winners gradually move towards the modes as ε decreases. As ε approaches 0 , only a few additional hypotheses escape from the barycenter to reach the modes, indicating that annealing does not solve the collapse issue of Relaxed-WTA.

Results on the UCI datasets confirm this qualitative analysis (Table 5 and Figure 6). In Figure 6, aMCL outperforms the best Relaxed-WTA variants on Distortion.

ε ≈ 0.98

ε  0.049

≈ 0.00098

C Experimental details in the synthetic data experiments

This Section provides a detailed description of the synthetic data experiments presented in the paper, particularly the results and visualizations in Figures 1, 3 and 4. Note that for each of those experiments, a two-hidden-layer Multi-Layer Perceptron (MLP) with 256 neurons per layer and ReLU activation functions was used, except for the last layer. The final layer of the MLP was duplicated to account for multiple hypotheses { f k ( x ) } n k =1 (followed by a tanh activation to restrict the output in the square [ -1 , 1] 2 ) and scores { γ k ( x ) } n k =1 (followed by sigmoid activation), similar to [42, 43], with n = 49 hypotheses in this study. At each epoch, 100,000 points were sampled from the corresponding synthetic datasets, and the models were trained for t epoch = 1000 epochs. Note also that the visualization of the output space in Figure 1 and Figure 4 is restricted to the square [ -1 , 1] 2 .

Settings of Figure 1 . In this Figure, the dataset consists of a mixture of three two-dimensional Gaussians located at µ 0 = ( -0 . 5 , -0 . 5) , µ 1 = (0 , 0 . 5) and µ 2 = (0 . 5 , -0 . 5) , each with a standard

deviation of 0 . 1 . In this Figure, each model was trained with a Stochastic Gradient Descent optimizer with a constant learning rate of 0 . 01 . For the Relaxed WTA run, we used ε = 0 . 1 , and for the Annealed MCL run, we used an exponential scheduler in the form T ( t ) = T 0 ρ t , with T 0 = 0 . 6 and ρ = 0 . 99 .

Settings of Figure 3 and Figure 4. In these experiments, the dataset is a conditional variant of that used in Figure 1, with a ground-truth distribution of the form p ( y | x ) = 1 3 ∑ 2 i =0 N (˜ µ i ( x ) , σ 2 ) , with σ = 0 . 1 , ˜ µ i ( x ) = xµ i for x ∈ [0 , 1] , and µ i is defined as above for i ∈ [ [0 , 2] ] . Here, aMCL was trained with a linear scheduler defined by T ( t ) = T 0 (1 -t t epoch ) , with T 0 = 1 . 0 . Note that in the results of Figure 4, we have ( t 1 , t 2 , t 3 , t 4 ) = (50 , 745 , 870 , 1000) with associated temperature values ( T ( t 1 ) , T ( t 2 ) , T ( t 3 ) , T ( t 4 )) ≃ (0 . 950 , 0 . 256 , 0 . 131 , 0 . 001) . Here, aMCL was trained with Adam optimizer [34], with a constant learning rate of 0 . 01 . In Figure 3, T c 0 (upper bound) correspond to the right-hand-side of (15) and was computed as 2 λ max ( ∫ Y yy t p ( y | x = 1)d y ) ≃ 0 . 46 , approximating the covariance matrix over 1,000 samples. Note that the distortion plotted in the curve of Figure 3 corresponds to the hard distortion (1) averaged over a validation set of 25,000 samples every 5 epoch.

D Additional results from UCI Datasets experiments

The UCI Regression Datasets [15] serve as a standard benchmark for conditional distribution estimation. This Section provides additional details and results from the experiments presented in Section 5.1.

Experimental setup. The experiments were conducted using the protocol from [29]. Each dataset is divided into 20 train-test folds, except the protein dataset, which had 5 folds, and the Year Prediction MSD dataset which used a single test-set split. We used the same neural network backbone for each baseline: a one-hidden layer MLP with ReLU activation function, containing 50 hidden units except for the Protein and Year datasets, for which 100 hidden units were used. In the WTA models, the final layer of the MLP was duplicated to account for multiple hypotheses { f k ( x ) } n k =1 (followed by no activation) and scores { γ k ( x ) } n k =1 (followed by sigmoid activation). The WTA-based methods (Relaxed-MCL, MCL, and aMCL) were trained with n = 5 . All models were trained with the Adam optimizer over 1,000 epochs with a constant learning rate of 0 . 01 . The data loading pipeline was adapted from the implementation of [27]. The best models for each dataset are highlighted in bold in Table 2 and Table 1, with the second-best models underlined.

Baselines. Table 2 of the main paper includes results from three baselines reported from Table 1 of [38], which we use as references for those benchmarks. 'PBP' stands for Probabilistic Back Propagation [29], and 'MC-dropout' corresponds to Monte Carlo Dropout [21]. Relaxed-MCL was trained with ε = 0 . 1 in Tables 1 and 2. Other values of ε are experimented in Appendix D.2.

Metrics. We used RMSE and distortion as metrics. RMSE is defined as RMSE = √ 1 N ∑ N i =1 ∥ ˆ y i -y i ∥ 2 , where ˆ y i denotes the estimated conditional mean, defined as ∑ n k =1 γ k ( x ) f k ( x ) for the WTA variants, and N is the number of samples in each test set. The distortion of the multi-hypotheses models was computed as Distortion = 1 N ∑ N i =1 min k ∈ [ [1 ,n ] ] ∥ ˆ y k -y i ∥ 2 .

Temperature scheduler. For training aMCL in those experiments, an exponential temperature scheduler was used, defined as T ( t ) = T 0 ρ t , with ρ = 0 . 95 and T 0 = 0 . 5 , where t denotes the current epoch and t epoch = 1 , 000 is the maximum number of epochs. When the temperature of the exponential scheduler was lower than 5 e -4 , we set the training back to the vanilla Winner-takes-all mode.

Evaluation details. During the evaluation, following [29], both input and output variables were normalized using the training data's means and standard deviations. For evaluation, the original scale of the output predictions was restored using the transformation f k θ ( x ) ↦→ µ train + σ train f k θ ( x ) where µ train and σ train are the empirical mean and standard deviation of the output variable, computed across the training set.

The results concerning the distortion metric on the UCI datasets are presented in Table 1, and their analysis is provided in Section 5.1.

D.1 Impact of the scheduler type in aMCL

The impact of the scheduler type on the performance on the UCI dataset is provided in Table 4. We see that for large dataset sizes, the exponential scheduler outperforms the linear one. Further study on the impact of the speed of the scheduler on the performance is left for further work.

D.2 Impact of ε in Relaxed WTA

RMSE compares the barycenter of the predicted distribution with the target positions. For this metric, Relaxed-MCL outperforms other approaches on the UCI datasets when ε is high (see Table 5 and Figure 6). This outcome is expected, as Relaxed-WTA is biased towards the distribution barycenter under this regime (see Figure 1 of the paper).

However, using RMSE for comparison discards valuable spatial distribution information of the hypotheses. Distortion, defined in (1), addresses this limitation by measuring quantization performance, which our theoretical analysis is based on. When focusing on the Distortion metric, Table 5 shows that lower values of epsilon ( e.g. , ε = 0 . 1 or annealed epsilon) consistently perform better than for ε = 0 . 5 across nearly all settings. This improvement is attributed to the reduced bias toward the barycenter in this configuration. Additionally, for Distortion, aMCL (trained with an exponential scheduler) generally outperforms Relaxed-WTA on the UCI datasets for the tested Relaxed WTA variants (see Tables 1, 5 and Figure 6). This is especially true on large datasets (Year, Protein).

Overall, aMCL demonstrates a strong balance between Distortion and RMSE compared to the baselines, especially on the largest datasets. Figure 6 further supports this trend, providing additional analysis on Year and Protein, where aMCL, MCL, and Relaxed-MCL are compared across different ε values in the (RMSE, Distortion) space.

Year

Protein

35

E=0.5

MCL

7

=0.5

MCL

aMCL

aMCL

30

6

25

RelaxedWTA

Relaxed WTA

Distortion

Distortion

5

20

4

15

E

0.2

3

=0.2

10

=0.1

2

0.1

=3

0.05

=0.05

5

1

8.900

8.925

8.950

8.975

9.000

9.025

9.050

9.075

9.100

4.20

4.25

4.30

4.35

4.40

4.45

4.50

RMSE

RMSE

E Application to speech separation

This Appendix provides a detailed description of the speech source separation experiments. Speech separation consists of isolating the audio signals of each individual speaker from a mixture in which they are simultaneously active. This task is of major interest for speech processing applications such as Automatic Speech Recognition [50, 44], Speaker Diarization [71], or singing voice extraction from music [31]. Most of the recent speech separation systems are deep neural networks trained in a supervised setting, where the ground truth of the separated speaker tracks is known [72].

E.1 Task description

Formally, let y 1 , . . . , y m ∈ R l denote the raw speech signals from m individual speakers, with l denoting the number of time frames in the corresponding audio recordings. Under anechoic (no reverberation) and clean (no background noise) conditions, the mixture can be expressed as x = ∑ m s =1 y s . Hence, the task is to provide estimates ˆ y 1 , . . . , ˆ y m of the isolated speech tracks from the mixture signal x using a neural network f θ . The quality of the estimation ℓ (ˆ y k , y s ) is commonly measured using the audio-domain distance Signal-to-Distortion Ratio (SDR) and its scale-invariant (SI) variant SI-SDR [70, 39]. Since there is a fundamental ambiguity concerning the best association ˆ y k ↦→ y s , separation systems are trained using the Permutation Invariant Training (PIT) loss:

where Σ is the set of permutations of [ [1 , m ] ] .

Instead, we propose to use the MCL framework to solve this assignment problem. Naming L MCL sep the resulting loss, L aMCL sep its annealed variant, and L Relaxed -WTA sep its relaxed variant, we define

This gives the separation task a new interpretation: the targets y s are samples drawn from a common distribution p ( y | x ) conditional on the mix x , the estimated speech signals ˆ y k = f k ( x ) are the output of neural networks f k , and the mix x is their common input.

Notice that (39) is subject to collapse. Furthermore, when n = m , (38) always finds the best assignation, and acts as a topline, while MCL acts as a baseline. The goal of these experiments is therefore threefold:

applying our proposed aMCL algorithm to a more realistic and data-intensive setting;

extending the analysis of aMCL to a more general setting where the underlying metric (SI-SDR) is non-Euclidean;

ensuring that MCL is a valid substitute for PIT in the context of source separation, thus opening up a new avenue for MCL research in this domain.

E.2 Experimental settings

E.2.1 Separation model architecture

We use the Dual-Path Recurrent Neural Network (DPRNN) [46] separation architecture for our experiments. This model follows the encoder-masker-decoder structure, originally proposed in ConvTasNet [47]. Each component of the network is described below:

Encoder. The encoder transforms the raw audio signal x ∈ R l , with l the number of time frames, into a sequence of features. It is implemented as a 1-D convolutional layer with F channels, kernel size K , and stride S . It outputs the sequence X ∈ R F × T where T is the number of encoded frames.

Masker. The masker of the DPRNN model first applies a sliding window of width W to the sequence of encoded features X . The input sequence is transformed from R F × T to X ′ ∈ R F × P × W , where P is the number of windows. DPRNN is composed of two recurrent layers, processing the sequence X ′ over two different directions (dual-path):

intra -chunk: processes the dimension W ,

inter -chunk: processes the dimension P .

Thus, the model learns local and global dependencies. The RNN is implemented as a bidirectional Long Short-Term Memory layer (LSTM) with hidden dimension H . The sequence X ′ is processed by B successive DPRNN blocks. The output of the last block is denoted X ′′ ∈ R H × P × W . Before decoding, the sequence is mapped back to the original dimensions using overlap-and-add. A final convolutional layer maps the sequence X ′′ to the masks M ∈ R n × F × T , where n is the number of predictions. The k -th source latent sequence ˆ Y k = M k ⊙ X is inferred by the element-wise product ⊙ between the k -th mask M k and the encoded sequence X .

Decoder. The decoder is implemented as a 1-D transposed convolutional layer with n groups, F channels, and a kernel of size K and stride S . It maps the latent sequences ˆ Y ∈ R n × F × T to the audio domain, and we note ˆ y ∈ R n × l the corresponding vector. Therefore, using the MCL framework notations, the output of the prediction model k is f k ( x ) = ˆ y k .

Scoring. We do not use scoring models in these experiments. Indeed, we will focus on the case where n = m and all predictions should be matched to a corresponding source. This ensures a fair comparison with the PIT baseline, which uses the same setting.

The model configuration used for our experiments is presented in Table 6. It corresponds to the model configuration from the original paper. The variables l , T , and P , which depend on the unknown input length l , are not specified.

E.2.2 Training configuration

Dataset. The source separation experiments are conducted on the Wall Street Journal mix dataset (WSJ0-mix) [30] which is the current benchmark dataset for speech source separation. It consists of clean utterances of read speech linearly combined to build synthetic speech mixtures. Several versions of the dataset are available to separate from 2-speaker to 5-speaker mixtures. In the present

work, we focus on the 2- and 3-speaker scenarios. Each version features 20000, 5000, and 3000 mixtures for training, validation, and testing respectively. The audio signals are samples at 8kHz.

Training objectives. Four types of systems are compared in the following experiments:

PIT: training with L PIT sep defined in (38),

MCL: training with L MCL sep defined in (39),

aMCL: training with L aMCL sep defined in (40),

Relaxed-WTA: training with L Relaxed -WTA sep defined in (41).

We use the SI-SDR score as the separation quality measure

where ⟨· , ·⟩ denotes the scalar product on R l .

Training parameters. The separation models are trained on 5-second audio segments. Since the utterances from the training set might have a longer duration, 5-second segments are randomly cropped inside the training utterances. The utterances shorter than 5s are removed. The batch size is set to 22 to reduce training time compared to the standard setup [47, 46, 68], in which a batch size of 4 is used. We verified that our models match the originally reported scores [46] when using this original setting. To analyze robustness to initialization, we train each method with 3 different seeds, and we report the inter-seed average scores and standard deviations. More precisely, the prediction scores are first averaged across the evaluation dataset for each single seed. Then these dataset-wise scores are used to compute the final average and standard deviations. Each model is trained on t epoch = 200 epochs, without early stopping. All the training set is processed during each epoch. Unless otherwise stated, the temperature scheduler is chunk linear:

with initial temperature T 0 ≈ 0 . 1 and t max = 100 . Note that at the end of the training, when t ∈ [ t max , t epoch ] , annealing is disabled by setting T ( t ) = 0 in order to fine-tune the hypotheses on the target MCL objective.

The neural network weights are updated using the Adam optimizer, with a learning rate set to 10 -3 . The learning rate is halved after every 5 epochs with no improvement in the validation metric. Separation models are trained on Nvidia A40 GPU cards. The DPRNN has 3.7 million parameters.

Evaluation. The separation scores are computed on the evaluation subset of the WSJ0-mix datasets. Two types of scores are presented:

PIT SI-SDR: the assignment is performed with PIT (38) as in standard source separation approaches;

MCL SI-SDR: the assignment is performed with MCL (39), and this score can be seen as a form of quantization error.

E.3 Experimental results

This Section presents the results of the experiments conducted on speech separation with the aMCL framework, and is organized as follows:

Appendix E.3.1: comparison of the separation performance of each approach;

Appendix E.3.2: impact of the number of hypotheses on the separation performance with MCL and aMCL;

Appendix E.3.3: analysis of the phase transitions of aMCL in the context of speech separation.

E.3.1 Separation performance on 2- and 3-speaker mixtures

Table 7 shows the average score of each separation model in the 2- and 3-speaker scenarios.

First, we observe that PIT and aMCL achieve similar performances, both in terms of the average PIT SI-SDR and MCL SI-SDR scores, and both in the 2- and the 3- speaker settings (see lines 1 and 3 of Table 7). Looking at the inter-seed variance, we further see that the differences in performances between these approaches are not significant. Note that MCL and aMCL have a O ( mn ) complexity, while PIT is O ( m 3 ) with the Hungarian algorithm [16]. Therefore aMCL reaches the same performance as PIT with a gain in terms of complexity. Further work includes an analysis of this complexity gap when the number of speakers is high, similarly to [66].

Second, we see that aMCL achieves better performance than MCL in the 2-speaker scenario, both in terms of PIT SI-SDR and MCL SI-SDR metrics (see lines 2 and 3 of Table 7). More precisely, we see that this performance discrepancy between MCL and both PIT and aMCL is due to a higher variance of the MCL training method. In fact, MCL has a higher inter-seed variance than aMCL and PIT in all settings. Indeed this method is known to be sensitive to initialization [49, 63]. In contrast, we see that aMCL is more robust to initialization, having a lower inter-seed variance and higher average score. This phenomenon was highlighted in Section A, which provided a theoretical analysis of the Euclidean case. This experiment suggests that our conclusions hold even for the source separation task, where the underlying loss function ℓ is non-Euclidean.

In the subsequent Sections, we further study these two claims. In Section E.3.2, we demonstrate the advantages of aMCL over PIT by letting the number of predictions n exceed the number of sources m . In Section E.3.3, we show that other conclusions of our theoretical analysis concerning phase transitions also seem to hold in the non-Euclidean setting.

↑

(a) PIT SI-SDR (

)

(b) MCL SI-SDR (

↑

)

E.3.2 Impact of the number of hypotheses in MCL and aMCL

In this Section, we study the impact of the number of hypotheses on the source separation performances. We focus on the 3-speaker scenario, and we increase the number of predictions. More precisely, we consider the cases n ∈ { 3 , 5 , 10 } . If n = 3 , the number of hypotheses is the same as the number m of sources to separate. If n ∈ { 5 , 10 } case the number of hypotheses is higher. Since the number of predictions and sources differ, it is not possible to compute the PIT SI-SDR metric, which is ill-defined in this setting. Therefore, we only present results for the MCL SI-SDR metric. In this experiment, we use a single seed across the considered methods to ensure a fair comparison. We report dataset-wise score statistics in Figure 7.

First, we observe that the separation performance increases as the number of predictions n increases, both for MCL and aMCL. This shows that the additional hypotheses are effectively used by both methods, despite the risk of collapse or convergence towards a suboptimal configuration. Recalling that MCL, aMCL, and PIT performances are similar in the case n = m = 3 , this demonstrates a strong advantage of the MCL family compared to PIT. Indeed, we can improve MCL SI-SDR score beyond what is achievable by PIT training, by increasing the number of predictions n . Moreover, this improvement comes at minimal computational complexity cost, as discussed in Section E.3.1. This property can be exploited to handle settings where the number of speakers is unknown. This is left to further work.

Second, we observe that there are no significant discrepancies between MCL and aMCL performances in this setting. This suggests that the main difference between these two approaches is their sensitivity to initialization.

3 Hyp.

5 Hyp.

10 Hyp.

0

5

10

15

20

MCL SISDR [dB]

(a) MCL

3 Hyp.

5 Hyp.

10 Hyp.

0

5

10

15

20

MCL SISDR [dB]

(b) aMCL

E.3.3 Phase transition in aMCL training

In this Section, we analyze the training trajectory of aMCL in speech separation. We focus on exponential schedulers defined by

where ρ = 0 . 9 is the decay factor, and T 0 ∈ { 0 . 1 , 5 , 23 } is the initial temperature. Figure 8 indicates the negative MCL SI-SDR score computed during training on a validation set. In order to compare meaningfully the different temperature schedules, we plot the separation score against the temperature, regardless of the training step when this temperature is reached.

We see that that aMCL exhibits phase transitions in this setting. Indeed, when we use a scheduler with high initial temperature T 0 ∈ { 5 , 23 } , the negative MCL SI-SDR metric stays on a plateau and then decreases after some critical temperature T c 0 has been reached. As expected, this critical temperature seems to be independent of the initial temperature T 0 > T c 0 . Moreover, when using a scheduler with initial temperature T 0 ≈ 0 . 1 < T c 0 , aMCL does not undergo a phase transition. This suggests that the phase transition phenomenon exhibited in Section A also exists in the non-Euclidean case, even though the properties of the barycenter and the shape of the Voronoi cells have to be redefined. Further work will analyze more exhaustively the impact of the scheduler type on this phase transition phenomenon. We hypothesize that the temperature decay speed at the critical temperature T c 0 will play a key role in this analysis.

10

-

4

10

-

2

10

0

Temperature (log scale)

-

17

.

5

-

15

.

0

-

12

.

5

-

10

.

0

-

7

.

5

-

5

.

0

-

2

.

5

0

.

0

MCL Negative SISDR [dB]

T

0

≈

0

.

1

T

0

≈

5

T

0

≈

23

T

c

0

(observed)

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: We strictly describe the scope of our article (introduction of a new training method, theoretical analysis, experimental validation).

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: we include this discussion at the end of Section 6 in a special entry. The main arguments are hyperparameter tuning (temperature schedule) and algorithm duration (due to annealing).

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: We strove to make apparent all the hypotheses of the provided theoretical results, either at the beginning of the Proposition that encapsulates them or factorized at the beginning of the theoretical Section A in the appendix. Proofs are complete, or we give a specific reference toward an external source proving it.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: All the experimental settings are extensively described in the Appendices C, D, and E to provide enough detail to reproduce the experiments.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: We provide code for reproducibility of our results. Moreover, the datasets featured in this work (UCI and WSJ0) are publicly available.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: The experimental details are provided in Appendices C, D, and E covering the synthetic data, UCI regression benchmark, and source separation studies, respectively.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: We always report indications of the statistical significance of the provided results, either in the main paper or in the Appendix.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: We mentioned in the Appendix E the number of parameters of the models we used as well as the hardware architecture used in our main experiments.

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: This research does not involve human subjects or participants, we referred the licensed dataset that we used (namely wsj0). We are not aware of any potential misuse of our algorithm.

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [Yes]

Justification: We did not discuss the potential societal impacts of our work, since we believe that it should not be clearly different from those of the previous learning algorithms based on Multiple choice learning.

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: No such safeguard has been put in place, because we do not believe that our method has a high risk of misuse.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: We have properly refered to the prior work (articles or code framework, datasets) we base our research on.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [NA]

Justification: We do not introduce any new assets.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer:[NA]

Justification: We did not use crowdsourcing.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: Our research did not involve experiments on human subjects.

Peer reviews:
comments: SUMMARY:
The paper introduces Annealed Multiple Choice Learning (aMCL), a method that integrates simulated annealing with Multiple Choice Learning (MCL), in applications where the output label may be ambiguous, and many values may be plausible given the same input.

The authors show that problems arise with the use of winner-takes-all in the optimization step of MCL, where the output space is partitioned in Voronoi cells based on the individual predictors. Instead, the authors suggest to soft-weight the predictors to replace hard-WTA assignments, in a way that depends on a temperature parameter. Low values of the temperature make the weighting scheme behave close to hard-WTA (sharpening the softmin distribution of weights).

STRENGTHS:
- The integration of simulated annealing and MCL is novel and effective;  I also like that the objective in eq. (6) becomes fully differentiable, without any non-differentiable assignment or operators.

- Extensive theoretical analysis.

- Extensive validation over a large number of datasets, although from the results in seems that epsilon-MCL seems to perform better, overall.

WEAKNESSES:
- The introduction of the temperature schedule as hyperparameter complicates the use of the method, while it seems that epsilon-MCL would be both easier (one hyperparameter, but fixed, without schedule) and overall better performing.
- I am not sure how representative the UCI datasets are of real world problems, but that is not a big problem.

- Minor: typo?  'broaden impact' -> 'broadeR impact'?

QUESTIONS:
N/A

LIMITATIONS:
- Since the authors argue that WTA is sensitive in particular to initialization, it would have been interesting to see experiments comparing the robustness of vanilla MCL vs aMCL over different initializations.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-10

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nI read the response and found it convincing, and I appreciate the new results.  I decided to raise my score.  \n--\nThis paper aims to tackle the local minima problems in MCL optimization.  Inspired by simulated annealing, a temperature-controlled soft assignment is used and then soft-cells are directly optimized in the MCL process.  Some theory and 2d analysis sheds insight into why this approach works.  This work seems well motivated and reasonable, but the experimental gains on real datasets are quite small.  Additionally, the performance seems similar to the older epsilon-MCL approach, in which some gradient is used to increase the weighting on non-selected cells.  \n\nnotes from reading the paper: \n  -Multiple choice learning handles ambiguous tasks by producing a set of hypotheses.  \n  -Hypotheses trained using winner-take-all technique, to encourage diversity.  \n  -This paper aims to apply simulated annealing ideas to MCL.  \n  -Try to increase variance in the annealing process, to address MCL's issue of falling into local minima based on the initialization.  \n  -aMCL uses softmin assignment, followed by gradient steps on the soft cells.\n\nSTRENGTHS:\n-The basic theoretical analysis is nice, providing relevant bounds on the performance of the algorithm.  \n  -The analytical results also nicely complement the theory.\n\nWEAKNESSES:\n-The empirical results seem lackluster.  Epsilon-MCL and aMCL seem to work roughly equally well on the UCI datasets.  For the speech separation datasets, the improvement over MCL seems quite small, and Epsilon-MCL doesn't have results presented.\n\nQUESTIONS:\n-Because it could be difficult to tune the temperature schedule over the course of training, I wonder if it could be possible to move the temperature schedule into the inference process (a bit like what's done in diffusion generative models)?  I.e., we would train a few models with different temperature levels, and force the cell assignment to be consistent with the higher temperature level.\n\nLIMITATIONS:\nLimitations aren't discussed in much detail in the paper.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-11"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes to apply deterministic simulated annealing to multiple choice learning (MCL) as a means to mitigate some of the drawbacks associated with the winner-takes-all (WTA) scheme used to train MCLs, such as sensitivity to initialization and hypothesis collapse. They demonstrate that the proposed annealed CML (aCML) method works well in practice, on par with previous approaches, and allows for an interesting and sound theoretical analysis of the training trajectory.\n\nSTRENGTHS:\n- The proposed method, aMCL, is very well motivated both theoretically and based on existing open research questions in multiple choice learning. Moreover, the mathematical development of the paper is very clear and seems sound.\n- The paper is very well written and easy to follow.\n\nWEAKNESSES:\n- The main weakness of the paper is that the experimental results, with the exception of the quite enlightening toy experiments, are somewhat underwhelming and aCML does not seem to outperform the baselines in most cases. I think the paper could benefit from a more in-depth discussion as to why that is the case.\n- Admittedly, I am not an expert on the CML literature and maybe guilty of hindsight bias, but the idea of applying simulated annealing to CML seems quite intuitive and, one might argue, somewhat incremental. For instance, simply annealing $\\epsilon$ in $\\epsilon$-CML (Rupprecht et al., 2017) seems quite natural, and perhaps a good baseline for the authors to compare against. That does not take away from the solid mathematical motivation and theoretical results in the paper though.\n\n### Minor issues\n- Line 201: “in no longer” should probably be “is no longer”.\n- The work of Rupprecht et al. is mentioned under different names (Relaxed WTA, $\\epsilon$-WTA and $\\epsilon$-CML) which can be a bit confusing.\n- It is not entirely clear to me what the authors want to show in Figure 2.\n- The robustness result in Figure 8 is quite interesting and could be expanded upon (perhaps with more random seeds) or at least mentioned in the main paper.\n\n### References\nChristian Rupprecht, Iro Laina, Robert DiPietro, Maximilian Baust, Federico Tombari, Nassir Navab, and Gregory D Hager. Learning in an uncertain world: Representing ambiguity through multiple hypotheses. In Proceedings of the IEEE international conference on computer vision, pages 3591–3600, 2017.\n\nQUESTIONS:\n1. Do the authors compare against the score-based method of Letzelter et al.? In line 232, it is suggested that this was one of the baselines, but it was not mentioned again in the paper, unless I missed it.\n2. On a similar note, could the authors elaborate on the relation between aMCL and rMCL (Letzelter et al., 2023)? It seems to be both have similar objectives, but the assignments in aCML are a function of the temperature of the system and the loss function, while in rMCL the assignments are learnable. If that is correct, would be fair to say that, given enough data and sufficient learnable parameters, we can always expect rMCL to outperform aCML?\n3. It is somewhat surprising that aCML does not outperform $\\epsilon$-CML in many cases. Do the authors have any intuition as to why that could be? Is it because the temperature schedule is hard to tune or aCML would require a larger number of epochs to fully converge? Or maybe the datasets are not particularly sensitive to bias of $\\epsilon$-CML?\n4. Have the authors considered stochastic simulated annealing as well?\n\n### References\n\nLetzelter, Victor, et al. \"Resilient Multiple Choice Learning: A learned scoring scheme with application to audio scene analysis.\" Advances in neural information processing systems 36 (2023).\n\nLIMITATIONS:\nThe limitations section is very well written and covers all possible shortcomings of the model and analysis proposed in the paper that I could think of.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-07-19"}