Large Batch Training › Key Parameters

Understanding each hyperparameter and its role helps you make principled decisions when configuring large batch training.

• batch_size: The number of training examples processed per gradient update. Should be chosen to balance hardware utilization and the critical batch size limit. The critical batch size varies by task; empirical measurement (as shown above) gives the best estimate.

• learning_rate ($\eta$): The global learning rate. Must be scaled linearly with batch size: $\eta' = k\eta$ when batch size is multiplied by $k$. The base learning rate $\eta$ should be validated at a small batch size before scaling.

• warmup_epochs: Number of epochs to linearly increase the learning rate from a small initial value to the target. Typically 5 epochs for ImageNet-scale tasks, often 1-3% of total training steps for LLM pre-training. A good rule of thumb is to warm up for at least as many steps as $k$ (the scaling factor).

• weight_decay ($\lambda$): L2 regularization coefficient. Incorporated directly into the LAMB update as a decoupled weight decay term. Decoupled weight decay (AdamW-style, not the original Adam L2 approach) is important for LAMB to behave correctly.

• beta1, beta2: Exponential decay rates for Adam's first and second moment estimates (LAMB). Standard values are 0.9 and 0.999. These generally do not need adjustment when scaling batch size.

• trust_ratio_clip: (Optional) Maximum value of the LARS/LAMB trust ratio to prevent instability for layers near zero weight norm. A common default is to clip trust ratios to the range $[0, 10]$.