Mimetic Initialization of MLPs

TL;DR

This paper investigates whether mimetic initialization, previously effective for spatial mixing layers, can extend to MLP channel mixing by proposing a simple first-layer mean shift in $W_1$ guided by weight-space structure observed in pretrained models. It analyzes covariance patterns across populations of networks and demonstrates that a small nonzero mean initialization improves early training on tasks like CIFAR-10 and ImageNet-1k, and can complement existing mimetic inits for convolutional and self-attention layers. The findings suggest that a portion of the benefits of pretraining arises from the geometry of the weight space that facilitates optimization, especially in data-limited settings, and that population-level weight-space analysis can uncover new initialization signals. While benefits are strongest in short training, the approach offers a practical, model-agnostic increment to training efficiency and opens avenues for exploring other simple, population-informed initializations.

Abstract

Mimetic initialization uses pretrained models as case studies of good initialization, using observations of structures in trained weights to inspire new, simple initialization techniques. So far, it has been applied only to spatial mixing layers, such convolutional, self-attention, and state space layers. In this work, we present the first attempt to apply the method to channel mixing layers, namely multilayer perceptrons (MLPs). Our extremely simple technique for MLPs -- to give the first layer a nonzero mean -- speeds up training on small-scale vision tasks like CIFAR-10 and ImageNet-1k. Though its effect is much smaller than spatial mixing initializations, it can be used in conjunction with them for an additional positive effect.

Figures (4)

• Figure 1: Full empirical covariance matrices for unrolled MLP weights $W_1, W_2$ (jointly).
• Figure 2: CIFAR-10 experiments for ConvNeXt and ViT in conjunction with previous mimetic inits for convolutional and self-attention layers. The effect of the MLP init is significant in early training, but eventually evens out. The effect of conv. and attn. mimetic init, however, remains. Note: each point in the graph represents the accuracy of a completed training run of $x$ epochs, mean/std. reported over 5 seeds.
• Figure 3: Our MLP init improves early ImageNet-1k training in a data-efficient ViT, and maintains a 0.5% advantage over baseline for the full training time.
• Figure 4: Sweep of the bias parameter $b$ of our init; across several baselines, $b=0$ is noticeably suboptimal. Mean/std. over 5 seeds per $b$.