Variational Learning is Effective for Large Deep Networks

TL;DR

This paper provides extensive empirical evidence demonstrating that variational learning, specifically using the IVON optimizer, is effective for training large neural networks like GPT-2 and ResNets, outperforming or matching Adam.

Contribution

The paper introduces IVON, an optimizer that makes variational learning practical for large networks, with improved uncertainty estimates and new applications in LLM finetuning and model merging.

Findings

IVON matches or outperforms Adam in large network training

IVON provides better predictive uncertainty estimates

Variational learning is shown to be effective for large neural networks

Abstract

We give extensive empirical evidence against the common belief that variational learning is ineffective for large neural networks. We show that an optimizer called Improved Variational Online Newton (IVON) consistently matches or outperforms Adam for training large networks such as GPT-2 and ResNets from scratch. IVON's computational costs are nearly identical to Adam but its predictive uncertainty is better. We show several new use cases of IVON where we improve finetuning and model merging in Large Language Models, accurately predict generalization error, and faithfully estimate sensitivity to data. We find overwhelming evidence that variational learning is effective.