Abhinandan

Posted on Oct 12

Context Encoders: Learning to Inpaint and What It Teaches a Network

#ai #machinelearning #opensource #programming

Our visual world is very diverse, yet highly structured, and
humans have an uncanny ability to make sense of this structure.

This single line tells us a very deep meaning of how we human perceive visual identity throughout the nature. We humans have a strange ability of prediction, we can predict possibility of a certain things out of nowhere whenever there is ambiguity. For example, take this following image of a building, where in the first part is clearly blanked but in the second part, it was well predicted by an human artist.

As ever incrementing technology, where things are pretty much changing on day to day basis, it's very important for us to step back a little and introspect, did we really learn the fundamentals based on which this big AI orgs. are experimenting and giving somewhat true and false promises. In this technical blog, we will deep dive into one of the state-of-art paper which marked the foundational step of self-supervised learning, "Context Encoders: Feature Learning by Inpainting".

Consider the above shown image, although the central part of the image is missing but most of us can easily imagine it's content (what can be there) by looking at it's sorroundings. This is what make humans truely advanced, where we just don't need to see everything, our visual system can predict the missing part by looking only some portion of it. In light of this question, it's pretty much considerable to ask, can we have advanced computer vision technique, which can also do it? And the answer of this fundamental question first given by context encoders, which do the same by inpainting.

What is Inpainting?

By defination, inpainting is the process of restoring or modifying missing, damaged, or unwanted parts of an image or physical artwork so that the altered area blends seamlessly with the surroundings.

So, by defination, if a model/network that fills in a missing patch of an image has to understand what the image contains, and it needs no human labels to learn that. That is the bet behind Context Encoders. Context Encoders pursues two goals at once,

Semantic Inpainting: enerating plausible content for a hole, not just smearing nearby pixels into it.
Unsupervised Feature Learning: showing that an encoder trained only on this fill-in-the-blank task learns features that transfer to classification, detection and segmentation.

In this technical blog, we will deep dive into the core of context encoders and implement it using PyTorch and CelebA dataset from scratch.

The Code Idea

Computer vision has made tremendous progress on semantic image understanding tasks such as classification, object detection and segmentation in the past decade. Recently, Convolutional Neural Network (CNN) and Vision Transformer (ViT) have greately advanced the performance of these tasks. The success of such models on image classification paved the way to tackle harder problems, including unsupervised understanding (pre self-supervised learning) and generation of natural images.

A context encoder is a convolutional encoder-decoder trained to predict missing image content from the pixels around it. Think of a jigsaw puzzle with a piece missing, you can guess the piece only if you know the picture is a face, so you expect an eye, a nose and skin tones in the right places.

The training philosophy of context encoders are pretty much simple. Take an image x and binary mask M where 1 marks the hole. Feed the network the damaged image then ask it to produce an output that matches the original content inside the hole. Below is the representation image of context encoders.

Because the target is the image itself, labels are free. The hard part is that a hole has many valid fillings. For that reason, loss design has a very fundamental role as much as the network.

Architecture & The Encoder

The reason of the model/network being called context encoder, as it consist of an encoder capturing the context of an image into a compact latent feature representation and a decoder which uses that representation to produce the missing image content. We can imagine the architecture somewhat close to autoencoders.

The architecture pipeline of context encoder is an encoder, a channel-wise fully-connected layer, then a decoder. The encoder reads the damaged image and compresses it into a latent feature map. The decoder expands that map back into pixels.

Originally paper's encoder comes from AlexNet, where the first five convolutional layers plus the pooling layer after them. A 227×227 input becomes a 6×6×256 representation. Unlike AlexNet, it is not trained for ImageNet classification. It is trained from scratch, with random weights, on the inpainting task alone.

Unlike a classic autoencoder, there is no narrow bottleneck. A normal autoencoder must reconstruct its whole input, so it needs a small code to avoid copying. Here the input has a hole and the target is the missing content, so the network is already forced to infer rather than copy. The paper's latent size is 6×6×256 on both the encoder and decoder side.

In our implemented 128×128 CelebA version, we keep the AlexNet flavor but resize the strides so the latent map is 8×8×256. Each conv is followed by BatchNorm and LeakyReLU. BatchNorm is our addition for stable training, not something the paper's text describes. Hence, the encoder architecture can be coded as,

class Encoder(nn.Module):
    def __init__(self):
        super().__init__()
        def block(cin, cout, k, s, p):
            return nn.Sequential(
                nn.Conv2d(cin, cout, k, s, p, bias=False),
                nn.BatchNorm2d(cout),
                nn.LeakyReLU(0.2, inplace=True))
        self.net = nn.Sequential(
            block(3,   64,  7, 2, 3),            # 128 -> 64
            nn.MaxPool2d(3, 2, 1),               # 64  -> 32
            block(64,  192, 5, 1, 2),            # 32
            nn.MaxPool2d(3, 2, 1),               # 32  -> 16
            block(192, 384, 3, 1, 1),            # 16
            block(384, 256, 3, 1, 1),            # 16
            block(256, 256, 3, 1, 1),            # 16  (conv5)
            nn.MaxPool2d(3, 2, 1),               # 16  -> 8  (pool5)
        )

    def forward(self, x):
        return self.net(x)

The Channel-Wise Fully Connected Layer

This layer is context encoders key architectural idea. It lets information travel across the whole image without an explosion in parameters. But why in the first place, did we really need this channel-wise fully connected layer?

The problem: Convolutional layers connect all feature maps together, but never connect distant locations within one feature map. Information in one corner cannot directly reach the opposite corner. For inpainting that is a real limit, because the pixels that best predict the center of a face may sit far from it. Fully-connected (inner product) layers solve this, since every activation connects to every other.

The catch is cost: Connecting a 256-channel feature map of size nxn to the decoder with a full FC layer needs m²n⁴ parameters (m channels, n×n maps). The paper notes this would exceed 100 million parameters, making training on GPUs of the time impractical. The fix is a fully-connected layer with groups, where each feature map gets its own n²×n² weight matrix, and nothing connects different feature maps. The parameter count drops from m²n⁴ to m·n⁴.

For our 8×8×256 latent this means 256 × 8⁴ ≈ 1.05 million weights, versus about 268 million for a full FC layer. Since this layer never mixes channels, the paper follows it with a stride-1 convolution to propagate information across channels. We use a 1×1 convolution for that step. In code, the layer is a single einsum over per-channel weight matrices. Hence, channel-wise fully connected layer implemented as,

class ChannelWiseFC(nn.Module):
    def __init__(self, channels, size):
        super().__init__()
        self.n2 = size * size
        self.weight = nn.Parameter(torch.empty(channels, self.n2, self.n2))
        self.bias = nn.Parameter(torch.zeros(channels, self.n2))
        bound = 1 / self.n2 ** 0.5
        nn.init.uniform_(self.weight, -bound, bound)

    def forward(self, x):
        B, m, n, _ = x.shape
        x = x.reshape(B, m, self.n2)
        x = torch.einsum('bmi,mij->bmj', x, self.weight)
        return (x + self.bias).reshape(B, m, n, n)

The Decoder

The decoder turns the processed latent map back into pixels. The encoder features pass through the channel-wise FC layer and then a series of five up-convolutional layers, each with learned filters and a ReLU activation.

An up-convolution is simply a convolution that produces a higher-resolution output. It can be understood two ways, upsampling followed by a convolution, or a convolution with fractional stride (a transposed convolution). Stacking up-convolutions with non-linearities builds a non-linear, learned upsampling of the encoder's features until the output reaches the target size.

class Decoder(nn.Module):
    def __init__(self, channels=256, size=8):
        super().__init__()
        self.cw_fc = ChannelWiseFC(channels, size)
        self.mix = nn.Conv2d(channels, channels, 1)       

        def up(cin, cout):                               
            return nn.Sequential(
                nn.ConvTranspose2d(cin, cout, 4, 2, 1),
                nn.BatchNorm2d(cout),
                nn.ReLU(inplace=True))
        self.up = nn.Sequential(
            up(256, 128),                                 # 8   -> 16
            up(128, 64),                                  # 16  -> 32
            up(64, 64),                                   # 32  -> 64
            up(64, 32),                                   # 64  -> 128
            nn.ConvTranspose2d(32, 3, 3, 1, 1),           # 5th up-conv, stride 1 -> RGB
            nn.Tanh(),                                    # matches [-1, 1] normalization
        )
    def forward(self, z):
        z = self.mix(self.cw_fc(z))
        return self.up(z)

Reconstruction Loss

The original paper train context encoders by regressing to the
ground truth content of the missing (dropped out) region. In practicality there are multiple ways to fill the missing region which are consistent with the context. The paper model this behaviour by having a decoupled joint loss function to handle both continuity within the context and multiple modes in the output. The first loss of it's kind is reconstruction loss.

The reconstruction loss is a normalized masked L2 distance. It is responsible for capturing the overall structure of the missing region and coherence with regards to its context, but tend to average together the multiple modes in prediction. Reconstruction loss compares the prediction with the ground truth, but only inside the hole.

The mask M zeroes out every pixel the network was shown, so the model gets no credit for copying context. The paper tried both L1 and L2 and found no significant difference. In our code we divide by the number of masked elements so the loss keeps the same scale across central, block and region masks that cover different areas. The following is the implementation of reconstruction loss,

class ReconstructionLoss(nn.Module):
    def forward(self, pred, target, mask):
        diff = mask * (target - pred)
        sq_err = diff.pow(2).sum(dim=(1, 2, 3))
        norm = mask.sum(dim=(1, 2, 3)) * pred.size(1) + 1e-8
        return (sq_err / norm).mean()

Reconstruction loss alone has a known weakness, it encourages a rough outline of the missing object but fails to capture high-frequency detail. The paper's explanation is that L2 (or L1) prefers a blurry answer. When several fillings are plausible, the safest prediction for L2 is the mean of the distribution, since it minimizes average pixel error. The mean of many sharp faces is a blurry face.

Adversarial Loss

To fix the blur of reconstruction loss, the paper adds a GAN-style adversarial loss. A discriminator D learns to tell real images from the context encoder's outputs, while the context encoder (acting as the generator G) tries to fool it. Producing a single sharp, realistic face then scores better than the blurry average, because D can spot the blur. The standard GAN objective is a minimax game between G and D as,

The paper adapts it with two deliberate choices,

Condition only the generator on context, not the discriminator. A conditional GAN would show D the context too, but D then easily exploits the visible discontinuity between generated and original regions to classify fakes. Conditioning only G avoids that shortcut.
No noise vector. The authors found results improved when the generator was not given noise. The context itself supplies enough variation.

Hence, the resulting objective for D becomes,

Because D sees the entire output of the context encoder, not just the hole, the objective encourages the whole output to look realistic. The following adversarial loss can be implemented as,

class AdversarialLoss(nn.Module):
    def __init__(self):
        super().__init__()
        self.bce = nn.BCEWithLogitsLoss()

    def d_loss(self, d_real, d_fake):
        return (self.bce(d_real, torch.ones_like(d_real)) +
                self.bce(d_fake, torch.zeros_like(d_fake)))

    def g_loss(self, d_fake):
        # non-saturating form: minimize -log D(F(...))
        return self.bce(d_fake, torch.ones_like(d_fake))

We use the common non-saturating generator loss (maximize log D of the fake) instead of minimizing log(1 - D), which gives stronger gradients early in training. This is a standard GAN practice rather than something stated in the paper's loss equations.

For training purpose, full objective is a weighted sum of the two losses. The reconstruction term dominates, and the adversarial term acts as a sharpening nudge. Hence,

The generator and discriminator are optimized jointly with alternating updates, one step on D using real images and detached fake outputs, then one step on G with the joint loss. The paper notes that adversarial training diverged for its AlexNet feature-learning experiments, so it uses the adversarial loss only for the inpainting experiments.

Masking, Training, Results

The original paper uses three masking strategy on images before feeding into context encoders. These are,

Central Region: The simplest such shape is the central square patch in the image. This is simple, but the network can overfit to the hole's location and learn features that do not generalize.

Random Block: Several small rectangles at random positions. Hole location varies, though the sharp rectangle edges still give the network a cue.

Random Region: Irregular blobs, drawn as random-walk brush strokes. Closest to real occlusions, and the paper reports it works best for feature learning.

In our implementation each mask covers at most about a quarter of the image. During training we pick a mask type at random for every batch.

class ContextEncoder(nn.Module):
    def __init__(self):
        super().__init__()
        self.encoder = Encoder()   # AlexNet-style, 3x128x128 -> 256x8x8
        self.decoder = Decoder()   # channel-wise FC -> 1x1 conv -> 5 up-convs

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

We train the overall context encoder model, with 20 EPOCHS. On inpainting, the paper shows that the reconstruction loss alone gives blurry fills, while adding the adversarial loss produces much sharper, more realistic completions. It also shows the features are useful, an encoder pretrained this way and then fine-tuned on PASCAL VOC gave competitive results among unsupervised pretraining methods of its time on classification, detection and segmentation, without any ImageNet labels.

Expect modest results from a small reproduction. With a few thousand CelebA images and a couple of dozen epochs, completions will usually be plausible but soft. Faces help, since they are highly structured, but fine texture (hair strands, skin detail) needs far more data and training. That reflects data and compute, not a bug in the architecture.

The approach has real limitations. Output sharpness depends on a delicate GAN balance. A single deterministic output cannot represent the many valid fillings of a large hole. And because the channel-wise FC layer is tied to a fixed latent size, the network cannot natively handle images of different resolutions.

The ideas carried forward. GAN-based inpainting methods added global and local discriminators, contextual attention and gated convolutions to handle larger and irregular holes. Masked autoencoders and related masked-modeling methods return to the same core bet, hide part of the input, predict it, and learn good representations as a side effect. Context encoders were an early, clear demonstration that this works for images. Following are the results of model training on central block masking.