Controllable Diversity in Normalization-Based Implicit Ensembles via Softmax-Temperature Modulation

Technical Summary: Controllable Diversity in Normalization-Based Implicit Ensembles via Softmax-Temperature Modulation

Problem Statement
Deep ensembles remain the gold standard for uncertainty quantification in deep learning, offering reliable calibration and out-of-distribution (OOD) detection by training multiple independent models. However, their computational and memory costs scale linearly with the number of ensemble members, becoming prohibitive for large-scale models. Implicit ensembles address this by sharing a single backbone across members, introducing per-member modulation or partitioning to simulate ensemble behavior in a single training run.

A critical limitation of existing implicit ensemble methods is the lack of control over member diversity. Diversity is a primary determinant of ensemble quality, yet current approaches fix it at initialization (e.g., adapter-based methods like BatchEnsemble) or hard-code it into the architecture (e.g., partitioning methods like Packed Ensembles). Consequently, there is no mechanism to shape or navigate the trade-off between accuracy and calibration during training. Furthermore, existing methods often require bespoke architectural redesigns, making them difficult to apply to pretrained models.

Methodology: $\sigma$N-Ens
The authors introduce $\sigma$N-Ens, a normalization-based implicit ensemble that treats each member as a distinct task within a multi-task learning (MTL) framework.

1. Architecture ($\sigma$Norm): The method replaces standard normalization layers (BatchNorm or LayerNorm) in a shared backbone with task-specific layers. For each member $t$, the normalization is modulated by a sigmoid-bounded scaler:
   $$\sigma\text{Norm}_t(x) = \sigma(\gamma_t) \odot \text{norm}_t(x)$$
   where $\sigma(\gamma_t) = \frac{1}{1 + e^{-\gamma_t}}$. These scalers act as per-feature importances. Because only normalization layers are replicated, the backbone weights (convolutional or transformer) remain fully shared, allowing the method to wrap any standard architecture and adapt pretrained models via short fine-tuning.

2. Controllable Diversity via Softmax-Temperature: To actively shape diversity during training, the authors introduce a softmax-temperature regularizer applied to the importance vectors ($\sigma(\gamma)$) across members for each feature.
   $$\mathcal{L}_{div} = \lambda \sum_{\ell \in \mathcal{R}} \sum_{i} \sum_{m=1}^{M} \log \left[ \text{softmax}_{\tau}(\sigma(\gamma_{\cdot, \ell, i})) \right]_m$$
   Here, $\tau$ acts as a temperature hyperparameter that sets the equilibrium level of sharing.

   • Low $\tau$: Drives the distribution toward one-hot, forcing features to be owned by a single member (maximal diversity, akin to hard partitioning).
   • High $\tau$: Drives the distribution toward uniform, encouraging full sharing of features.
   • Intermediate $\tau$: Creates a "soft $M$-choose-$k$" allocation, where features are shared among a subset of $k$ members.
     Unlike standard regularization strength ($\lambda$), which determines how hard the network is pushed toward a target, $\tau$ selects which equilibrium the system settles into. This allows practitioners to trace an accuracy-calibration frontier by sweeping $\tau$.

3. Modulation Uncertainty: The authors reframe the epistemic uncertainty expressed by this ensemble as modulation uncertainty. Since the backbone is shared and deterministic, the uncertainty arises solely from the reweighting of shared features by the modulation parameters. This is structurally distinct from the uncertainty derived from differing feature extractors in full deep ensembles.

Key Contributions

• Controllable Diversity: The paper presents the first mechanism in the implicit ensemble family to shape member diversity during training via a single hyperparameter ($\tau$), rather than fixing it at initialization or construction. This enables navigation of the accuracy-calibration Pareto frontier.
• Modulation Uncertainty: The authors identify and define "modulation uncertainty" as the specific epistemic component expressed by shared-backbone ensembles. They establish that this uncertainty maintains calibration under input corruption but exhibits weaker OOD detection compared to full ensembles, as the shared backbone cannot diverge on inputs it does not encode.
• Architecture Agnosticism and Efficiency: By operating solely on normalization layers, $\sigma$N-Ens applies to both convolutional and transformer backbones without invasive redesign. It can convert existing pretrained models into calibrated ensembles with minimal fine-tuning, avoiding the need to train members from scratch.

Results
The method was evaluated on CIFAR-10/100, ImageNet, and SST-2 using ResNets, WideResNets, and Vision Transformers (ViT).

• Performance vs. Cost: $\sigma$N-Ens matches or outperforms deep ensembles in terms of accuracy and Negative Log-Likelihood (NLL) while using a fraction of the parameter count (e.g., on CIFAR-100 with WideResNet, it achieves state-of-the-art accuracy with significantly fewer parameters).
• Scalability: Unlike partitioning methods (e.g., Packed Ensembles) which degrade and collapse as the number of members ($M$) increases due to capacity constraints, $\sigma$N-Ens scales gracefully. As $M$ grows, the soft sharing mechanism pools capacity effectively, allowing performance to improve or saturate rather than collapse.
• Calibration under Shift: The method maintains robust calibration under distribution shift (input corruption), with Expected Calibration Error (ECE) remaining nearly flat as corruption severity increases.
• OOD Detection: Consistent with the theoretical analysis, OOD detection is weaker than in full deep ensembles, particularly when using pretrained backbones where the feature representation is fixed. However, detection performance improves when training from scratch, as members shape the shared representation under diversity pressure.
• Hyperparameter Sensitivity: The final performance is largely insensitive to the regularization strength $\lambda$ once converged, making $\tau$ the sole critical hyperparameter for practitioners to tune.

Significance
The paper argues that controllable diversity transforms implicit ensembles from a static architectural choice into a dynamic, tunable system. By decoupling diversity from initialization and architecture, $\sigma$N-Ens offers a practical alternative to training full deep ensembles, especially as model sizes grow. The ability to convert a single pretrained model into a calibrated ensemble with a short fine-tune addresses a significant gap in deploying uncertainty-aware deep learning in resource-constrained or risk-sensitive domains. The work clarifies the trade-offs inherent in shared-backbone ensembles, explicitly linking the structural limitation (shared feature extractors) to the observed weakness in OOD detection, while highlighting the strength in calibration robustness.