Prevent VAE Posterior Collapse with Free Bits

A variational autoencoder can appear to train normally while its latent representation becomes nearly useless. The decoder learns to explain the data without depending on the latent variable, the encoder moves toward the prior, and the KL divergence shrinks toward zero. This failure mode is called posterior collapse.

Free bits is a small change to the VAE objective that can reduce one source of that collapse. It stops the KL term from rewarding the optimizer for squeezing an already-small amount of latent information even closer to zero. The technique is simple, but its name and common shorthand can lead to a misleading mental model. Free bits does not force a latent variable to contain a chosen amount of information. It changes the optimization pressure below a threshold.

This article explains that distinction, works through the objective with a small example, and shows what to measure before deciding whether free bits is actually helping.

Why a VAE can learn to ignore its latent variable

A VAE has two learned parts. The encoder produces an approximate posterior q(z|x) over a latent variable z, and the decoder models the data from a sampled latent value.

A common form of the evidence lower bound (ELBO) is:

ELBO = expected log likelihood - KL(q(z|x) || p(z))

Training usually minimizes the negative ELBO, so the same objective can be written as:

loss = reconstruction loss + KL(q(z|x) || p(z))

The two terms ask for different things. The reconstruction term rewards latent information that helps the decoder explain the input. The KL term regularizes the approximate posterior toward the prior p(z).

If the decoder can model the data well without using z, the model has an easy way to reduce the objective: make q(z|x) close to p(z). Then the KL term approaches zero. When the approximate posterior is effectively independent of x, samples of z no longer tell the decoder much about which input produced them.

This is especially relevant when the decoder is expressive enough to predict much of the structure from its own context. It is not limited to one architecture, and a low KL by itself does not prove that every latent representation is useless. The practical question is whether the latent variable carries information the application needs.

The free-bits mental model: stop rewarding tiny KL values

Suppose the latent space is divided into groups. A group might be one latent dimension, several dimensions, or a larger block chosen by the implementation. Let KL_j be the average KL contribution for group j.

Without free bits, the regularization term is:

KL_total = sum(KL_j)

With a threshold lambda, free bits replaces each contribution with a floor:

KL_free = sum(max(lambda, KL_j))

The training loss becomes:

loss = reconstruction_loss + KL_free

Consider one group whose KL is 0.08 while lambda = 0.20:

ordinary KL penalty: 0.08
free-bits penalty:    0.20

At first glance, that looks like a stronger penalty. The gradient tells the more useful story. While KL_j remains below lambda, max(lambda, KL_j) is constant with respect to that KL value. The objective therefore provides no KL gradient pushing 0.08 down toward zero. The encoder can increase the group’s KL up to the threshold without increasing this part of the objective.

Once the group’s KL rises above the threshold, the normal KL pressure returns:

KL_j = 0.35
lambda = 0.20
penalty = 0.35

That is the central idea: free bits creates a region in which using a small amount of latent capacity does not cost additional KL penalty.

Work through a small latent example

Assume a VAE has four latent groups. After averaging the per-example KL values over the batch, they are:

group 1: 0.03 nats
group 2: 0.11 nats
group 3: 0.27 nats
group 4: 0.64 nats

The ordinary total is:

0.03 + 0.11 + 0.27 + 0.64 = 1.05 nats

Now apply free bits with lambda = 0.20 per group:

group 1: max(0.20, 0.03) = 0.20
group 2: max(0.20, 0.11) = 0.20
group 3: max(0.20, 0.27) = 0.27
group 4: max(0.20, 0.64) = 0.64

total = 1.31 nats

Groups 1 and 2 now sit inside the free region. Their KL values can move upward without changing KL_free until they cross 0.20. Groups 3 and 4 still receive the usual pressure toward the prior because they are already above the threshold.

This example also exposes an easy reporting mistake. The optimized free-bits penalty is 1.31, but the model’s actual raw KL is still 1.05. If you log only the clipped value, you can no longer tell how much KL the approximate posterior really has.

Track both.

Apply the threshold at the intended aggregation level

The formula is short; the reduction order is where implementations often diverge.

For a diagonal Gaussian posterior, code commonly computes a KL contribution for every example and latent dimension. A useful conceptual sequence is:

1. compute KL per example and latent dimension
2. average across the batch
3. optionally combine dimensions into groups
4. apply the free-bits floor to each group
5. sum the floored group values

In pseudocode:

kl_per_example_dim = gaussian_kl(mean, log_variance)
kl_per_dim = mean_over_batch(kl_per_example_dim)
kl_for_loss = sum(max(free_bits, kl_per_dim))
loss = reconstruction_loss + kl_for_loss

This is a teaching sketch, not a framework-specific API. The exact tensor reductions must match the likelihood reduction and the intended definition of a training example.

Why does the order matter? Compare flooring an average with averaging individually floored values. For two examples with KL values 0.0 and 0.6, using lambda = 0.2 gives:

floor after averaging:
max(0.2, (0.0 + 0.6) / 2) = 0.3

floor before averaging:
(max(0.2, 0.0) + max(0.2, 0.6)) / 2 = 0.4

Those are different objectives and produce different gradients. Copying a threshold from another implementation without matching its grouping and reduction convention can therefore change the training recipe substantially.

“Bits” may actually mean nats

The technique is called free bits, but the numeric unit follows the logarithm used in the KL calculation.

If the implementation uses natural logarithms, as many machine-learning objectives do, the KL is measured in nats. A threshold of 0.2 in that objective means 0.2 nats, not 0.2 bits. Converting between them requires the usual logarithm-base conversion.

You rarely need to convert units during ordinary training if the whole implementation is internally consistent. You do need to know the unit when reproducing a threshold from a paper, comparing two codebases, or reporting the amount of KL used by the model.

The group size matters for the same reason. A threshold of 0.2 per dimension is not equivalent to 0.2 for a group of 16 dimensions. Record both the threshold and where it is applied.

Free bits does not guarantee useful representations

A common explanation says free bits “forces each latent dimension to carry information.” That overstates what the objective does.

Below the threshold, the KL contribution becomes flat in the modified objective. The optimizer is no longer rewarded for reducing that KL further, but nothing in the free-bits term itself forces reconstruction gradients to increase it. If the decoder has no reason to use a latent group, its raw KL can remain below the threshold.

There is another subtlety: KL divergence is a useful rate-like diagnostic in a VAE, but a particular KL value does not by itself prove that the latent encodes the semantic factors you care about. A model can spend latent capacity on nuisance details, and different dimensions can be redundant.

So treat free bits as an optimization tool, not as a representation-quality guarantee.

Diagnose posterior collapse before changing the objective

Before adding free bits, log enough information to establish the failure mode. At minimum, separate the reconstruction term from the raw KL term. A single combined loss can hide a KL that has fallen close to zero.

For a factorized latent, per-dimension or per-group KL statistics are also useful. A healthy total can conceal many inactive dimensions, while a low total may be acceptable for a task that genuinely needs little latent information.

Then test whether the decoder depends on z. Depending on the application, useful checks include replacing latent samples with samples from the prior, shuffling latent codes across examples, or measuring downstream quality from the learned representation. If those interventions barely affect the outputs or task metric, the decoder may not be using the latent information in the way you intended.

Do not diagnose collapse from sample quality alone. A strong decoder can produce plausible samples while ignoring z, because unconditional modeling and useful latent representation are different goals.

Choose the threshold as a capacity trade-off

There is no architecture-independent free-bits threshold that is correct for every VAE. The useful value depends on the latent grouping, data, decoder, likelihood scale, optimization recipe, and what you need from the representation.

A threshold that is too small may leave nearly the same pressure toward zero as ordinary ELBO training. A threshold that is too large can weaken the prior-matching pressure over a substantial region, allowing the posterior to use more rate than the application benefits from. That can change reconstruction quality, prior samples, and the geometry of the latent space.

Tune the threshold against measurements that reflect the goal rather than against raw KL alone. If you need controllable latent factors, test those controls. If you need compression, measure the rate-quality trade-off. If you need generation from the prior, inspect or score prior samples as well as reconstructions.

The threshold is part of the objective, so changing it creates a different training run. It should be selected on validation behavior rather than adjusted until a training curve looks aesthetically pleasing.

Compare free bits with nearby interventions

Free bits is one of several ways to change the balance between reconstruction and KL regularization.

KL annealing starts with a reduced KL influence and increases it during training. It changes when the model pays the KL cost. Free bits instead changes the shape of that cost below a threshold. The methods address related optimization pressure but are not interchangeable.

A beta-VAE-style weight multiplies the KL term by a coefficient. With a coefficient below one, every KL value receives weaker pressure; with free bits, the pressure is specifically flattened below the threshold and remains ordinary above it. Their gradients are therefore different even when their total losses happen to match at one point.

Architectural changes attack another cause. If the decoder can solve the task without z, reducing decoder shortcuts or changing how latent information reaches the decoder may be more direct. No objective tweak can guarantee that an architecture will use a latent variable in a useful way.

Start with the simplest intervention that matches the diagnosis. If the model uses z early and then loses it as KL regularization takes hold, an objective-level method is a reasonable experiment. If the decoder ignores z from the start, inspect the architecture and data path as well.

Common implementation mistakes

Several mistakes can make a free-bits experiment hard to interpret:

• Logging only the floored KL. Keep the raw KL separate so you can see whether latent usage actually changed.
• Flooring the wrong quantity. Per-example, per-dimension, per-group, and total-KL floors define different objectives.
• Mixing sum and mean reductions. Changing batch size or sequence length can silently change the relative scale of reconstruction and KL terms if their reductions are inconsistent.
• Assuming the threshold is a guaranteed minimum. The flat region removes downward KL pressure; it does not create upward pressure by itself.
• Copying a number without its units or grouping. A value expressed per group in nats cannot be transferred directly to a per-dimension objective measured differently.

These details are more consequential than the few lines of code needed to implement the clamp.

Use free bits when the diagnosis matches

Free bits is a good candidate when a VAE needs informative latent variables, raw KL collapses toward zero or many latent groups become inactive, and the decoder has reconstruction signal that could benefit from using z. It is especially useful when you want a local change to the KL objective rather than a redesign of the model.

It is less compelling when the latent is already used adequately, when generation quality rather than latent representation is the only goal, or when the real problem is an architecture that gives the decoder no practical reason to depend on z. In those cases, adding another hyperparameter can complicate training without addressing the bottleneck.

Treat latent usage as something to measure

The most useful way to work with free bits is to separate three questions: what objective are you optimizing, how much latent rate is the model actually using, and does that latent information improve the behavior you care about?

Implement the floor at an explicit aggregation level, log both raw and modified KL, and compare against an ordinary-ELBO baseline under the same evaluation. If the raw KL rises but the latent still has little effect on outputs or downstream tasks, the experiment has revealed something useful: the problem is not simply that the KL penalty was too eager to push small values toward zero.