PaperlayerView exact revision on arXiv

Source-linked AI summary

Colorization Transformer

Manoj Kumar, Dirk Weissenborn, Nal Kalchbrenner

arXiv:2102.04432v2cs.CVcs.AIcs.LG

TL;DR

Image colorization is inherently stochastic and requires semantic understanding, motivating models that can generate diverse plausible colors from grayscale inputs. ColTran uses a conditional autoregressive transformer for coarse coloring followed by two fully parallel upsamplers, and reports state-of-the-art ImageNet fidelity with human evaluators preferring the best of three generated colorings over ground truth in more than 60% of cases.

  • Problem

    Image colorization is a stochastic, semantically demanding one-to-many task requiring models to generate diverse plausible colorings from grayscale images.

  • Method

    ColTran sequentially predicts coarse low-resolution colorings with a conditional autoregressive transformer, then uses parallel color and spatial upsamplers to produce high-resolution images.

  • Results

    ColTran establishes a new state of the art on ImageNet using FID and Mechanical Turk evaluation, with more than 60% of cases favoring the highest-rated of three generated colorings over ground truth.

  • Takeaways & Limitations

    Self-attention can generate diverse, high-fidelity colorizations that are largely indistinguishable from ground truth to human raters.

  • Takeaways & Limitations

    Parallel upsampling uses argmax rather than sampling at each pixel, slightly limiting potential colorization diversity while reducing minor pixel-level inconsistencies.

Abstract

from arXiv · show

We present the Colorization Transformer, a novel approach for diverse high fidelity image colorization based on self-attention. Given a grayscale image, the colorization proceeds in three steps. We first use a conditional autoregressive transformer to produce a low resolution coarse coloring of the grayscale image. Our architecture adopts conditional transformer layers to effectively condition grayscale input. Two subsequent fully parallel networks upsample the coarse colored low resolution image into a finely colored high resolution image. Sampling from the Colorization Transformer produces diverse colorings whose fidelity outperforms the previous state-of-the-art on colorising ImageNet based on FID results and based on a human evaluation in a Mechanical Turk test. Remarkably, in more than 60% of cases human evaluators prefer the highest rated among three generated colorings over the ground truth. The code and pre-trained checkpoints for Colorization Transformer are publicly available at https://github.com/google-research/google-research/tree/master/coltran

1 INTRODUCTION

Colorization is a stochastic, semantically demanding task, and ColTran applies self-attention in a multistage design for high-resolution diverse colorization. The model combines coarse autoregressive coloring with parallel upsampling and reports state-of-the-art ImageNet results.

  • Motivation: Image colorization requires semantic scene understanding and world knowledge because it is inherently stochastic.The task supports applications ranging from restoring grayscale media to representation learning and unsupervised data augmentation.
  • Motivation: Probabilistic generative models address colorization’s one-to-many nature by producing multiple diverse colorings from a learned distribution.They improve over earlier deterministic approaches while preserving sampling-based diversity.
  • Approach: ColTran is a probabilistic colorization model composed only of axial self-attention blocks.Axial attention captures a global receptive field with two layers and lower O(D) rather than O(D^2) complexity, while supporting efficient accelerator implementation.
  • Approach: The pipeline decomposes high-resolution colorization into coarse low-resolution autoregressive coloring followed by parallel color and spatial super-resolution.This decomposition simplifies high-resolution generation and enables fully parallel processing in the upsampling stages.
  • Results: More than 60% of Mechanical Turk evaluations preferred the highest-rated of three generated colorings over the ground truth.The paper reports this result alongside a new state of the art on ImageNet using FID and 2AFC human evaluation.

2 RELATED WORK

Prior colorization methods span interactive, exemplar-based, and generative approaches. Generative methods condition diverse color predictions on grayscale inputs, with ColTran most closely related to autoregressive models and multistage upsampling methods.

  • Interactive and exemplar-based methods: Early colorization methods relied on human-provided scribbles or exemplar images from which colors were copied.These approaches require either user hints or a reference source image.
  • Generative colorization: Generative colorization extends unconditional image-generation models with grayscale conditioning using normalizing flows, VAEs, or GANs.Examples include cINN, VAE-MDN, SCC-DC, and cGAN.
  • Relation to ColTran: ColTran is most closely related to autoregressive approaches such as PixColor and PIC.Its multistage generation process also relates to Subscale Pixel Networks, while differing in representation, ordering, and use of fully parallel networks.

3 BACKGROUND: AXIAL TRANSFORMER

Axial Transformer models pixels autoregressively with alternating masked row and column attention, efficiently summarizing prior pixels. Its encoder supplies conditioning from prior channels, while ColTran extends the design with grayscale conditioning and staged sampling.

  • Axial attention: Self-attention alternates row and column operations to model image dependencies without global attention across every pixel.This preserves broad receptive-field modeling while avoiding the full quadratic pixel-to-pixel operation.
  • Autoregressive modeling: Axial Transformer uses masked row and column attention to model each pixel from preceding pixels in raster order.Causal masking removes attention to future positions by setting A_m,n = 0 when n > m.
  • Decoder: The outer decoder summarizes previous rows, shifts its output downward, and produces context containing only prior-row information.Its computation is represented by e = Embeddings(x), s_o = MaskedColumn(Row(e)) × N, and o = ShiftDown(s_o).
  • Decoder: The inner decoder shifts embeddings right to mask the current pixel, combines them with outer context, and predicts p(x_i,j) from masked row attention.The resulting state contains information from prior rows and earlier pixels in the same row.
  • Channel conditioning: For multichannel images, an encoder independently processes prior channels and sums their outputs into a conditioning context.That context biases both the outer and inner decoder inputs.
  • ColTran extension: ColTran’s architecture contains an autoregressive colorizer, a color upsampler, and a spatial upsampler optimized independently.The upsamplers use row/column attention and deterministically increase color and spatial resolution in parallel.
  • Sampling: Sampling runs the encoder once per channel, the outer decoder once per row, and the inner decoder once per pixel.This semi-parallel procedure avoids reevaluating the entire network for every generated pixel.

4 PROPOSED ARCHITECTURE

ColTran decomposes stochastic image colorization into coarse autoregressive coloring followed by parallel color and spatial upsampling. Its conditional transformer layers incorporate grayscale context throughout the core, while auxiliary prediction and parallel upsamplers support high-fidelity generation.

  • Overall architecture: ColTran transforms a grayscale image into a colored image through intermediate low-resolution representations, simplifying high-resolution colorization and enabling larger models.The intermediate representations are spatially downsampled and use 3-bit color channels, yielding 512 coarse colors per pixel.
  • Overall architecture: The architecture factors colorization into a coarse-color model, a parallel color upsampler, and a parallel spatial upsampler.These components model p(xs↓c↓|xg), p(xs↓|xs↓c↓, xg), and p(x|xs↓), respectively.
  • ColTran core: ColTran Core predicts one of 512 coarse colors per pixel autoregressively, conditioned on the low-resolution grayscale image and previously predicted pixels.The conditional distribution follows raster order.
  • ColTran core: Conditional transformer layers apply grayscale-derived conditioning through conditional attention, MLP, and layer-normalization components.Conditioning includes scale and shift operations for q, k, and v, conditional MLP outputs, and context-dependent normalization parameters.
  • ColTran core: Conditioning every layer improves encoder contextual representations, and ColTran outperforms native Axial Transformer baselines in the reported experiments.The paper attributes this to stronger gradient signals through the encoder.
  • Upsampling and objective: An auxiliary parallel prediction model regularizes the grayscale encoder, while color and spatial upsamplers predict per-channel 256-level intensities in parallel.Parallel upsampling uses argmax predictions to reduce minor pixel-level color inconsistencies, with a slight limitation on diversity.
  • Upsampling and objective: The architecture is trained with a negative log-likelihood combining autoregressive, auxiliary parallel, color-upsampling, and spatial-upsampling terms.λ controls the relative contribution of the autoregressive and auxiliary parallel predictions.

5 EXPERIMENTS

Experiments show that ColTran’s conditional components and auxiliary parallel model improve coarse colorization, while its upsamplers achieve high-fidelity results and its full system outperforms prior methods on FID and human evaluation.

  • ColTran core ablations: ColTran-B and its 2x and 4x wider MLP baselines are outperformed by ColTran, showing conditional layers improve loss beyond parameter-count increases alone.The comparison uses wider baselines with increased parameter counts to test whether gains arise solely from model size.
  • ColTran core ablations: cAtt contributes most among the conditional components, although removing cLN, cMLP, or cAtt each worsens final performance.This conclusion comes from a leave-one-out retraining study.
  • ColTran core ablations: Multiplicative conditional interactions have a much stronger impact than additive interactions, although both improve performance.The comparison evaluates architectures using conditional scales only versus conditional shifts only.
  • Other ablations: At λ = 0.01, the auxiliary parallel model reaches better coarse-sample FID faster and achieves a marginal but consistent final improvement; higher λ values degrade performance.FID, rather than the holdout metric, is used to select λ because the parallel model is harder to optimize independently.
  • Upsamplers: 16.4 FID is obtained by the upsamplers on fine-grained 256 × 256 images, compared with 15.5 FID between two random holdout-sample sets.The upsamplers receive ground-truth 64×64 coarse images and produce images without visible artifacts.
  • Frechet Inception Distance: 19.37 FID is achieved by ColTran on ImageNet, improving over the previous best FID of 24.32 and outperforming ColTran-S at 22.06 and ColTran-B at 19.98.FID is computed from colorizations of 5000 grayscale 256×256 ImageNet validation images without overlap with the ground-truth distribution.
  • Other ablations: A 0.57 correlation is observed between log-likelihood and FID, with absolute log-likelihood improvements of 0.01–0.02 sometimes producing significant FID gains.The correlation is measured across architectural variants after 150K training steps.
  • Qualitative evaluation: 42.3% is ColTran’s best-seed fooling rate versus 35.4% for PixColor, while ColTran Oracle reaches 62%.The oracle selects the best of three colorizations using human ratings.

6 CONCLUSION

ColTran uses conditional transformer layers and self-attention to generate diverse, high-fidelity ImageNet colorizations. The authors report that its outputs are largely indistinguishable from ground truth to human raters, with implementation resources publicly available.

  • ColTran entirely relies on self-attention and introduces conditional transformer layers for generative image colorization.The authors report that ablations show this mechanism outperforms several baselines.
  • ColTran generates diverse, high-fidelity colorizations on ImageNet that are largely indistinguishable from ground truth to human raters.
  • The implementation, hyperparameters, and pre-trained ImageNet checkpoints for the colorizer and upsamplers are publicly available.
  • EMA checkpoints were extremely crucial for generating high-quality samples, with FID improving steadily over training when EMA was applied.

C NUMBER OF PARAMETERS AND INFERENCE SPEED

ColTran's core samples 64×64 grayscale images in minutes, while its parallel upsamplers operate in milliseconds. Its 74M-parameter model is smaller than the compared PixColor system, and reuse of activations could further accelerate sampling.

  • A batch of 20 64×64 grayscale images takes around 3.5–5 minutes to sample on a P100 GPU, while upsamplers take milliseconds.Sampling 28×28 colorizations takes around 30 seconds in the reported comparison.
  • The naive implementation recomputes decoder activations for every generated pixel.The authors suggest reusing encoder- and row-level activations as future engineering work to speed sampling further.
  • 74M parameters comprise the ColTran core, color upsampler, and spatial upsampler, compared with 83M for PixColor.ColTran allocates 46M, 14M, and 14M parameters to these components, respectively.

D LOWER COMPUTE REGIME

In a lower-compute configuration, reducing hardware and batch sizes degraded FID but retained performance better than the other compared models. Top-K sampling improved fidelity at the cost of diversity.

  • FID dropped from 19.71 to 20.9 in the lower-compute retraining, remaining significantly better than the other models.The reduced setup used four TPUv2 chips and smaller configurations, including a two-block spatial upsampler.
  • Top-K sampling improves colorization fidelity and removes artifacts from unnatural colors, but reduces colorization diversity.The reported settings include K = 4 and K = 8 versus the baseline K = 512.

F ADDITIONAL ABLATIONS:

The additional materials describe architectural ablations, autoregressive generation, attention equations, cross-dataset demonstrations, and a small baseline hyperparameter sweep. They also document that the ImageNet-trained model was applied to LSUN and Celeb-A without fine-tuning.

  • Additional ablations: Additional conditional-transformer ablations using gated layers or global conditioning did not help.
  • Autoregressive generation: Autoregressive models factorize a joint distribution into conditionals, train with masked future symbols, and generate sequentially.For images, symbols typically represent RGB pixel-channels generated in raster-scan order.
  • Row/column self-attention: The row self-attention block applies attention independently across each column, followed by residual connections and a feed-forward MLP with ReLU activation.
  • Cross-dataset evaluation: The ImageNet-trained model colorizes 256×256 LSUN and 64×64 Celeb-A grayscale images without fine-tuning on either dataset.
  • Visualizations: Figure 10 orders generated colorizations from lower to higher fooling rate, with ground truth shown beneath them.
  • Hyperparameter sweep: A small sweep of baseline axial-transformer configurations was fixed after selecting an optimal configuration for later architecture design.The tested settings varied hidden size and block count.

K ANALYSIS OF MTURK RATINGS

The MTurk analysis organizes three-sample colorizations into diverse successes, realism gradients, and failure cases, while probability maps provide a per-pixel uncertainty proxy.

  • Figure 11 compares ground truth with three samples in columns showing diverse results, increasing realism, and failure cases.
  • Per-pixel maximum predicted probability over 512 colors is displayed as a proxy for uncertainty.
  • More than 60% fool rates occur for all samples in the diverse examples, including variation in high-level structure and low-level details.
  • Samples for complex objects can differ by 80% between the best and worst MTurk ratings, indicating variable success across attempts.
  • Failure cases have a 0% fool rate for all samples and involve highly complex structures that the model cannot colorize reliably.

M MORE SAMPLES

Additional, non-cherry-picked samples demonstrate a wide diversity of colorizations produced by ColTran.

  • The additional ColTran samples were not cherry-picked.
  • The displayed samples cover a wide diversity of colorizations.
  • These examples provide further visual evidence of ColTran's varied colorization outputs.
Loading 2102.04432v2…