Single-Model Ensembling via Test-Time Stochastic Ensembles

Ensemble methods are a well-established paradigm of machine learning and have recently emerged as a successful technique for obtaining inference-time gains in deep neural networks. In contrast to prior work, we create a test-time stochastic ensemble (TTSE) from a single transformer by generating multiple stochastic forward passes from a frozen checkpoint and aggregating their outputs, without any additional training, checkpoints, or weight updates. We carefully investigate the precise design of the stochasticity and ensembling mechanism, finding that (1) adding isotropic Gaussian noise to Query and Key vectors is the most effective; (2) averaging probabilities is better than averaging logits; and (3) increasing the sample size enhances the differentiability of correct reasoning traces from incorrect ones. Across next-token prediction and reasoning tasks, our method consistently improves upon the deterministic-inference baseline and alternative stochastic methods, suggesting that successful ensemble methods do not actually require training multiple models, which are very expensive in practice.