Autoencoders and Variational Autoencoders

Week 12 of 18 · Generative and self-supervised · 7 days

Week 12 · Generative and self-supervised

Autoencoders and Variational Autoencoders

Week 12 · Day 1 of 7

Learning Without Labels

Rebuilding the input through a bottleneck, and why the bottleneck matters

Everything so far has needed labels. An autoencoder needs none: it is trained to reproduce its own input through a narrow middle, and what it learns on the way is a compressed representation of the data.

The idea

Autoencoder: An encoder that maps the input to a vector narrower than the input, and a decoder that maps that vector back. Because the middle cannot hold everything, training forces it to keep whatever structure is most common and discard the rest.

import torch
from torch import nn
from torch.utils.data import DataLoader, Subset
from torchvision import datasets, transforms

tf = transforms.ToTensor()
train_all = datasets.MNIST('data', train=True, download=True, transform=tf)
test_all = datasets.MNIST('data', train=False, download=True, transform=tf)
train_set = Subset(train_all, range(8000))
val_set = Subset(test_all, range(2000))
train_loader = DataLoader(train_set, batch_size=128, shuffle=True)
val_loader = DataLoader(val_set, batch_size=512)

RAMP = ' .:-=+*#%@'
def show(img, step=2):
    """Print a 28 by 28 image as characters."""
    t = img.detach().reshape(28, 28)
    t = (t - t.min()) / (t.max() - t.min() + 1e-9)
    for row in t[::step]:
        print(''.join(RAMP[min(9, int(v * 9.999))] for v in row.tolist()))
class AutoEncoder(nn.Module):
    def __init__(self, latent=16):
        super().__init__()
        self.encoder = nn.Sequential(
            nn.Flatten(),
            nn.Linear(784, 256), nn.ReLU(),
            nn.Linear(256, 64), nn.ReLU(),
            nn.Linear(64, latent))
        self.decoder = nn.Sequential(
            nn.Linear(latent, 64), nn.ReLU(),
            nn.Linear(64, 256), nn.ReLU(),
            nn.Linear(256, 784), nn.Sigmoid())

    def forward(self, x):
        return self.decoder(self.encoder(x))

def train_ae(model, epochs=8, lr=1e-3, noise=0.0, loader=None):
    torch.manual_seed(0)
    opt = torch.optim.AdamW(model.parameters(), lr=lr)
    loss_fn = nn.MSELoss()
    for _ in range(epochs):
        model.train()
        for xb, _ in (loader or train_loader):
            target = xb.flatten(1)
            inp = xb if noise == 0 else (xb + noise * torch.randn_like(xb)
                                         ).clamp(0, 1)
            opt.zero_grad()
            loss_fn(model(inp), target).backward()
            opt.step()
    model.eval()
    total = seen = 0
    with torch.no_grad():
        for xb, _ in val_loader:
            total += loss_fn(model(xb), xb.flatten(1)).item() * len(xb)
            seen += len(xb)
    return total / seen
torch.manual_seed(0)
model = AutoEncoder(latent=16)
x = next(iter(val_loader))[0][:4]
code = model.encoder(x)
print('input   ', tuple(x.shape), '= %d numbers per image' % (28 * 28))
print('code    ', tuple(code.shape), '= 16 numbers per image')
print('output  ', tuple(model(x).shape))
print('\ncompression %.0f to 1' % (784 / 16))
print('parameters %d' % sum(p.numel() for p in model.parameters()))

input    (4, 1, 28, 28) = 784 numbers per image
code     (4, 16) = 16 numbers per image
output   (4, 784)

compression 49 to 1
parameters 437664

Training it, with no labels anywhere

import torch
from torch import nn
from torch.utils.data import DataLoader, Subset
from torchvision import datasets, transforms

tf = transforms.ToTensor()
train_all = datasets.MNIST('data', train=True, download=True, transform=tf)
test_all = datasets.MNIST('data', train=False, download=True, transform=tf)
train_set = Subset(train_all, range(8000))
val_set = Subset(test_all, range(2000))
train_loader = DataLoader(train_set, batch_size=128, shuffle=True)
val_loader = DataLoader(val_set, batch_size=512)

RAMP = ' .:-=+*#%@'
def show(img, step=2):
    """Print a 28 by 28 image as characters."""
    t = img.detach().reshape(28, 28)
    t = (t - t.min()) / (t.max() - t.min() + 1e-9)
    for row in t[::step]:
        print(''.join(RAMP[min(9, int(v * 9.999))] for v in row.tolist()))
class AutoEncoder(nn.Module):
    def __init__(self, latent=16):
        super().__init__()
        self.encoder = nn.Sequential(
            nn.Flatten(),
            nn.Linear(784, 256), nn.ReLU(),
            nn.Linear(256, 64), nn.ReLU(),
            nn.Linear(64, latent))
        self.decoder = nn.Sequential(
            nn.Linear(latent, 64), nn.ReLU(),
            nn.Linear(64, 256), nn.ReLU(),
            nn.Linear(256, 784), nn.Sigmoid())

    def forward(self, x):
        return self.decoder(self.encoder(x))

def train_ae(model, epochs=8, lr=1e-3, noise=0.0, loader=None):
    torch.manual_seed(0)
    opt = torch.optim.AdamW(model.parameters(), lr=lr)
    loss_fn = nn.MSELoss()
    for _ in range(epochs):
        model.train()
        for xb, _ in (loader or train_loader):
            target = xb.flatten(1)
            inp = xb if noise == 0 else (xb + noise * torch.randn_like(xb)
                                         ).clamp(0, 1)
            opt.zero_grad()
            loss_fn(model(inp), target).backward()
            opt.step()
    model.eval()
    total = seen = 0
    with torch.no_grad():
        for xb, _ in val_loader:
            total += loss_fn(model(xb), xb.flatten(1)).item() * len(xb)
            seen += len(xb)
    return total / seen
torch.manual_seed(0)
model = AutoEncoder(latent=16)
print('validation reconstruction error %.5f' % train_ae(model))

x, _ = next(iter(val_loader))
with torch.no_grad():
    out = model(x[:1])
print('\noriginal:')
show(x[0])
print('\nrebuilt from 16 numbers:')
show(out[0])

validation reconstruction error 0.02912

original:
                            
                            
                            
                            
      %@@@@@########*:      
            : :::: @@=      
                  @@-       
                :@@:        
                %@:         
              :@@:          
             #@%.           
           .%@=             
          :@@@:             
          =@%               

rebuilt from 16 numbers:
                            
                            
                            
                            
    .:-+*##+#*##%@%%*:      
    .::=**+==:-:-+%@%:      
        ..      .+#%=       
               .=%#-        
              .+%%-         
              *@%.          
            .%@@:           
           -*#*:            
         .=*#*-             
                            

Note what the loss function never saw

The target is the input. There is no label in that training loop at all, which means it will run on any pile of unlabelled data you have. That is the whole appeal, and it is the thread running through this week and the next two: labels are expensive, raw data is not, and a great deal can be learned from raw data alone.

How narrow can the middle be?

import torch
from torch import nn
from torch.utils.data import DataLoader, Subset
from torchvision import datasets, transforms

tf = transforms.ToTensor()
train_all = datasets.MNIST('data', train=True, download=True, transform=tf)
test_all = datasets.MNIST('data', train=False, download=True, transform=tf)
train_set = Subset(train_all, range(8000))
val_set = Subset(test_all, range(2000))
train_loader = DataLoader(train_set, batch_size=128, shuffle=True)
val_loader = DataLoader(val_set, batch_size=512)

RAMP = ' .:-=+*#%@'
def show(img, step=2):
    """Print a 28 by 28 image as characters."""
    t = img.detach().reshape(28, 28)
    t = (t - t.min()) / (t.max() - t.min() + 1e-9)
    for row in t[::step]:
        print(''.join(RAMP[min(9, int(v * 9.999))] for v in row.tolist()))
class AutoEncoder(nn.Module):
    def __init__(self, latent=16):
        super().__init__()
        self.encoder = nn.Sequential(
            nn.Flatten(),
            nn.Linear(784, 256), nn.ReLU(),
            nn.Linear(256, 64), nn.ReLU(),
            nn.Linear(64, latent))
        self.decoder = nn.Sequential(
            nn.Linear(latent, 64), nn.ReLU(),
            nn.Linear(64, 256), nn.ReLU(),
            nn.Linear(256, 784), nn.Sigmoid())

    def forward(self, x):
        return self.decoder(self.encoder(x))

def train_ae(model, epochs=8, lr=1e-3, noise=0.0, loader=None):
    torch.manual_seed(0)
    opt = torch.optim.AdamW(model.parameters(), lr=lr)
    loss_fn = nn.MSELoss()
    for _ in range(epochs):
        model.train()
        for xb, _ in (loader or train_loader):
            target = xb.flatten(1)
            inp = xb if noise == 0 else (xb + noise * torch.randn_like(xb)
                                         ).clamp(0, 1)
            opt.zero_grad()
            loss_fn(model(inp), target).backward()
            opt.step()
    model.eval()
    total = seen = 0
    with torch.no_grad():
        for xb, _ in val_loader:
            total += loss_fn(model(xb), xb.flatten(1)).item() * len(xb)
            seen += len(xb)
    return total / seen
print('%8s %10s %16s' % ('latent', 'params', 'reconstruction'))
for latent in [2, 8, 32, 128]:
    torch.manual_seed(0)
    model = AutoEncoder(latent=latent)
    err = train_ae(model, epochs=6)
    print('%8d %10d %16.5f'
          % (latent, sum(p.numel() for p in model.parameters()), err))

  latent     params   reconstruction
       2     435858          0.05140
       8     436632          0.03442
      32     439728          0.03465
     128     452112          0.03403

The error falls sharply from two dimensions to eight and then stops moving. Below eight the bottleneck is genuinely the constraint; above it, something else is, and on this architecture and this budget the extra width buys nothing. Finding that knee is worth doing once, because everything past it is capacity you are paying for and not using.

A bottleneck that is not a bottleneck teaches nothing

Widen the middle until it is as large as the input and the model can learn the identity function, which reconstructs perfectly and has learned nothing whatsoever about the data. The constraint is not an inconvenience to be minimised, it is the entire mechanism. If your reconstruction error is suspiciously low, check that the code is actually narrower than the input once you count every channel.

Day 1 takeaway

An autoencoder learns by rebuilding its own input through a bottleneck, using no labels at all. The width of the bottleneck controls how much it must throw away, and a bottleneck that is not narrow teaches the model nothing.

Week 12 · Day 2 of 7

Denoising, Convolutions and Anomalies

Three variations on one loop, each solving a different problem

By Pawan Rama Mali1203 words

Three variations, each a small change to the same loop, and each solving a different problem.

Denoising

import torch
from torch import nn
from torch.utils.data import DataLoader, Subset
from torchvision import datasets, transforms

tf = transforms.ToTensor()
train_all = datasets.MNIST('data', train=True, download=True, transform=tf)
test_all = datasets.MNIST('data', train=False, download=True, transform=tf)
train_set = Subset(train_all, range(8000))
val_set = Subset(test_all, range(2000))
train_loader = DataLoader(train_set, batch_size=128, shuffle=True)
val_loader = DataLoader(val_set, batch_size=512)

RAMP = ' .:-=+*#%@'
def show(img, step=2):
    """Print a 28 by 28 image as characters."""
    t = img.detach().reshape(28, 28)
    t = (t - t.min()) / (t.max() - t.min() + 1e-9)
    for row in t[::step]:
        print(''.join(RAMP[min(9, int(v * 9.999))] for v in row.tolist()))
class AutoEncoder(nn.Module):
    def __init__(self, latent=16):
        super().__init__()
        self.encoder = nn.Sequential(
            nn.Flatten(),
            nn.Linear(784, 256), nn.ReLU(),
            nn.Linear(256, 64), nn.ReLU(),
            nn.Linear(64, latent))
        self.decoder = nn.Sequential(
            nn.Linear(latent, 64), nn.ReLU(),
            nn.Linear(64, 256), nn.ReLU(),
            nn.Linear(256, 784), nn.Sigmoid())

    def forward(self, x):
        return self.decoder(self.encoder(x))

def train_ae(model, epochs=8, lr=1e-3, noise=0.0, loader=None):
    torch.manual_seed(0)
    opt = torch.optim.AdamW(model.parameters(), lr=lr)
    loss_fn = nn.MSELoss()
    for _ in range(epochs):
        model.train()
        for xb, _ in (loader or train_loader):
            target = xb.flatten(1)
            inp = xb if noise == 0 else (xb + noise * torch.randn_like(xb)
                                         ).clamp(0, 1)
            opt.zero_grad()
            loss_fn(model(inp), target).backward()
            opt.step()
    model.eval()
    total = seen = 0
    with torch.no_grad():
        for xb, _ in val_loader:
            total += loss_fn(model(xb), xb.flatten(1)).item() * len(xb)
            seen += len(xb)
    return total / seen
torch.manual_seed(0)
model = AutoEncoder(latent=32)
# input: a corrupted image. target: the clean one.
train_ae(model, epochs=8, noise=0.4)

x, _ = next(iter(val_loader))
torch.manual_seed(1)
noisy = (x[:1] + 0.4 * torch.randn_like(x[:1])).clamp(0, 1)
with torch.no_grad():
    cleaned = model(noisy)
print('noisy input:')
show(noisy[0])
print('\ndenoised:')
show(cleaned[0])
print('\nclean original:')
show(x[0])

noisy input:
         .   * .  : -  -=   
 -  *         *       .    :
- +- .: @  # # .: .=  :   ::
.*=     -=:*    *  + = - -# 
*   =+@@@@@-@ =%%@@+-.  .  =
-   .=.   ..    *:+@@@ :=  @
    :   ..++ :   =@:  * -  -
 #      .        @@  .   . :
 ... -# -  * .  @@=-=  :    
=::: .. :+   ::=@#  ..  . :+
  : =-     .-@@@ :.. = + -+ 
     +      =+@. @-= %.-= = 
 =:  :    *@%@ :  -.     .  
 *  :*  * @@% -=:    .  = :#

denoised:
                            
                            
                            
              ...           
     ::=++*%@%@%@%%#-.      
     .:==*##+=-=*%%#-.      
        .::.   :*%#=.       
          ....:+%@+         
              =%%+.         
             =%%+:          
            :##+-           
          :=*#*+:           
         .-=+#+:            
          ..:-:             

clean original:
                            
                            
                            
                            
      %@@@@@########*:      
            : :::: @@=      
                  @@-       
                :@@:        
                %@:         
              :@@:          
             #@%.           
           .%@=             
          :@@@:             
          =@%               

One line changed, and the model is now doing something else

The architecture is identical. The only difference is what sits on each side of the loss: corrupted input, clean target. That substitution is the basis of a good deal of modern self-supervised learning. Hide part of the input, train the model to restore it, and the representation you get for free is useful for tasks you have not thought of yet. Week 14 builds on exactly this.

Convolutional

import torch
from torch import nn
from torch.utils.data import DataLoader, Subset
from torchvision import datasets, transforms

tf = transforms.ToTensor()
train_all = datasets.MNIST('data', train=True, download=True, transform=tf)
test_all = datasets.MNIST('data', train=False, download=True, transform=tf)
train_set = Subset(train_all, range(8000))
val_set = Subset(test_all, range(2000))
train_loader = DataLoader(train_set, batch_size=128, shuffle=True)
val_loader = DataLoader(val_set, batch_size=512)

RAMP = ' .:-=+*#%@'
def show(img, step=2):
    """Print a 28 by 28 image as characters."""
    t = img.detach().reshape(28, 28)
    t = (t - t.min()) / (t.max() - t.min() + 1e-9)
    for row in t[::step]:
        print(''.join(RAMP[min(9, int(v * 9.999))] for v in row.tolist()))
class ConvAutoEncoder(nn.Module):
    def __init__(self):
        super().__init__()
        self.encoder = nn.Sequential(
            nn.Conv2d(1, 16, 3, stride=2, padding=1), nn.ReLU(),  # 14
            nn.Conv2d(16, 32, 3, stride=2, padding=1), nn.ReLU())  # 7
        self.decoder = nn.Sequential(
            nn.ConvTranspose2d(32, 16, 3, stride=2, padding=1,
                               output_padding=1), nn.ReLU(),        # 14
            nn.ConvTranspose2d(16, 1, 3, stride=2, padding=1,
                               output_padding=1), nn.Sigmoid())     # 28

    def forward(self, x):
        return self.decoder(self.encoder(x))

torch.manual_seed(0)
model = ConvAutoEncoder()
x, _ = next(iter(val_loader))
print('input ', tuple(x[:2].shape))
print('code  ', tuple(model.encoder(x[:2]).shape))
print('output', tuple(model(x[:2]).shape))
print('parameters %d' % sum(p.numel() for p in model.parameters()))

opt = torch.optim.AdamW(model.parameters(), lr=1e-3)
loss_fn = nn.MSELoss()
for _ in range(6):
    model.train()
    for xb, _ in train_loader:
        opt.zero_grad()
        loss_fn(model(xb), xb).backward()
        opt.step()
model.eval()
with torch.no_grad():
    err = sum(loss_fn(model(xb), xb).item() * len(xb)
              for xb, _ in val_loader) / len(val_set)
print('\nreconstruction error %.5f' % err)

input  (2, 1, 28, 28)
code   (2, 32, 7, 7)
output (2, 1, 28, 28)
parameters 9569

reconstruction error 0.00229

ConvTranspose2d is the upsampling counterpart of a strided convolution. The output_padding argument exists because a stride of 2 maps several input sizes to the same output size, so the reverse operation is ambiguous and you have to say which one you meant. Getting it wrong is the usual reason a decoder produces a 27 by 27 image.

Anomaly detection from reconstruction error

import torch
from torch import nn
from torch.utils.data import DataLoader, Subset
from torchvision import datasets, transforms

tf = transforms.ToTensor()
train_all = datasets.MNIST('data', train=True, download=True, transform=tf)
test_all = datasets.MNIST('data', train=False, download=True, transform=tf)
train_set = Subset(train_all, range(8000))
val_set = Subset(test_all, range(2000))
train_loader = DataLoader(train_set, batch_size=128, shuffle=True)
val_loader = DataLoader(val_set, batch_size=512)

RAMP = ' .:-=+*#%@'
def show(img, step=2):
    """Print a 28 by 28 image as characters."""
    t = img.detach().reshape(28, 28)
    t = (t - t.min()) / (t.max() - t.min() + 1e-9)
    for row in t[::step]:
        print(''.join(RAMP[min(9, int(v * 9.999))] for v in row.tolist()))
class AutoEncoder(nn.Module):
    def __init__(self, latent=16):
        super().__init__()
        self.encoder = nn.Sequential(
            nn.Flatten(),
            nn.Linear(784, 256), nn.ReLU(),
            nn.Linear(256, 64), nn.ReLU(),
            nn.Linear(64, latent))
        self.decoder = nn.Sequential(
            nn.Linear(latent, 64), nn.ReLU(),
            nn.Linear(64, 256), nn.ReLU(),
            nn.Linear(256, 784), nn.Sigmoid())

    def forward(self, x):
        return self.decoder(self.encoder(x))

def train_ae(model, epochs=8, lr=1e-3, noise=0.0, loader=None):
    torch.manual_seed(0)
    opt = torch.optim.AdamW(model.parameters(), lr=lr)
    loss_fn = nn.MSELoss()
    for _ in range(epochs):
        model.train()
        for xb, _ in (loader or train_loader):
            target = xb.flatten(1)
            inp = xb if noise == 0 else (xb + noise * torch.randn_like(xb)
                                         ).clamp(0, 1)
            opt.zero_grad()
            loss_fn(model(inp), target).backward()
            opt.step()
    model.eval()
    total = seen = 0
    with torch.no_grad():
        for xb, _ in val_loader:
            total += loss_fn(model(xb), xb.flatten(1)).item() * len(xb)
            seen += len(xb)
    return total / seen
torch.manual_seed(0)
# Train on fours and nines only, then show it every digit.
targets = train_all.targets[:8000]
keep = ((targets == 4) | (targets == 9)).nonzero().flatten().tolist()
narrow = Subset(train_all, keep)
print('training on %d images of fours and nines only' % len(narrow))

model = AutoEncoder(latent=8)
train_ae(model, epochs=10,
         loader=DataLoader(narrow, batch_size=128, shuffle=True))

model.eval()
errors = {}
with torch.no_grad():
    for xb, yb in val_loader:
        per = ((model(xb) - xb.flatten(1)) ** 2).mean(1)
        for digit, e in zip(yb.tolist(), per.tolist()):
            errors.setdefault(digit, []).append(e)

print('\n%8s %14s %10s' % ('digit', 'mean error', 'seen?'))
for digit in range(10):
    e = sum(errors[digit]) / len(errors[digit])
    print('%8d %14.5f %10s'
          % (digit, e, 'yes' if digit in (4, 9) else 'no'))

training on 1594 images of fours and nines only

   digit     mean error      seen?
       0        0.10658         no
       1        0.05656         no
       2        0.09327         no
       3        0.08882         no
       4        0.04909        yes
       5        0.07873         no
       6        0.07898         no
       7        0.05782         no
       8        0.07786         no
       9        0.04494        yes

Day 2 takeaway

Corrupt the input and keep the clean target and you have a denoising autoencoder, which learns a better representation than plain reconstruction. Use convolutions for images and ConvTranspose2d to upsample. And reconstruction error on data the model never saw is an anomaly score that needs no labels.

Week 12 · Day 3 of 7

The Latent Space Problem

Why you cannot sample from a plain autoencoder

By Pawan Rama Mali864 words

The latent space of a plain autoencoder is useful for compression and useless for generation, and the reason is worth seeing before the fix.

What the code actually looks like

import torch
from torch import nn
from torch.utils.data import DataLoader, Subset
from torchvision import datasets, transforms

tf = transforms.ToTensor()
train_all = datasets.MNIST('data', train=True, download=True, transform=tf)
test_all = datasets.MNIST('data', train=False, download=True, transform=tf)
train_set = Subset(train_all, range(8000))
val_set = Subset(test_all, range(2000))
train_loader = DataLoader(train_set, batch_size=128, shuffle=True)
val_loader = DataLoader(val_set, batch_size=512)

RAMP = ' .:-=+*#%@'
def show(img, step=2):
    """Print a 28 by 28 image as characters."""
    t = img.detach().reshape(28, 28)
    t = (t - t.min()) / (t.max() - t.min() + 1e-9)
    for row in t[::step]:
        print(''.join(RAMP[min(9, int(v * 9.999))] for v in row.tolist()))
class AutoEncoder(nn.Module):
    def __init__(self, latent=16):
        super().__init__()
        self.encoder = nn.Sequential(
            nn.Flatten(),
            nn.Linear(784, 256), nn.ReLU(),
            nn.Linear(256, 64), nn.ReLU(),
            nn.Linear(64, latent))
        self.decoder = nn.Sequential(
            nn.Linear(latent, 64), nn.ReLU(),
            nn.Linear(64, 256), nn.ReLU(),
            nn.Linear(256, 784), nn.Sigmoid())

    def forward(self, x):
        return self.decoder(self.encoder(x))

def train_ae(model, epochs=8, lr=1e-3, noise=0.0, loader=None):
    torch.manual_seed(0)
    opt = torch.optim.AdamW(model.parameters(), lr=lr)
    loss_fn = nn.MSELoss()
    for _ in range(epochs):
        model.train()
        for xb, _ in (loader or train_loader):
            target = xb.flatten(1)
            inp = xb if noise == 0 else (xb + noise * torch.randn_like(xb)
                                         ).clamp(0, 1)
            opt.zero_grad()
            loss_fn(model(inp), target).backward()
            opt.step()
    model.eval()
    total = seen = 0
    with torch.no_grad():
        for xb, _ in val_loader:
            total += loss_fn(model(xb), xb.flatten(1)).item() * len(xb)
            seen += len(xb)
    return total / seen
torch.manual_seed(0)
model = AutoEncoder(latent=2)          # two numbers, so we can print them
train_ae(model, epochs=10)

model.eval()
codes, labels = [], []
with torch.no_grad():
    for xb, yb in val_loader:
        codes.append(model.encoder(xb))
        labels.append(yb)
codes = torch.cat(codes)
labels = torch.cat(labels)

print('latent range: x %.2f to %.2f, y %.2f to %.2f'
      % (codes[:, 0].min(), codes[:, 0].max(),
         codes[:, 1].min(), codes[:, 1].max()))
print('\naverage position of each digit:')
for digit in range(10):
    mean = codes[labels == digit].mean(0)
    print('  %d  (%7.2f, %7.2f)' % (digit, mean[0], mean[1]))

latent range: x -11.54 to 9.20, y -16.03 to 2.54

average position of each digit:
(  -5.48,   -2.98)
(   4.38,   -7.70)
(  -2.20,   -4.17)
(  -2.05,   -2.76)
(  -1.59,   -5.16)
(  -2.19,   -4.33)
(  -2.61,   -4.03)
(  -0.77,   -6.77)
(  -1.76,   -4.91)
(  -1.33,   -5.84)

There is no reason for the space between clusters to mean anything

Nothing in the loss asked the codes to be arranged sensibly. It asked only that each one decodes back to its own image. The clusters end up wherever the optimiser put them, separated by regions no training image ever occupied, and a decoder given a point in one of those gaps has never been asked to produce anything meaningful there.

Decoding a point nobody trained on

import torch
from torch import nn
from torch.utils.data import DataLoader, Subset
from torchvision import datasets, transforms

tf = transforms.ToTensor()
train_all = datasets.MNIST('data', train=True, download=True, transform=tf)
test_all = datasets.MNIST('data', train=False, download=True, transform=tf)
train_set = Subset(train_all, range(8000))
val_set = Subset(test_all, range(2000))
train_loader = DataLoader(train_set, batch_size=128, shuffle=True)
val_loader = DataLoader(val_set, batch_size=512)

RAMP = ' .:-=+*#%@'
def show(img, step=2):
    """Print a 28 by 28 image as characters."""
    t = img.detach().reshape(28, 28)
    t = (t - t.min()) / (t.max() - t.min() + 1e-9)
    for row in t[::step]:
        print(''.join(RAMP[min(9, int(v * 9.999))] for v in row.tolist()))
class AutoEncoder(nn.Module):
    def __init__(self, latent=16):
        super().__init__()
        self.encoder = nn.Sequential(
            nn.Flatten(),
            nn.Linear(784, 256), nn.ReLU(),
            nn.Linear(256, 64), nn.ReLU(),
            nn.Linear(64, latent))
        self.decoder = nn.Sequential(
            nn.Linear(latent, 64), nn.ReLU(),
            nn.Linear(64, 256), nn.ReLU(),
            nn.Linear(256, 784), nn.Sigmoid())

    def forward(self, x):
        return self.decoder(self.encoder(x))

def train_ae(model, epochs=8, lr=1e-3, noise=0.0, loader=None):
    torch.manual_seed(0)
    opt = torch.optim.AdamW(model.parameters(), lr=lr)
    loss_fn = nn.MSELoss()
    for _ in range(epochs):
        model.train()
        for xb, _ in (loader or train_loader):
            target = xb.flatten(1)
            inp = xb if noise == 0 else (xb + noise * torch.randn_like(xb)
                                         ).clamp(0, 1)
            opt.zero_grad()
            loss_fn(model(inp), target).backward()
            opt.step()
    model.eval()
    total = seen = 0
    with torch.no_grad():
        for xb, _ in val_loader:
            total += loss_fn(model(xb), xb.flatten(1)).item() * len(xb)
            seen += len(xb)
    return total / seen
torch.manual_seed(0)
model = AutoEncoder(latent=2)
train_ae(model, epochs=10)
model.eval()

with torch.no_grad():
    codes = torch.cat([model.encoder(xb) for xb, _ in val_loader])
    # a point midway between two real codes
    a, b = codes[0], codes[1]
    for name, point in [('a real code', a),
                        ('halfway to another', (a + b) / 2),
                        ('a random point', torch.randn(2) * codes.std(0)
                         + codes.mean(0))]:
        img = model.decoder(point.unsqueeze(0))
        print('%s:' % name)
        show(img[0], step=3)
        print()

a real code:
                            
                            
             .::::::..      
      ::=*%%%*===*##+-.     
      .:+*+-:. .+#%*-.      
      .::--===+#@@*:        
         ...:-#@%=          
         ..:=#*=-.          
         .:-==-:.           
                            

halfway to another:
                            
                            
          .:--=--:::..      
     ..:=*#*+---=**+-.      
      .-+*=-:..:+*#=:.      
      .-=+*****#@@#=        
       ..:---:=*@%+.        
        ...::=+**+:         
           .:===-:.         
                            

a random point:
                            
                            
       .::-=+#*+-:..        
      .......:===+-.        
          ..-+%%+-.         
          :-=%@*-.          
         .::-----     .     
        .....-+#+-:-:.      
            .:..            
                            

Day 3 takeaway

A plain autoencoder arranges its latent space however training happened to leave it. Points between clusters decode to nothing in particular, because nothing in the loss required otherwise. Sampling from such a space does not generate anything, which is what tomorrow fixes.

Week 12 · Day 4 of 7

The Variational Autoencoder

Encoding a distribution, the reparameterisation trick, and the KL term

By Pawan Rama Mali637 words

A variational autoencoder makes two changes. The encoder outputs a distribution rather than a point, and the loss adds a term pulling those distributions towards a standard normal. Together they make the latent space something you can sample from.

The reparameterisation trick

Reparameterisation: You cannot backpropagate through a random draw. So instead of sampling z directly, draw a standard normal epsilon and compute z = mu + sigma * epsilon. The randomness now sits in epsilon, which has no parameters, and the gradient flows cleanly through mu and sigma.

import torch

torch.manual_seed(0)
mu = torch.tensor([1.0], requires_grad=True)
log_var = torch.tensor([0.0], requires_grad=True)

# the wrong way: sample, then try to differentiate
z_bad = torch.normal(mu.detach(), torch.ones(1))
print('sampling directly gives a tensor with no graph:',
      z_bad.grad_fn is None)

# the right way
eps = torch.randn(1)
z = mu + (0.5 * log_var).exp() * eps
z.sum().backward()
print('reparameterised, the gradient reaches mu:  %.4f' % mu.grad.item())
print('and log_var:                               %.4f'
      % log_var.grad.item())
print('\nthe randomness is in eps, which has no parameters to learn.')

sampling directly gives a tensor with no graph: True
reparameterised, the gradient reaches mu:  1.0000
and log_var:                               -0.1467

the randomness is in eps, which has no parameters to learn.

The two-part loss

import torch
from torch import nn
from torch.utils.data import DataLoader, Subset
from torchvision import datasets, transforms

tf = transforms.ToTensor()
train_all = datasets.MNIST('data', train=True, download=True, transform=tf)
test_all = datasets.MNIST('data', train=False, download=True, transform=tf)
train_set = Subset(train_all, range(8000))
val_set = Subset(test_all, range(2000))
train_loader = DataLoader(train_set, batch_size=128, shuffle=True)
val_loader = DataLoader(val_set, batch_size=512)

RAMP = ' .:-=+*#%@'
def show(img, step=2):
    """Print a 28 by 28 image as characters."""
    t = img.detach().reshape(28, 28)
    t = (t - t.min()) / (t.max() - t.min() + 1e-9)
    for row in t[::step]:
        print(''.join(RAMP[min(9, int(v * 9.999))] for v in row.tolist()))
class VAE(nn.Module):
    def __init__(self, latent=16):
        super().__init__()
        self.body = nn.Sequential(nn.Flatten(),
                                  nn.Linear(784, 256), nn.ReLU(),
                                  nn.Linear(256, 128), nn.ReLU())
        self.to_mu = nn.Linear(128, latent)
        self.to_log_var = nn.Linear(128, latent)
        self.decoder = nn.Sequential(
            nn.Linear(latent, 128), nn.ReLU(),
            nn.Linear(128, 256), nn.ReLU(),
            nn.Linear(256, 784), nn.Sigmoid())

    def encode(self, x):
        h = self.body(x)
        return self.to_mu(h), self.to_log_var(h)

    def forward(self, x):
        mu, log_var = self.encode(x)
        z = mu + (0.5 * log_var).exp() * torch.randn_like(mu)
        return self.decoder(z), mu, log_var

def vae_loss(out, target, mu, log_var, beta=1.0):
    # how well it rebuilt the input, summed over pixels
    recon = nn.functional.binary_cross_entropy(out, target,
                                               reduction='sum')
    # how far each latent distribution is from a standard normal
    kl = -0.5 * torch.sum(1 + log_var - mu.pow(2) - log_var.exp())
    return recon + beta * kl, recon, kl

torch.manual_seed(0)
model = VAE()
x, _ = next(iter(val_loader))
out, mu, log_var = model(x[:8])
total, recon, kl = vae_loss(out, x[:8].flatten(1), mu, log_var)
print('reconstruction term %.1f' % recon.item())
print('kl term             %.1f' % kl.item())
print('total               %.1f' % total.item())

reconstruction term 4364.9
kl term             0.2
total               4365.1

What the second term is doing

The reconstruction term wants each image to have its own distinct code, as far from the others as possible, because that is easiest to decode. The KL term wants every code to be a standard normal, which would make them all identical and useless. Training balances the two, and the result is a latent space that is spread out enough to decode and packed tightly enough that the region between two codes is also somewhere the decoder has learned about.

Day 4 takeaway

A VAE encodes to a distribution rather than a point, samples with the reparameterisation trick so gradients flow, and adds a KL term pulling those distributions towards a standard normal. The two terms pull in opposite directions, and the tension is what makes the space samplable.

Week 12 · Day 5 of 7

Sampling and Interpolating

Generating digits from noise, and walking between two of them

By Pawan Rama Mali1101 words

Trained, sampled, and interpolated, which is the test a plain autoencoder failed yesterday.

Training

import torch
from torch import nn
from torch.utils.data import DataLoader, Subset
from torchvision import datasets, transforms

tf = transforms.ToTensor()
train_all = datasets.MNIST('data', train=True, download=True, transform=tf)
test_all = datasets.MNIST('data', train=False, download=True, transform=tf)
train_set = Subset(train_all, range(8000))
val_set = Subset(test_all, range(2000))
train_loader = DataLoader(train_set, batch_size=128, shuffle=True)
val_loader = DataLoader(val_set, batch_size=512)

RAMP = ' .:-=+*#%@'
def show(img, step=2):
    """Print a 28 by 28 image as characters."""
    t = img.detach().reshape(28, 28)
    t = (t - t.min()) / (t.max() - t.min() + 1e-9)
    for row in t[::step]:
        print(''.join(RAMP[min(9, int(v * 9.999))] for v in row.tolist()))
import time

class VAE(nn.Module):
    def __init__(self, latent=16):
        super().__init__()
        self.body = nn.Sequential(nn.Flatten(),
                                  nn.Linear(784, 256), nn.ReLU(),
                                  nn.Linear(256, 128), nn.ReLU())
        self.to_mu = nn.Linear(128, latent)
        self.to_log_var = nn.Linear(128, latent)
        self.decoder = nn.Sequential(
            nn.Linear(latent, 128), nn.ReLU(),
            nn.Linear(128, 256), nn.ReLU(),
            nn.Linear(256, 784), nn.Sigmoid())
        self.latent = latent

    def encode(self, x):
        h = self.body(x)
        return self.to_mu(h), self.to_log_var(h)

    def forward(self, x):
        mu, log_var = self.encode(x)
        z = mu + (0.5 * log_var).exp() * torch.randn_like(mu)
        return self.decoder(z), mu, log_var

def train_vae(model, epochs=10, beta=1.0):
    torch.manual_seed(0)
    opt = torch.optim.AdamW(model.parameters(), lr=1e-3)
    for epoch in range(epochs):
        model.train()
        for xb, _ in train_loader:
            target = xb.flatten(1)
            out, mu, log_var = model(xb)
            recon = nn.functional.binary_cross_entropy(out, target,
                                                       reduction='sum')
            kl = -0.5 * torch.sum(1 + log_var - mu.pow(2) - log_var.exp())
            opt.zero_grad()
            ((recon + beta * kl) / len(xb)).backward()
            opt.step()
    return model

torch.manual_seed(0)
start = time.time()
model = train_vae(VAE())
model.eval()

with torch.no_grad():
    total = kl_total = seen = 0
    for xb, _ in val_loader:
        out, mu, log_var = model(xb)
        total += nn.functional.binary_cross_entropy(
            out, xb.flatten(1), reduction='sum').item()
        kl_total += (-0.5 * torch.sum(
            1 + log_var - mu.pow(2) - log_var.exp())).item()
        seen += len(xb)
print('per image: reconstruction %.1f, kl %.1f'
      % (total / seen, kl_total / seen))
print('trained in %.0f seconds' % (time.time() - start))

per image: reconstruction 129.0, kl 12.1
trained in 7 seconds

Sampling from nothing but noise

import torch
from torch import nn
from torch.utils.data import DataLoader, Subset
from torchvision import datasets, transforms

tf = transforms.ToTensor()
train_all = datasets.MNIST('data', train=True, download=True, transform=tf)
test_all = datasets.MNIST('data', train=False, download=True, transform=tf)
train_set = Subset(train_all, range(8000))
val_set = Subset(test_all, range(2000))
train_loader = DataLoader(train_set, batch_size=128, shuffle=True)
val_loader = DataLoader(val_set, batch_size=512)

RAMP = ' .:-=+*#%@'
def show(img, step=2):
    """Print a 28 by 28 image as characters."""
    t = img.detach().reshape(28, 28)
    t = (t - t.min()) / (t.max() - t.min() + 1e-9)
    for row in t[::step]:
        print(''.join(RAMP[min(9, int(v * 9.999))] for v in row.tolist()))
class VAE(nn.Module):
    def __init__(self, latent=16):
        super().__init__()
        self.body = nn.Sequential(nn.Flatten(),
                                  nn.Linear(784, 256), nn.ReLU(),
                                  nn.Linear(256, 128), nn.ReLU())
        self.to_mu = nn.Linear(128, latent)
        self.to_log_var = nn.Linear(128, latent)
        self.decoder = nn.Sequential(
            nn.Linear(latent, 128), nn.ReLU(),
            nn.Linear(128, 256), nn.ReLU(),
            nn.Linear(256, 784), nn.Sigmoid())
        self.latent = latent
    def encode(self, x):
        h = self.body(x)
        return self.to_mu(h), self.to_log_var(h)
    def forward(self, x):
        mu, log_var = self.encode(x)
        z = mu + (0.5 * log_var).exp() * torch.randn_like(mu)
        return self.decoder(z), mu, log_var

torch.manual_seed(0)
model = VAE()
opt = torch.optim.AdamW(model.parameters(), lr=1e-3)
for _ in range(10):
    model.train()
    for xb, _ in train_loader:
        out, mu, log_var = model(xb)
        recon = nn.functional.binary_cross_entropy(
            out, xb.flatten(1), reduction='sum')
        kl = -0.5 * torch.sum(1 + log_var - mu.pow(2) - log_var.exp())
        opt.zero_grad()
        ((recon + kl) / len(xb)).backward()
        opt.step()
model.eval()

torch.manual_seed(3)
with torch.no_grad():
    for i in range(2):
        z = torch.randn(1, model.latent)
        print('sample %d, from a random point in the latent space:' % (i + 1))
        show(model.decoder(z)[0], step=2)
        print()

sample 1, from a random point in the latent space:
                            
                            
                            
       ..::------:::.:...   
     ....:::-=+*++++-:::.   
         ..:-+##**+=:...    
         .:-+*++++-:.       
         .-+#%%%#+-..       
         .-=+#@@#+:..       
      ...::-=*@@*-.         
      .::::-=+*=:.          
       ...::---:..          
      .........             
                            

sample 2, from a random point in the latent space:
                            
                            
                            
             .....          
           :-+#%#*=:        
          .=++*+=--.        
          -=---==-.         
          ..:-++*-.         
             .=#+.          
            .=@*.           
           :#@*:            
          -#@*:             
         :+*=:              
                            

No image went in

The only input was a draw from a standard normal. The decoder has learned a map from that distribution to something resembling a digit, which is what makes this a generative model rather than a compressor. A plain autoencoder cannot do this, because nothing ever told it what the distribution of its codes should be.

Interpolating between two digits

import torch
from torch import nn
from torch.utils.data import DataLoader, Subset
from torchvision import datasets, transforms

tf = transforms.ToTensor()
train_all = datasets.MNIST('data', train=True, download=True, transform=tf)
test_all = datasets.MNIST('data', train=False, download=True, transform=tf)
train_set = Subset(train_all, range(8000))
val_set = Subset(test_all, range(2000))
train_loader = DataLoader(train_set, batch_size=128, shuffle=True)
val_loader = DataLoader(val_set, batch_size=512)

RAMP = ' .:-=+*#%@'
def show(img, step=2):
    """Print a 28 by 28 image as characters."""
    t = img.detach().reshape(28, 28)
    t = (t - t.min()) / (t.max() - t.min() + 1e-9)
    for row in t[::step]:
        print(''.join(RAMP[min(9, int(v * 9.999))] for v in row.tolist()))
class VAE(nn.Module):
    def __init__(self, latent=16):
        super().__init__()
        self.body = nn.Sequential(nn.Flatten(),
                                  nn.Linear(784, 256), nn.ReLU(),
                                  nn.Linear(256, 128), nn.ReLU())
        self.to_mu = nn.Linear(128, latent)
        self.to_log_var = nn.Linear(128, latent)
        self.decoder = nn.Sequential(
            nn.Linear(latent, 128), nn.ReLU(),
            nn.Linear(128, 256), nn.ReLU(),
            nn.Linear(256, 784), nn.Sigmoid())
    def encode(self, x):
        h = self.body(x)
        return self.to_mu(h), self.to_log_var(h)
    def forward(self, x):
        mu, log_var = self.encode(x)
        z = mu + (0.5 * log_var).exp() * torch.randn_like(mu)
        return self.decoder(z), mu, log_var

torch.manual_seed(0)
model = VAE()
opt = torch.optim.AdamW(model.parameters(), lr=1e-3)
for _ in range(10):
    model.train()
    for xb, _ in train_loader:
        out, mu, log_var = model(xb)
        recon = nn.functional.binary_cross_entropy(
            out, xb.flatten(1), reduction='sum')
        kl = -0.5 * torch.sum(1 + log_var - mu.pow(2) - log_var.exp())
        opt.zero_grad()
        ((recon + kl) / len(xb)).backward()
        opt.step()
model.eval()

x, y = next(iter(val_loader))
with torch.no_grad():
    mu, _ = model.encode(x[:2])
print('interpolating from a %d to a %d:\n' % (y[0], y[1]))
with torch.no_grad():
    for t in [0.0, 0.5, 1.0]:
        z = (1 - t) * mu[0] + t * mu[1]
        print('t = %.1f' % t)
        show(model.decoder(z.unsqueeze(0))[0], step=3)
        print()

interpolating from a 7 to a 2:

t = 0.0
                            
                            
                            
    .=#%@@@@@@@@@@@@*:      
       .::..   .-#@@=.      
               =%@%-        
             :%@%-          
           .*%%+.           
         .=**+-             
                            

t = 0.5
                            
                            
    .-+##**=--::....        
  .:*%%+=-:::=*%@%*-.       
         .:-=*#*+=-:.       
         ..---==**=-.       
           .::=+=-:...      
          .:::---:-::..     
     ....:------:......     
               .            

t = 1.0
                            
        .:-+#%##+-          
       :+*#*#*+=-:          
             .=*+-.         
           :=#%#=:          
          .----:---..       
       ..:-------=++*#+=-   
       :-=+##*#*+===-:..    
                            
                            

Day 5 takeaway

A trained VAE generates from a draw of standard normal noise, and the path between two encoded images passes through points that decode to plausible digits. Both are things a plain autoencoder cannot do, and both come from the KL term.

Week 12 · Day 6 of 7

Posterior Collapse and Blur

The beta trade-off, and why the samples are soft

By Pawan Rama Mali964 words

Two things that go wrong with VAEs, both well known, and the parameter that trades between them.

Posterior collapse

import torch
from torch import nn
from torch.utils.data import DataLoader, Subset
from torchvision import datasets, transforms

tf = transforms.ToTensor()
train_all = datasets.MNIST('data', train=True, download=True, transform=tf)
test_all = datasets.MNIST('data', train=False, download=True, transform=tf)
train_set = Subset(train_all, range(8000))
val_set = Subset(test_all, range(2000))
train_loader = DataLoader(train_set, batch_size=128, shuffle=True)
val_loader = DataLoader(val_set, batch_size=512)

RAMP = ' .:-=+*#%@'
def show(img, step=2):
    """Print a 28 by 28 image as characters."""
    t = img.detach().reshape(28, 28)
    t = (t - t.min()) / (t.max() - t.min() + 1e-9)
    for row in t[::step]:
        print(''.join(RAMP[min(9, int(v * 9.999))] for v in row.tolist()))
class VAE(nn.Module):
    def __init__(self, latent=16):
        super().__init__()
        self.body = nn.Sequential(nn.Flatten(),
                                  nn.Linear(784, 256), nn.ReLU())
        self.to_mu = nn.Linear(256, latent)
        self.to_log_var = nn.Linear(256, latent)
        self.decoder = nn.Sequential(
            nn.Linear(latent, 256), nn.ReLU(),
            nn.Linear(256, 784), nn.Sigmoid())
    def encode(self, x):
        h = self.body(x)
        return self.to_mu(h), self.to_log_var(h)
    def forward(self, x):
        mu, log_var = self.encode(x)
        z = mu + (0.5 * log_var).exp() * torch.randn_like(mu)
        return self.decoder(z), mu, log_var

print('%8s %14s %12s %16s'
      % ('beta', 'reconstruction', 'kl', 'active dims'))
for beta in [0.1, 1.0, 8.0, 60.0]:
    torch.manual_seed(0)
    model = VAE()
    opt = torch.optim.AdamW(model.parameters(), lr=1e-3)
    for _ in range(6):
        model.train()
        for xb, _ in train_loader:
            out, mu, log_var = model(xb)
            recon = nn.functional.binary_cross_entropy(
                out, xb.flatten(1), reduction='sum')
            kl = -0.5 * torch.sum(1 + log_var - mu.pow(2) - log_var.exp())
            opt.zero_grad()
            ((recon + beta * kl) / len(xb)).backward()
            opt.step()
    model.eval()
    with torch.no_grad():
        per_dim, r, k, n = [], 0.0, 0.0, 0
        for xb, _ in val_loader:
            out, mu, log_var = model(xb)
            per_dim.append(-0.5 * (1 + log_var - mu.pow(2)
                                   - log_var.exp()).mean(0))
            r += nn.functional.binary_cross_entropy(
                out, xb.flatten(1), reduction='sum').item()
            k += (-0.5 * torch.sum(
                1 + log_var - mu.pow(2) - log_var.exp())).item()
            n += len(xb)
        kl_per_dim = torch.stack(per_dim).mean(0)
    # a dimension carrying no information has a KL of about zero
    active = int((kl_per_dim > 0.01).sum())
    print('%8.1f %14.1f %12.1f %16d'
          % (beta, r / n, k / n, active))

    beta reconstruction           kl      active dims
1          108.8         52.6               16
0          123.5         16.4               16
0          175.7          1.9               16
0          201.4          0.0                0

Posterior collapse: When the KL term is strong enough, the cheapest way to satisfy it is for the encoder to ignore the input entirely and output the standard normal for every image. Those latent dimensions carry no information at all, and the decoder learns to produce the average image regardless of input.

The active dimensions column counts how many latent coordinates still vary between images. As beta rises, reconstruction gets worse and dimensions switch off. That trade is the single most important setting in a VAE, and it is why beta appears in the name of half the papers about them.

Why VAE samples are blurry

import torch

# Two equally plausible sharp images, and what minimising squared error
# against both produces.
a = torch.zeros(8, 8)
a[2:6, 2:4] = 1.0
b = torch.zeros(8, 8)
b[2:6, 5:7] = 1.0

best_single = (a + b) / 2
print('two plausible outputs, and the one that minimises the mean')
print('squared error against both:\n')
for name, img in [('option a', a), ('option b', b),
                  ('the average', best_single)]:
    print(name)
    for row in img:
        print('  ' + ''.join('%3.0f' % (v * 100) for v in row.tolist()))
    print()

for name, img in [('option a', a), ('the average', best_single)]:
    err = 0.5 * ((img - a) ** 2).mean() + 0.5 * ((img - b) ** 2).mean()
    print('%-14s expected error %.4f' % (name, err.item()))

two plausible outputs, and the one that minimises the mean
squared error against both:

option a
0  0  0  0  0  0  0
0  0  0  0  0  0  0
0100100  0  0  0  0
0100100  0  0  0  0
0100100  0  0  0  0
0100100  0  0  0  0
0  0  0  0  0  0  0
0  0  0  0  0  0  0

option b
0  0  0  0  0  0  0
0  0  0  0  0  0  0
0  0  0  0100100  0
0  0  0  0100100  0
0  0  0  0100100  0
0  0  0  0100100  0
0  0  0  0  0  0  0
0  0  0  0  0  0  0

the average
0  0  0  0  0  0  0
0  0  0  0  0  0  0
0 50 50  0 50 50  0
0 50 50  0 50 50  0
0 50 50  0 50 50  0
0 50 50  0 50 50  0
0  0  0  0  0  0  0
0  0  0  0  0  0  0

option a       expected error 0.1250
the average    expected error 0.0625

The blur is the loss doing its job

When several outputs are equally plausible, the one that minimises expected squared error is their average, and the average of several sharp images is a blurry one. A VAE is not failing when its samples look soft; it is correctly optimising a loss that rewards hedging.

That is precisely the gap the next week's models attack. A GAN replaces the pixel loss with a second network that judges whether the output looks real, and averaging two plausible images looks obviously fake to such a judge. Diffusion takes a different route to the same end.

Day 6 takeaway

Beta trades reconstruction quality against how well the latent space matches a standard normal, and pushing it too far switches latent dimensions off entirely. VAE samples are blurry because squared error rewards averaging over plausible outputs, which is the problem week 13 exists to solve.

Week 12 · Day 7 of 7

What They Are Actually For

Linear probes, an honest comparison, and the checklist

By Pawan Rama Mali917 words

What these models are actually for, measured rather than asserted.

Representation quality

import torch
from torch import nn
from torch.utils.data import DataLoader, Subset
from torchvision import datasets, transforms

tf = transforms.ToTensor()
train_all = datasets.MNIST('data', train=True, download=True, transform=tf)
test_all = datasets.MNIST('data', train=False, download=True, transform=tf)
train_set = Subset(train_all, range(8000))
val_set = Subset(test_all, range(2000))
train_loader = DataLoader(train_set, batch_size=128, shuffle=True)
val_loader = DataLoader(val_set, batch_size=512)

RAMP = ' .:-=+*#%@'
def show(img, step=2):
    """Print a 28 by 28 image as characters."""
    t = img.detach().reshape(28, 28)
    t = (t - t.min()) / (t.max() - t.min() + 1e-9)
    for row in t[::step]:
        print(''.join(RAMP[min(9, int(v * 9.999))] for v in row.tolist()))
class AutoEncoder(nn.Module):
    def __init__(self, latent=16):
        super().__init__()
        self.encoder = nn.Sequential(
            nn.Flatten(),
            nn.Linear(784, 256), nn.ReLU(),
            nn.Linear(256, 64), nn.ReLU(),
            nn.Linear(64, latent))
        self.decoder = nn.Sequential(
            nn.Linear(latent, 64), nn.ReLU(),
            nn.Linear(64, 256), nn.ReLU(),
            nn.Linear(256, 784), nn.Sigmoid())

    def forward(self, x):
        return self.decoder(self.encoder(x))

def train_ae(model, epochs=8, lr=1e-3, noise=0.0, loader=None):
    torch.manual_seed(0)
    opt = torch.optim.AdamW(model.parameters(), lr=lr)
    loss_fn = nn.MSELoss()
    for _ in range(epochs):
        model.train()
        for xb, _ in (loader or train_loader):
            target = xb.flatten(1)
            inp = xb if noise == 0 else (xb + noise * torch.randn_like(xb)
                                         ).clamp(0, 1)
            opt.zero_grad()
            loss_fn(model(inp), target).backward()
            opt.step()
    model.eval()
    total = seen = 0
    with torch.no_grad():
        for xb, _ in val_loader:
            total += loss_fn(model(xb), xb.flatten(1)).item() * len(xb)
            seen += len(xb)
    return total / seen
from sklearn.linear_model import LogisticRegression
import numpy as np

torch.manual_seed(0)
plain = AutoEncoder(latent=32)
train_ae(plain, epochs=8)
torch.manual_seed(0)
denoise = AutoEncoder(latent=32)
train_ae(denoise, epochs=8, noise=0.4)

def features(model, loader):
    model.eval()
    xs, ys = [], []
    with torch.no_grad():
        for xb, yb in loader:
            xs.append(model.encoder(xb))
            ys.append(yb)
    return torch.cat(xs).numpy(), torch.cat(ys).numpy()

def raw(loader):
    # one pass. Two passes over a shuffled loader return the
    # images in one order and the labels in another.
    xs, ys = [], []
    for xb, yb in loader:
        xs.append(xb.flatten(1))
        ys.append(yb)
    return torch.cat(xs).numpy(), torch.cat(ys).numpy()

raw_tr, raw_y = raw(train_loader)
raw_va, raw_vy = raw(val_loader)

print('%-34s %10s %10s' % ('features for a linear classifier',
                           'width', 'accuracy'))
clf = LogisticRegression(max_iter=1000).fit(raw_tr, raw_y)
print('%-34s %10d %10.4f' % ('raw pixels', raw_tr.shape[1],
                             clf.score(raw_va, raw_vy)))
for name, model in [('plain autoencoder code', plain),
                    ('denoising autoencoder code', denoise)]:
    Xtr, ytr = features(model, train_loader)
    Xva, yva = features(model, val_loader)
    clf = LogisticRegression(max_iter=300).fit(Xtr, ytr)
    print('%-34s %10d %10.4f'
          % (name, Xtr.shape[1], clf.score(Xva, yva)))

features for a linear classifier        width   accuracy
raw pixels                                784     0.8705
plain autoencoder code                     32     0.8310
denoising autoencoder code                 32     0.7610

Two results worth sitting with, because neither is what the textbook ordering predicts. The denoising code is worse than the plain one for this classifier, and both are behind raw pixels.

On MNIST, raw pixels are a strong baseline

MNIST digits are close to linearly separable in pixel space already, so a linear classifier on 784 raw numbers is hard to beat. What the autoencoder achieved is compression: 32 numbers reach within a few points of 784, which is the interesting result rather than a failure.

The denoising variant learning a worse representation here is a reminder that its advantage is not universal. It helps when the corruption forces the model to learn structure it would otherwise ignore; on clean, centred, high-contrast digits there is little such structure to find, and the noise mostly makes the job harder. Run the probe rather than assuming the ordering.

What each of them is for

Model	Good at	Not good at
Plain autoencoder	Compression, anomaly detection	Generating anything
Denoising autoencoder	Representations, actual denoising	Generating anything
VAE	Generating, interpolating, a principled latent space	Sharp samples
GAN (week 13)	Sharp samples	Coverage, stable training
Diffusion (week 13)	Sharp samples and coverage	Speed, at generation time

Compression is not what these are for

A 32 number code from 784 pixels sounds like a twenty-four to one compression, and JPEG will beat it comfortably on any measure you care to name, while working on images it has never seen and needing no training. An autoencoder only compresses well on data resembling its training set.

The value is elsewhere: a representation learned without labels, an anomaly score that needs no examples of anomalies, and in the VAE case a latent space you can sample from. Judge them on those.

The checklist

Make the bottleneck genuinely narrow, or the model learns the identity.
Try the denoising variant, which costs one line, and measure it rather than assuming it helps.
Use convolutions for images, and check your ConvTranspose2d output shapes.
For a VAE, use the reparameterisation trick and sum the reconstruction term over pixels rather than averaging, so the two terms are on comparable scales.
Watch the number of active latent dimensions; if it is falling, beta is too high.
Judge a generative model by sampling from noise, not by reconstruction error.
Compare any learned representation against raw pixels with a linear classifier on top, which is the standard probe and takes two lines.

Day 7 takeaway

Autoencoders learn representations and anomaly scores without labels; VAEs additionally give you a latent space you can sample from. Always measure a learned representation with a linear probe against raw pixels, because on an easy dataset the raw pixels win and the honest claim is compression rather than superiority. Judge a generative model by what it produces from noise.