notes

Log | Files | Refs

kl_regularization.txt (2852B)

      1 # KL Regularization
      2 
      3 KL regularization adds a Kullback-Leibler (KL) divergence penalty to a loss
      4 function. It measures how much one probability distribution diverges from
      5 another. Used as regularization, it pulls a learned distribution toward a
      6 target (usually a simple prior), preventing the model from collapsing to
      7 degenerate or overly complex representations.
      8 
      9 
     10 ## Definition
     11 
     12 For distributions q and p over the same variable z:
     13 
     14   KL(q || p) = E_q[log q(z) - log p(z)]
     15 
     16   - KL >= 0, and KL = 0 only when q = p almost everywhere
     17   - not symmetric: KL(q || p) != KL(p || q)
     18   - penalizes q for placing mass where p has little mass
     19 
     20 
     21 ## In variational autoencoders
     22 
     23 In a VAE, the encoder learns an approximate posterior q(z|x) (e.g. Gaussian
     24 with learned mean and variance). KL regularization matches it to a prior p(z)
     25 (usually N(0, I)):
     26 
     27   L = E_q[log p(x|z)] - KL(q(z|x) || p(z))
     28        ^ reconstruction          ^ KL regularization
     29 
     30 Effects:
     31   - makes the latent space smooth and continuous (nearby z decode to similar x)
     32   - lets you sample z ~ p(z) at inference to generate new data
     33   - prevents the encoder from memorizing x as a unique deterministic code
     34 
     35 Tradeoff: too much KL weight hurts reconstruction; too little yields
     36 unstructured latents or posterior collapse (encoder ignores z, decoder
     37 reconstructs from bias alone).
     38 
     39 beta-VAE scales the KL term by beta > 1 to push stronger regularization and
     40 encourage more disentangled factors (Higgins et al., 2017).
     41 
     42 
     43 ## Closed form (Gaussian case)
     44 
     45 When q(z|x) = N(mu, diag(sigma^2)) and p(z) = N(0, I):
     46 
     47   KL(q || p) = -0.5 * sum_i (1 + log(sigma_i^2) - mu_i^2 - sigma_i^2)
     48 
     49 This is cheap to compute and differentiable, so it is added directly to the
     50 training loss each batch.
     51 
     52 
     53 ## Other common uses
     54 
     55   - RL / policy optimization: KL(policy_new || policy_old) caps how far the
     56     policy moves per update (TRPO, PPO)
     57   - knowledge distillation: KL(student || teacher) transfers soft label
     58     distributions
     59   - variational inference generally: any model with an intractable posterior
     60     can be trained by minimizing KL between an approximate and true posterior
     61 
     62 
     63 ## Key papers
     64 
     65   Kingma & Welling (2013/2014)
     66     "Auto-Encoding Variational Bayes"
     67     arXiv:1312.6114
     68     https://arxiv.org/abs/1312.6114
     69     VAE loss = reconstruction - KL to prior; foundational use of KL as
     70     regularization in deep generative models.
     71 
     72   Higgins et al. (2017)
     73     "beta-VAE: Learning Basic Visual Concepts with a Constrained Variational
     74     Framework"
     75     arXiv:1804.03599
     76     https://arxiv.org/abs/1804.03599
     77     Tunable KL weight (beta) to balance reconstruction vs latent structure.
     78 
     79   Bowman et al. (2016)
     80     "Generating Sentences from a Continuous Space"
     81     arXiv:1511.06349
     82     https://arxiv.org/abs/1511.06349
     83     Early discussion of KL annealing to avoid posterior collapse in
     84     sequence VAEs.