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

We thank the reviewer for their suggestions to improve the quality of this manuscript.

---

Comparison with baselines

From the results in seems that epsilon-MCL perform better, overall.

We compare epsilon-MCL (refered to as Relaxed-WTA in the following), MCL and aMCL on the UCI datasets for two metrics: RMSE (line 239) and Distortion (1).

RMSE compares the barycenter of the predicted distribution with the target positions. For this metric, Relaxed-WTA outperforms other methods for high values of $\varepsilon$ (e.g., with $\varepsilon = 0.5$, see Table B). This is expected, since Relaxed-WTA is biased towards the distribution barycenter, especially for high values of $\varepsilon$ (see Figure 1 of the paper).

However, barycenter comparison discards information concerning the spatial distribution of the hypotheses. Distortion (1) corrects this issue by measuring quantization performance, and we used it in our theoretical analysis.

For this metric, aMCL outperforms Relaxed-WTA in most cases, and especially for large datasets (Year, Protein, see Table A). Hence, aMCL strikes a balance between RMSE and Distortion (See Figure B).

In order to verify this behavior on audio data, we have additionally trained Relaxed-WTA on the speech separation data. We show that aMCL still outperforms Relaxed-WTA in this more realistic setting (see Table C of the supplementary pdf).

It would have been interesting to see experiments comparing the robustness of vanilla MCL vs aMCL over different initializations.

We thank the reviewer for this suggestion. We have extended the experiments presented in Figure 8 of the Appendix by comparing PIT, MCL, aMCL and Relaxed-WTA for 3 different seeds on audio data (see Table C of the supplementary pdf).

First, we confirm that PIT, which performs perfect assignation, acts as a topline in this experiment. Moreover, the difference between PIT and aMCL is not statistically significant. Therefore, aMCL reaches the same performances as PIT, while improving from $\mathcal{O}(m^3)$ to $\mathcal{O}(nm)$ in terms of complexity (noting $m$ the number of speakers and $n$ the number of hypotheses) [A]. This complexity gap is best exploited when the number of speakers is high, similarly to [B], and this will be the object of further work.

Second, we see that aMCL is better performing than MCL when we average over seeds. Moreover, aMCL has lower inter-seed variance than MCL. This validates our theoretical analysis, which suggests that aMCL guides the hypotheses towards a better local minimum than MCL, independently of the initialization.

---

Limitations

The introduction of the temperature schedule as hyperparameter complicates the use of the method.

We thank the reviewer for raising this issue. The temperature schedule is indeed a new degree of freedom that may require some tuning.

However, we conjecture that the theoretical analysis of aMCL convergence will lead to the characterization of optimal temperature schedules, similarly to Hajek theorem for deterministic annealing [22]. We leave this analysis for future work.

How representative are the UCI datasets of real world problems ?

The UCI datasets correspond to real world tabular data for 1D predictions, and constitute a widely used benchmark to evaluate uncertainty quantification algorithms. We have included it our manuscript in order to compare ourselves to the customary multi-hypothesis baselines (see Table 1 of the main paper).

The datasets are sorted by size in Tables A and B. Interestingly, aMCL obtains the best Distortion results for large datasets (Year, Protein, see Table A and Figure B), which are the most realistic. This will be emphasized in the main paper.

[A] Jack Edmonds and Richard M Karp. Theoretical improvements in algorithmic efficiency for network flow problems. Journal of the ACM (JACM), 19(2):248–264, 1972.

[B] Hideyuki Tachibana. Towards listening to 10 people simultaneously: An efficient permutation invariant training of audio source separation using sinkhorn’s algorithm. In ICASSP 2021-2021 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pages 491–495. IEEE, 2021.

We thank the reviewer for their insightful comments. We provide here a detailed answer to the raised concerns.

---

Performance of aMCL

The experimental gains on real datasets are quite small. For the speech separation datasets, the improvement over MCL seems quite small.

To compare MCL and aMCL we provide two additional experimental validations. For the UCI datasets, we experimented with an additional temperature schedule. For the speech separation data, we trained PIT, MCL, and aMCL with 3 seeds to measure sensitivity to initialization. See Tables A and C of the supplementary pdf.

On the UCI datasets, we observe that aMCL is competitive with MCL for the Distortion metric, and better for the RMSE metric. This effect is exacerbated on large datasets. This suggests that with careful temperature schedule tuning, aMCL can outperform MCL.

On the audio data, we observe that aMCL has better inter-seed average PIT-SISDR score, and lower inter-seed variance than MCL in all settings. This is consistent with our theoretical analysis, which suggests that aMCL is less sensitive to initialization than MCL.

The performance seems similar to the older epsilon-MCL approach. Epsilon-MCL and aMCL seem to work roughly equally well on the UCI datasets. For the speech separation datasets, Epsilon-MCL doesn't have results presented.

We thank the reviewer for this suggestion. We provide two additional experiments. On the UCI datasets, we experiment with different values of $\varepsilon$ for the epsilon-MCL baseline (refered to as Relaxed-WTA in the following), in Table B and Figure B. On the speech separation data, we compare Relaxed-WTA with the other approaches (Table C).

On the UCI datasets and for the RMSE metric, we observe that Relaxed-WTA outperforms all other approaches for high values of $\varepsilon$, but not for lower values. This is expected since Relaxed-WTA is biased towards the barycenter for high values of $\varepsilon$, and RMSE compares the predicted distribution barycenter with the target positions.

However, barycenter comparison discards information concerning the spatial distribution of the hypotheses. Distortion (1) corrects this issue by measuring quantization performance, and we used it in our theoretical analysis.

For this metric, we observe that aMCL outperforms Relaxed-WTA in most cases, and especially for large datasets. A higher $\varepsilon$ in Relaxed-WTA comes at the cost of a higher Distortion error. Comparatively, aMCL strikes a balance between RMSE and Distortion (Figure B).

On the audio data, we observe an advantage of aMCL over Relaxed-WTA for the PIT SI-SDR metric on the 2-speaker dataset. Preliminary experiments suggests that this is also true on the 3-speaker dataset, where Relaxed-WTA reaches a PIT SI-SDR score of $9.64 \pm 0.03$ compared to $10.00 \pm 0.21$ for aMCL (average and standard deviation over 3 seeds).

---

Temperature schedule

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.

The idea of moving the temperature schedule into the inference process is interesting: with this, aMCL can be used to perform hierarchical clustering at test time.

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 allow to replay the temperature cooling for new test samples.

Replaying this trajectory has 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.

We appreciate the reviewer's suggestion, which presents an opportunity for further research. We hope our response addresses their question and we remain available for additional discussion on this topic.

It could be difficult to tune the temperature schedule over the course of training.

The temperature schedule is indeed a new degree of freedom that may require some tuning.

However, we conjecture that the theoretical analysis of aMCL convergence will lead to the characterization of optimal temperature schedules, similarly to Hajek theorem for deterministic annealing [22]. We leave this analysis for future work.

---

Limitations

Limitations aren't discussed in much detail in the paper.

We thank the reviewer for raising this issue. In addition to the "Limitations" paragraph that we had already provided at the end of the main paper, we will update Sections 3 and 5 in order to provide more insights into the specific challenges raised by aMCL. In particular, we will insist on the temperature schedule choice, and the potentially longer training time of optimal schedules.

We thank the reviewer for their detailed suggestions.

---

Experimental validation

It is surprising that aMCL does not outperform epsilon-MCL in many cases.

RMSE (line 239) compares the barycenter of the predicted distribution with the target positions. For this metric, epsilon-MCL (refered to as Relaxed-WTA in the following) outperforms other approaches on the UCI datasets if $\varepsilon$ is high ($\varepsilon = 0.5$ in the main paper). This is expected, as Relaxed-WTA is biased towards the distribution barycenter in this regime (see Figure 1 of the paper).

However, barycenter comparison using RMSE discards spatial distribution information of the hypotheses. Distortion (1) corrects this issue by measuring quantization performance, on which our theoretical analysis relies.

Focusing on the Distortion metric, aMCL (trained with an exponential scheduler) outperforms Relaxed-WTA both for the UCI datasets and for the audio data (see Tables A and C of the supplementary pdf). This is especially true on large datasets (Year, Protein, audio data).

aMCL does not seem to outperform the baselines in most cases.

Looking at the performance on the UCI datasets in Table A, we see that aMCL provides a good tradeoff between Distortion and RMSE compared to the baselines, especially on the largest datasets. Figure B confirms this trend with further analysis on Year and Protein, comparing aMCL, MCL, and Relaxed-MCL for several $\varepsilon$ in the (RMSE, Distortion) space.

On the audio data, we trained all compared methods on 3 different seeds (Table C). These experiments extend the findings presented in Figure 8 of the Appendix: aMCL has a better inter-seed average score and lower inter-seed variance than MCL for the PIT SI-SDR. This suggests that aMCL is more robust to initialization than MCL, which is consistent with our theoretical analysis.

Annealing in epsilon-MCL seems a good baseline.

Using annealing in Relaxed MCL is an interesting idea: it may reduce the bias toward the barycenter of this method.

On the synthetic dataset of Figure 1, Relaxed-WTA with annealed $\varepsilon$ has the following trajectory (Figure A). All hypotheses initially converge to the barycenter, then the winners gradually move towards the modes as $\varepsilon$ decreases. As $\varepsilon$ approaches 0, only 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 B and Figure B). In Figure B, aMCL outperforms the best Relaxed-WTA variants on Distortion.

---

Connection with the literature

Do the authors compare against the score-based method of Letzelter et al. ?

In synthetic and UCI experiments, 'MCL' refers to rMCL, the score-based method of Letzelter et al. (Figure 1, Table 1, and Table 3 of the paper). For audio experiments, scoring is not relevant since all sources are active.

Could the authors elaborate on the relation between aMCL and rMCL.

rMCL uses a hard assignment based on the Winner-takes-all scheme (3) for the prediction heads $\\{f_{k}\\}$. It is not learned, but determined by the hypotheses positions.

aMCL uses a soft assignment based on a temperature schedule (5) to train the prediction heads $\\{f_k\\}$ (6). With high temperatures, the assignment is uniform. As the temperature decreases, the assignment converges toward the WTA scheme.

Using a soft assignment with annealing guides the optimization process of aMCL towards a good local minimum. This is critical since optimizing the Distortion (1) is difficult by gradient descent (this task is NP-hard [1,9,42]).

Note that rMCL and aMCL are however identical with respect to the scoring heads $\\{\gamma_k\\}$, which are trained using the same loss $\mathcal{L}_{\mathrm{scoring}}$ (4).

Have the authors considered stochastic simulated annealing as well?

Stochastic simulated annealing [23,46] is a promising research direction due to its strong convergence properties (see Hajek theorem [22]).

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 $\tilde{f}$ based on a stochastic exploration criterion. The probability of accepting a neighbor state $\tilde{f}$ depends on the objective variation $D(\tilde{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, [A] defines the state as the hypothesis positions $\\{f_k\\}$ (similarly to the present work), while [22, 23, 46] defines it as the cluster assignation of each dataset sample.

In both cases, storing and updating this state using neural networks is costly. Moreover, evaluating $D(\tilde{f})$ requires going through a validation set, which is time-consuming.

Further investigation in this direction is left for future research.

---

Additional comments

It is not entirely clear to me what the authors want to show in Figure 2.

Figure 2 illustrates the training trajectory of MCL and aMCL in the rate-distortion space. It shows that MCL has a constant rate during training, i.e., a constant number of hypotheses with distinct positions. In contrast, aMCL's rate varies during training: at high temperatures, hypotheses merge into a single cluster, then recursively subdivide as the temperature decreases.

Crucially, this Figure shows that MCL always has a higher rate than aMCL. Yet, it is known that the optimization procedure is more difficult in this regime [52]. This motivates the use of annealing and will be made clearer in the next revision.

[A] Zeger, Kenneth, Jacques Vaisey, and Allen Gersho. "Globally optimal vector quantizer design by stochastic relaxation." IEEE Transactions on Signal Processing 40.2 (1992): 310-322.