1. buy bigger GPU.

Memory Optimization Techniques​

Optimizer Memory Reduction​

• 8-bit Adam - 2*8-bit vs 2*32-bit = ~75% less memory for optimizer states
• Not storing activations from forward pass but recomputing them in backward pass (memory-compute tradeoff)
• "Stateless" optimizers - SGD, Lion. In contrast to e.g. Adam these don't store additional momentum variables (which takes additional 2x model size of space). The catch is these are not drop-in replacements and Adam was specially made for faster convergence & more stable training.
• Optimizer states offloading to CPU, using torchao:

  optim = CPUOffloadOptimizer(model.parameters(), torch.optim.AdamW, fused=True)
  optim.load_state_dict(ckpt["optim"])
  

Parameter-Efficient Methods​

• LoRA (Low-Rank Adaptation) - thin layer over frozen pretrained layers of your model. It uses a trick of decomposing a large matrix into two smaller low-rank ($n\times m$ where $n << m$) matrices that gives huge memory savings. Slightly more formal:

  $\text{act}[(W + \Delta W)\cdot x + b]$, where:

  • $W=(d \times k)$ pretrained weight tensor
  • $\Delta W = AB$, (init LoRA tensors) where:
    • $A = (d\times \text{rank}) =\mathcal{N}(0, \sigma^2)$ at init
    • $B=(\text{rank} \times k)=0$ at init
    • at the begining $AB$ are no-op but thanks to $A$ being Gaussian there will be symmetry breaking
  • some good links about it: Sebastian Raschka, AI Coffe Break

• QLoRA - It differs in that base model is quantized (usually to 4-bit) but LoRA layers are kept in 16-bit precision

• After training its good to merge LoRA layers with the base model for less latency and computational overhead (but you can't swap if you have more of them).

• DoRA - decomposition of weight matrix into magnitude vector $m$ (euclidean distance) & directional matrix $V$ (angle) and train them separately.

• GaLore - LoRA but for gradient matrix. Supposedly works also for pretraining in contrast to LoRA, but I didn't tried it.

• Prefix Tuning - In place of the prompt you put a random init vector (the so-called prefix) and optimize it until you get the correct answer.

  • ✅ tiny amount of parameters to tune
  • ❌ takes context length (but alternatively you would put pre-prompt there)
  • ❌ interpretability - these are not words, but you can decode it to "nearest" words but it often gibberish

Mixed Precision Strategies​

https://sebastianraschka.com/blog/2023/llm-mixed-precision-copy.html

https://medium.com/@jbensnyder/solving-the-limits-of-mixed-precision-training-231019128b4b

All types other than FP32 are only faster on consumer cards due to Tensor cores which are only available on RTX cards. Tensor cores are used automatically when using mixed precision.

• FP32+FP16/BF16 (speed & memory*) - you store an extra copy of the model in 16-bit to be able to do faster calculations (~2x) (memory-compute trade-off). Gradients are also computed in 16-bit, but stored in FP32. BF16 should be more stable than FP16, as more range is more important for NNs than precision. BF16 stands from "Brain Float" from Google Brain btw.

  $$\begin{align*} \underbrace{(4+2\ \text{bytes})}_{\text{{model FP32+FP16}}} + \underbrace{(2\ \text{bytes})}_{\text{{activation FP16}}} + \underbrace{(4\ \text{bytes})}_{\text{{gradients FP32}}} = 12\ \text{bytes (+8 bytes for Adam)}\end{align*}$$

  *memory savings will be only visible if your batch size is sufficiently large and outweight additional cost of storing copy of weights in 16-bits

• turn on TF32 (Tensor Float) (speed) - not all operations are supported in 16-bit mixed-precision and have to be done in FP32. Turning on TF32 replaces FP32 in computation (storage still in FP32) at speeds similar to FP16. TF32 is supported on Ampere arch and newer. Fun fact is that TF32 is 19-bit format but has "32" in name.

  # The flag below controls whether to allow TF32 on matmul.
  # This flag defaults to False in PyTorch 1.12 and later.
  torch.backends.cuda.matmul.allow_tf32 = True
  
  # The flag below controls whether to allow TF32 on cuDNN. 
  # This flag defaults to True.
  torch.backends.cudnn.allow_tf32 = True
  

• FP32+FP8 (speed & memory) - there is possibility to do computations in 8-bit with FP32 accumulation but its more complicated to setup than AMP natively supported in PyTorch. Use HF Accelerate library to do this (from what I understand its a wrapper on 3 other packages [TransformersEngine, MS-AMP, and torchao] but not only this).

  $$\begin{align*} \underbrace{(1\ \text{byte})}_{\text{{model FP8 E4M3}}} + \underbrace{(1\ \text{byte})}_{\text{{activation FP8 E4M3}}} + \underbrace{(1\ \text{byte})}_{\text{{gradients FP8 E5M2}}} = 3\ \text{bytes (+8 bytes for Adam)}\end{align*}$$

Computational Efficiency​

Forward/Backward Pass Optimization​

• set gradients to None instead of default 0 (but this can cause some unexpected behaviors - None isn't a number so operations with it produces NaN)
• non-global Cross-Entropy calculation reduces memory usage spike at the end (especially beneficial for LLMs). paper

Hardware Utilization​

• avoid moving tensors to another device .to(device), create tensors directly on target device instead. If you don't have any synchronization later in code then you can use .to(non_blocking=True)
• use torch.compile() if it works, for me it usualy don't.
• torch.backends.cudnn.benchmark = True

Training Dynamics​

Initialization & Learning Rate​

• init weights properly (relaying on default pytorch init isn't always optimal)
• LR scheduler (OneCycleLR, lr warmup, etc.) and lr search. Similar for BS (batch size warmup etc.)
• some rule regarding relation of LR & BS - $LR \approx \sqrt{BS}$ (there is also issue of "critical batch size")

Batch Processing​

• max out batch size to fill whole VRAM (batch size finder)
• if you can't fit enough batch size for stable training (e.g. bsz=1) then use gradient accumulation. In pure pytorch it boils down to calling opt.step() less frequently what results in effectively higher batch size.

Stability Techniques​

• gradient clipping - prevents exploding gradients by capping their magnitude during backpropagation.
• weights clipping - regularization technique, similar to weight decay in some sense but this isn't connected with loss calc.
• start from pretrained model (transfer learning) and swap last layer.
  • Freeze whole model except last layer and after few epochs gradually unfreeze rest of the layers (ULMFiT, gradual unfreezing).
• constraint latent space of your model to align with some general pretrained model, stabilizes training and reduces overfitting (can be done by feature matching losses).

Model-Specific Optimizations​

LLM Techniques​

• SuperBPE - groups frequent word sequences into single tokens, improving efficiency and performance. Common word combinations get treated as one unit by the tokenizer, which reduces the number of easy-to-predict sequences. This creates a more balanced prediction difficulty across tokens, allowing the model to distribute computational effort more effectively. Author explaination

Diffusion Model Techniques​

• latent diffusion - diffuse in latent space, not pixel space. VAE encoder $\rightarrow$ latent (diffuse $N\times$) $\rightarrow$ VAE decoder. SD1.5 VAE maybe big but you can use TAESD ($(3\cdot512\cdot512)/(16\cdot64\cdot64)=12\times$ compression, 5MB for enc/dec each and minimal computational overhead)
  • if you want train on ImageNet: https://huggingface.co/datasets/fal/cosmos-imagenet (compressed to 2.45GB)
• Min-SNR - method of adding a weightning to the loss based on the SNR (signal to noise ratio) of the timestep. It prevents conflicting gradients from different deniosing phases (beggining, mid and final refinements)

Implementation Tricks​

PyTorch Data Handling​

• use .as_tensor() rather than .tensor(). torch.tensor() always copies data. If you have a numpy array that you want to convert, use torch.as_tensor() or torch.from_numpy() to avoid copying the data.
• try "channel last" format for tensors and model (NCHW => NHWC), sometimes it's faster. link
• slow dataloader optimizations Simo tweet, PyTorch forum

Other​

• normalization layers (stabilize and speeding up training - you can use higher LR)
• turn off bias before BatchNorm (bn already does shift)
• MoE - model architecture which selectively activates only part of the model. memory-computation tradeoff. MoE faster achives same loss under the same computational budged compared to dense models.
• MoD (Mixture of Depths) - learned skip connection for each transformer block, model learns to not waste compute on easy tokens.
• Stochastic Depth - each layer in deep ConvNet have probability of not being dropped from 1.0 (for first layer) to 0.5 (for last layer). Simply dropout whole layers (prevents vanishing gradients, faster training, better performance)
• checkpoint averaging - weighted average of previous checkpoints makes loss landscape more smooth and convex which speeds up training + reduces overfitting (applies to pretraining & finetuning)

(Vanilla) GAN​

In short:

• Generator $G$ - tries to generate realistic looking images (maximize $D(G(z))$).
• Discriminator $D$ - tries to distinguish between real and generated images (minimize $D(G(z))$ and maximize $D(x)$). Discriminator makes for learned loss function for generator.
• The original GAN uses binary cross-entropy loss.
• Discriminator has softmax at the end and outputs values between $[0, 1]$ (0 - fake, 1 - real).
• Best for Generator is for the Discriminator to output ~0.5, which means that D can't tell if an image is real or fake (Nash equilibrium).
• Loss isn't interpretable because both the generator and discriminator can improve in the same time, so even if G's image quality improves, the loss can stay the same.

DCGAN​

The simplest GAN that doesn't use FFN layers but Conv layers, takes a noise vector and gradually increases it through TransposeConv2d.

• Use conv layers instead of FFN layers.
• Use conv with a larger stride instead of pooling.
• Use batchnorm.
• ReLU in $G$ and LeakyReLU in $D$.
• $\beta_1=0.5$ (PyTorch default is $0.9$)

cGAN - Conditional GAN​

Refers to the general idea of adding class information to the GAN. In the context of the Generator, this usually involves adding labels to the noise (concatenating two vectors). The Discriminator checks if the image is real or fake, having additional information about the class. (The Discriminator returns a scalar for each image).

# Adding conditioning in the Generator
def forward(x, y):
	cond = self.embedding(y)
	x = torch.cat([x, cond], dim=1) # x.shape = [bs, z_dim], cond.shape = [bs, embed_dim]
	return self.backbone(x)

# Adding conditioning in the Discriminator
def forward(x, y):
	out = self.net(x)
  cond = self.embed(y)
  return (x*y).sum(dim=[1, 2, 3])

acGAN (Auxiliary Classifier GAN)​

In acGAN, the discriminator has two outputs: one for real/fake (as in cGAN) and one for classifying the class. The Discriminator doesn't receive classes as input because it has to predict them.

# Discriminator (Generator remains the same as in cGAN)
def forward(self, img):
  features = self.model(img)
  validity = self.real_fake_head(features)
  class_logits = self.class_head(features)
  return validity, class_logits
  
loss_fn_bce = nn.BCELoss()  # For real/fake
loss_fn_ce = nn.CrossEntropyLoss()  # For classification

for i in range(epochs):
	# ...
	# Discriminator
	d_loss_real = loss_fn_bce(real_validity, torch.ones_like(real_validity)) + \\
	              loss_fn_ce(real_class_logits, real_labels)
	d_loss_fake = loss_fn_bce(fake_validity, torch.zeros_like(fake_validity)) + \\
	              loss_fn_ce(fake_class_logits, labels)
	d_loss = d_loss_real + d_loss_fake
	# ...
	# Generator
	g_loss_validity = loss_fn_bce(fake_validity, torch.ones_like(fake_validity))
	g_loss_class = loss_fn_ce(fake_class_logits, labels)
	g_loss = g_loss_validity + g_loss_class
  # ...

WGAN & WGAN-GP (Wasserstein GAN with Gradient Penalty)​

Change of JS divergence (from Vanilla GAN) to Earth Mover Distance (aka Wasserstein distance).

• JS divergence works badly if two distributions have small or no overlap (mode collapse, vanishing grads).
• Wasserstein distance measures the "cost" of transforming one distribution into another, works even if there is no overlap between distributions, and gives meaningful values (decreasing loss of the discriminator says that Generator improves in quality).

In the original WGAN (without GP), weights are clipped (not the gradient) to ensure Lipschitz continuity. The parameter $c$ decides how strong the clipping is.

The purpose of the gradient penalty is to keep the critic in the space of 1-Lipschitz functions, so that the critic provides high-quality gradients and the Wasserstein distance is estimated better. (GP only for the critic).

LSGAN (Least Squares GAN)​

Using L2 loss instead of BCE loss to calculate the loss:

Progressive Growing GAN​

The Generator is gradually trained, first at a low resolution, and a new layer that increases the resolution is added periodically. Advantages:

• Training stability: Starting with small images causes the model to learn the most important features of the image (shapes, colors) at the beginning, rather than textures and background details. This reduces the risk of training derailment.
• Training speed: Training on small images is much cheaper, and the model quickly adapts to a higher resolution.
• Better image quality.

Pix2Pix​

An architecture created for img2img translation (e.g., sketch → image). The most important features are the use of adversarial loss + L1 loss (they didn't use L2 because of lower outlier resistance) and the introduction of PatchGAN (image discrimination with tiles instead of a scalar for the entire image - more detailed/local feedback for G).

Loss

• L1 for adhering to the outline but blurred (low-freq structure).
• Adv. loss is responsible for sharp edges (high-freq structure).

Model structure and other details

• UNet for the generator.
• PatchGAN - discrimination with tiles.
• Block structure: Conv + BatchNorm + ReLU.
• Dropout as a form of result diversity (enabled during inference; no noise vector as a source of randomness and diverse outputs).
• LR for D is 2x smaller than for G.
• Adam with b1=0.5 and b2=0.999.

CycleGAN​

The CycleGAN architecture was made for style-transfer without needing special data for mapping $(A \rightarrow B)$. To achieve this, it uses a combination of 3 losses:

1. Cycle-consistency loss

Transforms an image into something and back, and compares the difference between the original and the generated one (the circle/cycle closes).

2. Adversarial loss

The usual loss from the discriminator (each G has its own D).

3. Identity loss

The Generator sometimes receives an image that it should generate, so it shouldn't change anything in it, just return it (a form of regularization).

Features

• 2 generators and 2 discriminators.
• Doesn't need a dataset consisting of image pairs (it can be 2 arbitrary datasets from different domains, e.g., photos and paintings).
• Generators as autoencoders, not UNets (unlike Pix2Pix).
• Uses PatchGAN from Pix2Pix in the generator.

R3GAN​

Summary

Authors developed a new loss function that has mathematical guarantees of convergence and is stable. With stable training achieved, they modify the architecture (starting from StyleGAN2), remove all the now-superfluous tricks, and modernize the architecture. Ultimately, they arrive at a stable and simpler GAN architecture that rivals previous GANs and diffusion-based SOTA models.

Details

1. Re-formulated Loss: Instead of calculating losses for D and G separately, they are combined.

   • Traditional GAN Loss:

     $L(\theta, \psi) = \mathbb{E}_{z \sim p_z} \left[ f \left( D_\psi(G_\theta(z)) \right) \right] + \mathbb{E}_{x \sim p_D} \left[ f \left( -D_\psi(x) \right) \right]$ G and D are calculated separately, causing real and fake samples to be far apart (D pushes them apart).

   • "Relativistic" GAN Loss:

     $L(\theta, \psi) = \mathbb{E}_{z \sim p_z, \ x \sim p_D} \left[ f \left( D_\psi(G_\theta(z)) - D_\psi(x) \right) \right]$ Here, G and D are linked. D evaluates the "realness" relative to real samples, so fake samples should always be in the vicinity of real samples.

2. Addition of Zero-Centered Gradient Penalties (R1 and R2):

   • R1: Penalizes the gradient magnitudes of D on real samples.
   • R2: Penalizes the gradient magnitudes of D on fake samples.

   The goal is to reduce the gradients for G when it's close to generating real samples (eliminating training oscillations). When $p_{fake} = p_{real}$, the gradients from D should be 0.

3. Simplification and Modernization of the Architecture: They start with StyleGAN2 and remove everything that was added to increase stability and convergence. Then, they change the architecture to ResNeXt, but without normalization, because they use Fixup initialization, which doesn't require it.

StyleGAN​

A complex architecture that has many tricks to achieve stable training and good quality. Allows for interpolation between any two images. Many versions of StyleGAN have been released. In a nutshell:

1. The $z$ noise vector is transformed by the Mapping network into the vector $w$ (latent code).
2. The Synthesis network starts creating an image from a learnable tensor 4x4 (Const 4x4x512).
3. The gradually upsampled feature map passes through subsequent blocks into which the latent code $w$ (through AdaIN - a type of normalization) and the noise vector $B$ are injected.

Injecting noise through $B$ adds high-frequency details to the image (left with noise, right without noise).

Other stuff (tips, tricks, minor papers)​

Gated Shortcut (arXiv:2201.11351)​

The Gated shortcut decides what remains from the residual stream and what is added from the features. Used only in the Generator. $f_i$ - input $f_c$ - feature (what's on the output of conv, bn, relu) $f_o$ - output $*$ - conv, $\odot$ - concat, $\otimes$ - element-wise multiplication

1. $f_g = \text{sigmoid}(\text{conv}(f_c \odot f_i))$ - gate, value 0-1
2. $f_r = \text{conv}(f_c \odot f_i)$ - refinement, learned transformation of combined features; what to add
3. $f_o = \text{conv}[f_g \otimes f_c + (1-f_g) \otimes f_r]$ here: $f_g \otimes f_c$ means how much to take from the feature based on $f_i$ and $f_c$. $(1-f_g) \otimes f_r$ means the rest, which comes from the combination of $f_i$ and $f_c$.

Authors claims that it's better than other gated residuals because it can better add/remove data from the residual stream/features ($f_{EGS}$ is other way to implement a gated shortcut): $f_{EGS} = f_g \otimes f_i + (1-f_g) \otimes f_c$

In Eq. 6, $f_{EGS}$ can be interpreted as the weighted summation of the $f_i$ and $f_c$, where weight value is $f_g$. Thus, if $f_g$ is 0.5, it is equal to the scaled-identity shortcut; it cannot effectively keep (or remove) the relevant (or irrelevant) information in $f_c$. In contrast, instead of directly summing $f_i$, the proposed method produces the refinement feature, i.e. $f_r$.

This can probably be used in other architectures. The gating mechanism doesn't have to be useful only for GANs (e.g., LSTM).

Checkerboard Artifacts​

Transposed conv with overlapping stride can cause checkerboard artifacts - solution: use regular Conv2d + Upsampling2d("nearest"); then conv can overlap.