[Paper Review] Optimal Training of Mean Variance Estimation Neural Networks

Laurens Sluijterman, Eric Cator+1 more|arXiv (Cornell University)|Feb 17, 2023

Fault Detection and Control Systems4 citations

TL;DR

This paper proposes two key improvements for training Mean Variance Estimation (MVE) neural networks: a warm-up phase that fixes variance during initial mean-only training, and separate regularization for mean and variance predictions. The authors demonstrate through experiments on UCI datasets that these practices significantly improve model stability and performance, with separate regularization often yielding substantial gains over equal regularization.

ABSTRACT

This paper focusses on the optimal implementation of a Mean Variance Estimation network (MVE network) (Nix and Weigend, 1994). This type of network is often used as a building block for uncertainty estimation methods in a regression setting, for instance Concrete dropout (Gal et al., 2017) and Deep Ensembles (Lakshminarayanan et al., 2017). Specifically, an MVE network assumes that the data is produced from a normal distribution with a mean function and variance function. The MVE network outputs a mean and variance estimate and optimizes the network parameters by minimizing the negative loglikelihood. In our paper, we present two significant insights. Firstly, the convergence difficulties reported in recent work can be relatively easily prevented by following the simple yet often overlooked recommendation from the original authors that a warm-up period should be used. During this period, only the mean is optimized with a fixed variance. We demonstrate the effectiveness of this step through experimentation, highlighting that it should be standard practice. As a sidenote, we examine whether, after the warm-up, it is beneficial to fix the mean while optimizing the variance or to optimize both simultaneously. Here, we do not observe a substantial difference. Secondly, we introduce a novel improvement of the MVE network: separate regularization of the mean and the variance estimate. We demonstrate, both on toy examples and on a number of benchmark UCI regression data sets, that following the original recommendations and the novel separate regularization can lead to significant improvements.

Motivation & Objective

• To address convergence instability in MVE networks, a common issue in uncertainty estimation for regression.
• To validate the original Nix and Weigend (1994) recommendation of a warm-up period to stabilize training.
• To investigate whether separate regularization of mean and variance predictions improves model performance.
• To provide practical, empirically grounded training recommendations for MVE networks in real-world regression applications.

Proposed method

• Implement a warm-up phase where only the mean head is trained with a fixed, constant variance initialized via bias setting.
• Use gradient clipping during joint optimization of mean and variance to improve training stability.
• Apply separate L2 regularization coefficients for the mean and variance sub-networks, optimized independently via cross-validation.
• Train the MVE network using negative log-likelihood loss: $\mathcal{L}(\theta) = \sum_{i=1}^{N} \frac{1}{2}\log(\sigma^2_\theta(\bm{x}_i)) + \frac{1}{2}\frac{(y_i - \mu_\theta(\bm{x}_i))^2}{\sigma^2_\theta(\bm{x}_i)}$.
• Use softplus or exponential transformation on the variance output to ensure positivity.
• Evaluate performance on UCI regression benchmarks using RMSE and log-likelihood metrics across multiple random seeds and 5–10-fold cross-validation.

Experimental results

Research questions

• RQ1 Does following the original warm-up recommendation from Nix and Weigend (1994) resolve convergence issues in MVE network training?
• RQ2 Is there a performance benefit to optimizing the mean and variance jointly after warm-up, or is fixing the mean during variance training preferable?
• RQ3 Does applying separate regularization coefficients to the mean and variance sub-networks lead to improved predictive performance compared to shared regularization?
• RQ4 How do the optimal regularization strengths for mean and variance compare across different regression datasets?

Key findings

• The warm-up period, where variance is fixed and only the mean is trained initially, significantly improves training stability and prevents the network from focusing prematurely on high-performing regions.
• Without a warm-up, MVE networks often fail to learn in poorly fitted regions, leading to suboptimal RMSE and log-likelihood scores.
• Separate regularization of the mean and variance consistently improves performance across multiple UCI datasets, with the variance typically requiring a regularization strength an order of magnitude larger than that for the mean.
• For the energy and yacht datasets, separate regularization led to markedly better results than equal regularization, with improvements visible even when the optimal regularization constants were similar.
• The difference between fixing the mean or optimizing it jointly with the variance after warm-up was not substantial, though joint optimization was slightly more effective when separate regularization was not used.
• In cases like Boston Housing, where the problem is relatively simple, all strategies performed similarly, suggesting that warm-up and separate regularization are most beneficial in complex or noisy regression tasks.