🎓 A1Learner
🔍

PyTorch A Latent Space Walk

A latent space walk moves a code in small steps between two encoded points and decodes every step. The decoded points show whether the space changes smoothly or jumps.

Seven codes between two means

torch.linspace lists evenly spaced values between two ends. These seven codes run from the start point's latent mean to the end point's latent mean.

steps = torch.linspace(-1.12, 0.79, 7)
print([round(v, 2) for v in steps.tolist()])

One code becomes one point

The decoder turns any single code into a 2D point. This small untrained decoder shows the mechanics only.

torch.manual_seed(1)
decoder = nn.Sequential(nn.Linear(1, 16), nn.ReLU(), nn.Linear(16, 2))
z = torch.tensor([[0.0]])
print([round(v, 2) for v in decoder(z)[0].tolist()])

Decoding every step

The closing program repeats this decode for all seven codes. Each decoded point prints on its own line.

for step in range(7):
    frac = step / 6
    z = z_start + (z_end - z_start) * frac
    decoded = vae.decoder(z)

Example

import torch
from torch import nn

torch.manual_seed(0)
per_blob = 50
blob_a = torch.randn(per_blob, 2) * 0.5 + torch.tensor([-2.0, -2.0])
blob_b = torch.randn(per_blob, 2) * 0.5 + torch.tensor([2.0, 2.0])
points = torch.cat([blob_a, blob_b])


class VAE(nn.Module):
    def __init__(self):
        super().__init__()
        self.body = nn.Sequential(nn.Linear(2, 16), nn.ReLU())
        self.to_mean = nn.Linear(16, 1)
        self.to_logvar = nn.Linear(16, 1)
        self.decoder = nn.Sequential(nn.Linear(1, 16), nn.ReLU(), nn.Linear(16, 2))

    def encode(self, x):
        h = self.body(x)
        return self.to_mean(h), self.to_logvar(h)

    def forward(self, x):
        mean, logvar = self.encode(x)
        std = torch.exp(0.5 * logvar)
        z = torch.distributions.Normal(mean, std).rsample()
        return self.decoder(z), mean, logvar


torch.manual_seed(1)
vae = VAE()
opt = torch.optim.Adam(vae.parameters(), lr=0.02)
for epoch in range(300):
    rebuilt, mean, logvar = vae(points)
    recon = nn.functional.mse_loss(rebuilt, points)
    kl = -0.5 * torch.mean(1 + logvar - mean.pow(2) - logvar.exp())
    loss = recon + kl
    opt.zero_grad()
    loss.backward()
    opt.step()

with torch.no_grad():
    start = points[0].unsqueeze(0)
    end = points[99].unsqueeze(0)
    z_start = vae.encode(start)[0]
    z_end = vae.encode(end)[0]
    print("Start point:", [round(v, 2) for v in start[0].tolist()],
          "latent:", round(z_start.item(), 2))
    print("End point:", [round(v, 2) for v in end[0].tolist()],
          "latent:", round(z_end.item(), 2))
    print("The walk:")
    for step in range(7):
        frac = step / 6
        z = z_start + (z_end - z_start) * frac
        decoded = vae.decoder(z)
        print(f"Step {step}: ({decoded[0, 0].item():.2f}, {decoded[0, 1].item():.2f})")

Key points

  • The walk starts and ends at the encoder means of two real points
  • Even steps in the code give steady moves in the decoded points here
  • A smooth walk means the decoder treats nearby codes alike

In this example, the VAE trains on the two blobs with the recipe from the training lesson. The start point is (-2.56, -2.58), and its latent mean is -1.12. The end point is (1.97, 1.69), and its latent mean is 0.79. Seven codes step evenly from the first mean to the second. The decoded points move in one direction at every step, with no jump backwards. The first three steps stay near the first blob. Steps 3 and 4 cross the middle ground between the blobs. The last steps arrive at the second blob. The steps grow larger in the middle, but the path never teleports.

Note

The middle of the walk passes between the blobs, where no training point sits

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

During a latent walk, the decoded points jump from one blob straight to the other with nothing between. What does that show?