SmallInitEmb - LayerNorm(SmallInit(Embedding)) in a Transformer to improve convergence

Last update: Dec 25, 2022

Related tags

Overview

SmallInitEmb

LayerNorm(SmallInit(Embedding)) in a Transformer

I find that when training a transformer, the embedding matrix moves slowly, hence it's difficult for the model to jump out of the initial noisy embedding.

(initial embedding)
[[-0.0073  0.0062 -0.0261 ...  0.0086  0.0107 -0.008 ] ... ]
 (after 1 step, the directions of the embedding vectors are not moved much because the numbers change by ~LR = ~4e-4)
[[-0.0069  0.0066 -0.0265 ...  0.009   0.0111 -0.0084] ... ]

So I propose initializing the embedding matrix to tiny values, and put another LayerNorm after it (before all the SA & FFN layers):

if isinstance(module, (nn.Embedding)):
    nn.init.uniform_(module.weight, a=-1e-4, b=1e-4) # SmallInit(Emb)
...
if self.config.USE_SMALL_EMB and self.layer_id == 0:
    x = self.lnPre(x) # LN(SmallInit(Emb))
x = x + self.att(self.ln1(x))
x = x + self.ffn(self.ln2(x))

And then you get improved convergence (especially for BPE models) because the model can quickly jump out of the tiny initial embedding (small changes after 1 step -> significant changes of directions -> significant changes after LayerNorm).

Loss curve comparison: https://wandb.ai/blinkdl/SmallEmbTest

(the gap between LayerNorm(SmallEmb)) and baseline persists after more training)

Moreover, you can directly train PostLN models without warmup with SmallInit(Emb)

if isinstance(module, (nn.Embedding)):
    nn.init.uniform_(module.weight, a=-1e-4, b=1e-4) # SmallInit(Emb)
...
x = self.ln1(x) # this plays the same role as the lnPre in the above PreLN code
x = x + self.att(x)
x = self.ln2(x)
x = x + self.ffn(x)
(note you shall have another LN after the final ffn)

SmallInitEmb - LayerNorm(SmallInit(Embedding)) in a Transformer to improve convergence

Related tags

Overview

SmallInitEmb

Moreover, you can directly train PostLN models without warmup with SmallInit(Emb)

Owner

PENG Bo

AdamW optimizer and cosine learning rate annealing with restarts

CONetV2: Efficient Auto-Channel Size Optimization for CNNs

Unofficial PyTorch implementation of MobileViT.

Perspective: Julia for Biologists

Weakly-Supervised Semantic Segmentation Network with Deep Seeded Region Growing (CVPR 2018).

PyTorch code for the paper "Complementarity is the King: Multi-modal and Multi-grained Hierarchical Semantic Enhancement Network for Cross-modal Retrieval".

Implementation of Learning Gradient Fields for Molecular Conformation Generation (ICML 2021).

Code for the paper: Audio-Visual Scene Analysis with Self-Supervised Multisensory Features

Toolkit for collecting and applying prompts

VolumeGAN - 3D-aware Image Synthesis via Learning Structural and Textural Representations

Breast Cancer Classification Model is applied on a different dataset

Python code for the paper How to scale hyperparameters for quickshift image segmentation

Steerable discovery of neural audio effects

Code for CVPR2019 paper《Unequal Training for Deep Face Recognition with Long Tailed Noisy Data》

Generative Flow Networks

This is the reference implementation for "Coresets via Bilevel Optimization for Continual Learning and Streaming"

ConE: Cone Embeddings for Multi-Hop Reasoning over Knowledge Graphs

Selective Wavelet Attention Learning for Single Image Deraining

Transformers are Graph Neural Networks!

Element selection for functional materials discovery by integrated machine learning of atomic contributions to properties