C OLORIZATION T RANSFORMER
Manoj Kumar, Dirk Weissenborn & Nal Kalchbrenner
Google Research, Brain Team {mechcoder,diwe,nalk}@google.com
A BSTRACT
We present the Colorization Transformer, a novel approach for diverse high ﬁdelity image colorization based on self-attention.
Given a grayscale image, the colorization proceeds in three steps.
We ﬁrst 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 ﬁnely colored high resolution image.
Sampling from the Colorization
Transformer produces diverse colorings whose ﬁdelity 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 this url.
I NTRODUCTION
Image colorization is a challenging, inherently stochastic task that requires a semantic understanding of the scene as well as knowledge of the world.
Core immediate applications of the technique include producing organic new colorizations of existing image and video content as well as giving life to originally grayscale media, such as old archival images (Tsaftaris et al., 2014), videos (Geshwind, 1986) and black-and-white cartoons (S`ykora et al., 2004;
Qu et al., 2006; Cinarel & Zhang, 2017).
Colorization also has important technical uses as a way to learn meaningful representations without explicit supervision (Zhang et al., 2016; Larsson et al., 2016; Vondrick et al., 2018) or as an unsupervised data augmentation technique, whereby diverse semantics-preserving colorizations of labelled images are produced with a colorization model trained on a potentially much larger set of unlabelled images.
The current state-of-the-art in automated colorization are neural generative approaches based on log-likelihood estimation (Guadarrama et al., 2017; Royer et al., 2017; Ardizzone et al., 2019).
Probabilistic models are a natural ﬁt for the one-to-many task of image colorization and obtain better results than earlier determinisitic approaches avoiding some of the persistent pitfalls (Zhang et al., 2016).
Probabilistic models also have the central advantage of producing multiple diverse colorings that are sampled from the learnt distribution.
In this paper, we introduce the Colorization Transformer (ColTran), a probabilistic colorization model composed only of axial self-attention blocks (Ho et al., 2019b; Wang et al., 2020).
The main advantages of axial self-attention blocks are the ability to capture a global receptive ﬁeld with only two layers and O ( D √ D ) instead of O ( D 2 ) complexity.
They can be implemented efﬁciently using matrix-multiplications on modern accelerators such as TPUs (Jouppi et al., 2017).
In order to enable colorization of high-resolution grayscale images, we decompose the task into three simpler sequential subtasks: coarse low resolution autoregressive colorization, parallel color and spatial super-resolution.
For coarse low resolution colorization, we apply a conditional variant of Axial Transformer (Ho et al., 2019b), a state-of-the-art autoregressive image generation model that does not require custom kernels (Child et al., 2019).
While Axial Transformers support conditioning by biasing the input, we ﬁnd that directly conditioning the transformer layers can improve results signiﬁcantly.
Finally, by leveraging the semi-parallel sampling mechanism of Axial Transformers we are able to colorize images faster at higher resolution than previous work (Guadarrama et al., 2017) and as an effect this results in improved colorization ﬁdelity.
Finally, we employ fast parallel deterministic upsampling models to super-resolve the coarsely colorized image into the ﬁnal high resolution output.
In summary, our main contributions are:
• First application of transformers for high-resolution ( 256 × 256 ) image colorization.
• We introduce conditional transformer layers for low-resolution coarse colorization in Section 4.1.
The conditional layers incorporate conditioning information via multiple learnable components that are applied per-pixel and per-channel.
We validate the contribution of each component with extensive experimentation and ablation studies.
• We propose training an auxiliary parallel prediction model jointly with the low resolution coarse colorization model in Section 4.2.
Improved FID scores demonstrate the usefulness of this auxiliary model.
• We establish a new state-of-the-art on image colorization outperforming prior methods by a large margin on FID scores and a 2-Alternative Forced Choice (2AFC) Mechanical Turk test.
Remarkably, in more than 60% of cases human evaluators prefer the highest rated among three generated colorings over the ground truth.
R ELATED WORK
Colorization methods have initially relied on human-in-the-loop approaches to provide hints in the form of scribbles (Levin et al., 2004;
Ironi et al., 2005; Huang et al., 2005; Yatziv & Sapiro, 2006;
Qu et al., 2006;
Luan et al., 2007; Tsaftaris et al., 2014;
Zhang et al., 2017; Ci et al., 2018) and exemplar-based techniques that involve identifying a reference source image to copy colors from (Reinhard et al., 2001;
Welsh et al., 2002; Tai et al., 2005; Ironi et al., 2005; Pitié et al., 2007;
Morimoto et al., 2009; Gupta et al., 2012; Xiao et al., 2020).
Exemplar based techniques have been recently extended to video as well (Zhang et al., 2019a).
In the past few years, the focus has moved on to more automated, neural colorization methods.
The deterministic colorization techniques such as CIC (Zhang et al., 2016), LRAC (Larsson et al., 2016), LTBC (Iizuka et al., 2016), Pix2Pix (Isola et al., 2017) and DC (Cheng et al., 2015; Dahl, 2016) involve variations of CNNs to model per-pixel color information conditioned on the intensity.
Generative colorization models typically extend unconditional image generation models to incorporate conditioning information from a grayscale image.
Speciﬁcally, cINN (Ardizzone et al., 2019) use conditional normalizing ﬂows (Dinh et al., 2014), VAE-MDN (Deshpande et al., 2017; 2015) and
SCC-DC (Messaoud et al., 2018) use conditional VAEs (Kingma & Welling, 2013), and cGAN (Cao et al., 2017) use GANs (Goodfellow et al., 2014) for generative colorization.
Most closely related to
ColTran are other autoregressive approaches such as PixColor (Guadarrama et al., 2017) and PIC (Royer et al., 2017) with PixColor obtaining slightly better results than PIC due to its CNN-based upsampling strategy.
ColTran is similar to PixColor in the usage of an autoregressive model for low resolution colorization and parallel spatial upsampling.
ColTran differs from PixColor in the following ways.
We train ColTran in a completely unsupervised fashion, while the conditioning network in PixColor requires pre-training with an object detection network that provides substantial semantic information.
PixColor relies on PixelCNN (Oord et al., 2016) that requires a large depth to model interactions between all pixels.
ColTran relies on Axial Transformer (Ho et al., 2019b) and can model all interactions between pixels with just 2 layers.
PixColor uses different architectures for conditioning, colorization and super-resolution, while ColTran is conceptually simpler as we use self-attention blocks everywhere for both colorization and superresolution.
Finally, we train our autoregressive model on a single coarse channel and a separate color upsampling network that improves ﬁdelity (See: 5.3).
The multi-stage generation process in ColTran that upsamples in depth and in size is related to that used in Subscale Pixel Networks (Menick & Kalchbrenner, 2018) for image generation, with differences in the order and representation of bits as well as in the use of fully parallel networks.
The self-attention blocks that are the building blocks of ColTran were initially developed for machine translation (Vaswani et al., 2017), but are now widely used in a number of other applications including density estimation (Parmar et al., 2018; Child et al., 2019; Ho et al., 2019a; Weissenborn et al., 2019) and GANs (Zhang et al., 2019b)
B ACKGROUND : A XIAL T RANSFORMER 3. R OW AND COLUMN SELF - ATTENTION
Self-attention (SA) has become a standard building block in many neural architectures.
Although the complexity of self-attention is quadratic with the number of input elements (here pixels), it has become quite popular for image modeling recently (Parmar et al., 2018; Weissenborn et al., 2019) due to modeling innovations that don’t require running global self-attention between all pixels.
Following the work of (Ho et al., 2019b)
we employ standard qkv self-attention (Vaswani et al., 2017) within rows and columns of an image.
By alternating row- and column self-attention we effectively allow global exchange of information between all pixel positions.
For the sake of brevity we omit the exact equations for multihead self-attention and refer the interested reader to the Appendix H for more details.
Row/column attention layers are the core components of our model.
We use them in the autoregressive colorizer, the spatial upsampler and the color upsampler.
3. A XIAL T RANSFORMER
Ths Axial Transformer (Ho et al., 2019b) is an autoregressive model that applies (masked) row- and column self-attention operations in a way that efﬁciently summarizes all past information x i,<j and x <i, · to model a distribution over pixel x i,j at position i, j .
Causal masking is employed by setting all A m,n = 0 where n > m during self-attention (see Eq. 15).
Outer decoder.
The outer decoder computes a state s o over all previous rows x ≤ i, · by applying N layers of full row self-attention followed by masked column self-attention.
(Eq 2).
s o is shifted down by a single row, such that the output context o i,j at position i, j only contains information about pixels x <i, · from prior rows.
(Eq 3)
[Equation]
Inner decoder.
The embeddings to the inner decoder are shifted right by a single column to mask the current pixel x i,j .
The context o from the outer decoder conditions the inner decoder by biasing the shifted embeddings.
It then computes a ﬁnal state h , by applying N layers of masked row-wise self-attention to infuse additional information from prior pixels of the same row x
i,<j (Eq 4).
h i,j comprises information about all past pixels x <i and x i,<j .
A dense layer projects h into a distribution p ( x ij ) over the pixel at position ( i, j ) conditioned on all previous pixels x i,<j and x <i, · .
[Equation]
Encoder.
As shown above, the outer and inner decoder operate on 2-D inputs, such as a single channel of an image.
For multi-channel RGB images, when modeling the "current channel", the
Axial Transformer incorporates information from prior channels of an image (as per raster order) with an encoder.
The encoder encodes each prior channel independently with a stack of unmasked row/column attention layers.
The encoder outputs across all prior channels are summed to output a conditioning context c for the "current channel".
The context conditions the outer and inner decoder by biasing the inputs in Eq 1 and Eq 4 respectively.
[Figure]
Sampling.
The Axial Transformer natively supports semi-parallel sampling that avoids reevaluation of the entire network to generate each pixel of a RGB image.
The encoder is run once per-channel, the outer decoder is run once per-row and the inner decoder is run once per-pixel.
The context from the outer decoder and the encoder is initially zero.
The encoder conditions the outer decoder (Eq 1) and the encoder + outer decoder condition the inner decoder (Eq 4).
The inner decoder then generates a row, one pixel at a time via Eqs.
(4) to (6). After generating all pixels in a row, the outer decoder recomputes context via Eqs. (1) to (3) and the inner decoder generates the next row.
This proceeds till all the pixels in a channel are generated. The encoder, then recomputes context to generate the next channel.
P ROPOSED A RCHITECTURE
Image colorization is the task of transforming a grayscale image x
g ∈ R H
× W × 1 into a colored image x ∈ R H
× W × 3 .
The task is inherently stochastic; for a given grayscale image x g , there exists a conditional distribution over x , p ( x | x g ) .
Instead of predicting x directly from x g , we instead sequentially predict two intermediate low resolution images x s ↓ and x s ↓ c ↓ with different color depth ﬁrst.
Besides simplifying the task of high-resolution image colorization into simpler tasks, the smaller resolution allows for training larger models.
We obtain x s ↓ , a spatially downsampled representation of x , by standard area interpolation.
x s ↓ c ↓ is a 3 bit per-channel representation of x s ↓ , that is, each color channel has only 8 intensities.
Thus, there are 8 3 = 512 coarse colors per pixel which are predicted directly as a single “color” channel.
We rewrite the conditional likelihood p ( x | x g ) to incorporate the intermediate representations as follows:
[Equation]
ColTran core (Section 4.1), a parallel color upsampler and a parallel spatial upsampler (Section 4.3) model p ( x s ↓ c ↓ | x g ) , p ( x s ↓ | x s ↓ [Equation] ↓ , x g ) and p ( x | x s ↓ ) respectively.
In the subsections below, we describe
[Table]
y =
Softmax ( qk (cid:62) √ D ) v y =
Softmax ( q c k (cid:62) c √ D ) v c where ∀ z
= k , q , v z c =
( c U zs ) (cid:12) z
+ ( c U zb )
MLP y
=
ReLU ( x U 1 + b 1 ) U
+ bh = ReLU
( x U 1 + b 1 ) U
+ by =
( c U fs ) (cid:12) h
+ ( c U fb )
[Table]
y =
[Table]
Norm ( x ) + γ y = β c Norm ( x ) + γ c where ∀ µ = β c , γ c c ∈ R H × W × D → ˆ c ∈ R HW
× D µ = ( u · ˆ c ) U µd u ∈ R
HW these individual components in detail.
From now on we will refer to all low resolutions as M × N and high resolution as H × W .
An illustration of the overall architecture is shown in Figure 2.
4. C OL T RAN C ORE
In this section, we describe ColTran core, a conditional variant of the Axial Transformer (Ho et al., 2019b) for low resolution coarse colorization.
ColTran Core models a distribution p c ( x s ↓ c ↓ | x g ) over 512 coarse colors for every pixel, conditioned on a low resolution grayscale image in addition to the colors from previously predicted pixels as per raster order (Eq. 9).
[Equation]
Given a context representation c ∈ R M × N × D we propose conditional transformer layers in Table 1.
Conditional transformer layers have conditional versions of all components within the standard attention block (see Appendix H, Eqs.
Conditional Self-Attention.
For every layer in the decoder, we apply six 1 × 1 convolutions to c to obtain three scale and shift vectors which we apply element-wise to q , k and v of the self-attention operation (Appendix 3.1), respectively.
Conditional MLP.
A standard component of the transformer architecture is a two layer pointwise feed-forward network after the self-attention layer.
We scale and shift to the output of each MLP conditioned on c as for self-attention.
Conditional Layer Norm.
Layer normalization (Ba et al., 2016) globally scales and shifts a given normalized input using learnable vectors β , γ .
Instead, we predict β c and γ c as a function of c .
We ﬁrst aggregate c into a global 1-D representation c ∈ R L via a learnable, spatial pooling layer.
Spatial pooling is initialized as a mean pooling layer.
Similar to 1-D conditional normalization layers (Perez et al., 2017; De Vries et al., 2017; Dumoulin et al., 2016; Huang & Belongie, 2017), we then apply a linear projection on c to predict β c and γ c , respectively.
A grayscale encoder consisting of multiple, alternating row and column self-attention layers encodes the grayscale image into the initial conditioning context c g .
It serves as both context for the conditional layers and as additional input to the embeddings of the outer decoder.
The sum of the outer decoder’s output and c g condition the inner decoder.
Figure 2 illustrates how conditioning is applied in the autoregressive core of the ColTran architecture.
Conditioning every layer via multiple components allows stronger gradient signals through the encoder and as an effect the encoder can learn better contextual representations.
We validate this empirically by outperforming the native Axial Transformer that conditions context states by biasing (See Section 5.2 and Section 5.4).
4. A UXILIARY PARALLEL MODEL
We additionally found it beneﬁcial to train an auxiliary parallel prediction model that models (cid:101)
p c
( x s ↓ c ↓ ) directly on top of representations learned by the grayscale encoder which we found beneﬁcial for regularization (Eq. 10)
[Equation]
Intuitively, this forces the model to compute richer representations and global color structure already at the output of the encoder which can help conditioning and therefore has a beneﬁcial, regularizing effect on learning.
We apply a linear projection, U parallel ∈ R L × 512 on top of c g (the output of the grayscale encoder) into a per-pixel distribution over 512 coarse colors.
It was crucial to tune the relative contribution of the autoregressive and parallel predictions to improve performance which we study in Section 5.
4. C OLOR & S PATIAL U PSAMPLING
In order to produce high-ﬁdelity colorized images from low resolution, coarse color images and a given high resolution grayscale image, we train color and spatial upsampling models.
They share the same architecture while differing in their respective inputs and resolution at which they operate.
Similar to the grayscale encoder, the upsamplers comprise of multiple alternating layers of row and column self-attention.
The output of the encoder is projected to compute the logits underlying the per pixel color probabilities of the respective upsampler.
Figure 2 illustrates the architectures
Color Upsampler.
We convert the coarse image x s ↓ c ↓ ∈ R M × N × 1 of 512 colors back into a 3 bit
RGB image with 8 symbols per channel.
The channels are embedded using separate embedding matrices to x s ↓ c
↓ k ∈ R M
×
N × D , where k ∈ { R, G, B } indicates the channel.
We upsample each channel individually conditioning only on the respective channel’s embedding.
The channel embedding is summed with the respective grayscale embedding for each pixel and serve as input to the subsequent self-attention layers (encoder).
The output of the encoder is further projected to per pixel-channel probability distributions (cid:101)
p c ↑ ( x s ↓ k | x s ↓ c ↓ , x g ) ∈
R M × N × 256 over 256 color intensities for all k ∈
{ R, G, B } (Eq. 11).
[Equation]
Spatial Upsampler.
We ﬁrst naively upsample x s ↓ ∈ R M × N × 3 into a blurry, high-resolution RGB image using area interpolation.
As above, we then embed each channel of the blurry RGB image and run a per-channel encoder exactly the same way as with the color upsampler.
The output of the encoder is ﬁnally projected to per pixel-channel probability distributions (cid:101) p s ↑ ( x k | x s ↓ ,
x g ) ∈
R H
× W × 256 over 256 color intensities for all k ∈
R, G, B } .
(Eq. 12)
[Equation] In our experiments, similar to (Guadarrama et al., 2017), we found parallel upsampling to be sufﬁcient for high quality colorizations.
Parallel upsampling has the huge advantage of fast generation which would be notoriously slow for full autoregressive models on high resolution.
To avoid plausible minor color inconsistencies between pixels, instead of sampling each pixel from the predicted distribution in (Eq. 12 and Eq. 11), we just use the argmax.
Even though this slightly limits the potential diversity of colorizations, in practice we observe that sampling only coarse colors via ColTran core is enough to produce a great variety of colorizations.
Objective.
We train our architecture to minimize the negative log-likelihood (Eq. 13) of the data.
p c / (cid:101) p c , (cid:101) p s ↑ , (cid:101) p c ↑ are maximized independently and λ is a hyperparameter that controls the relative contribution of p c and (cid:101)
p c
[Equation]
E XPERIMENTS 5. T RAINING AND E VALUATION
We evaluate ColTran on colorizing 256 × 256 grayscale images from the ImageNet dataset (Russakovsky et al., 2015).
We train the ColTran core, color and spatial upsamplers independently on 16 TPUv2 chips with a batch-size of 224 , 768 and 32 for 450K, 300K and 150K steps respectively.
We use 4 axial attention blocks in each component of our architecture, with a hidden size ofand 4 heads.
We use RMSprop (Tieleman & Hinton, 2012) with a ﬁxed learning rate of 3 e − 4 .
We set apart 10000 images from the training set as a holdout set to tune hyperparameters and perform ablations.
To compute FID, we generate 5000 samples conditioned on the grayscale images from this holdout set.
We use the public validation set to display qualitative results and report ﬁnal numbers.
5. A BLATIONS OF C OL T RAN C ORE
The autoregressive core of ColTran models downsampled, coarse-colored images of resolution 64 ×with 512 coarse colots, conditioned on the respective grayscale image.
In a series of experiments we ablate the different components of the architecture (Figure 3).
In the section below, we refer to the conditional self-attention, conditional layer norm and conditional MLP subcomponents as cAtt, cLN and cMLP respectively.
We report the per-pixel log-likelihood over 512 coarse colors on the validation set as a function of training steps.
Impact of conditional transformer layers.
The left side of Figure 3 illustrates the signiﬁcant improvement in loss that ColTran core (with conditional transformer layers) achieves over the original
Axial Transformer (marked ColTran-B ).
This demonstrates the usefulness of our proposed conditional layers.
Because conditional layers introduce a higher number of parameters we additionally compare to and outperform the original Axial Transformer baselines with 2x and 4x wider MLP dimensions (labeled as ColTran-B 2x and ColTran-B 4x ).
Both ColTran-B 2x and ColTran-B 4x have an increased parameter count which makes for a fair comparison.
Our results show that the increased performance cannot be explained solely by the fact that our model has more parameters.
Importance of each conditional component.
We perform a leave-one-out study to determine the importance of each conditional component.
We remove each conditional component one at a time and retrain the new ablated model.
The curves no cLN , no cMLP and no cAtt in the middle of Figurequantiﬁes our results.
While each conditional component improves ﬁnal performance, cAtt plays the most important role.
Multiplicative vs Additive Interactions.
Conditional transformer layers employ both conditional shifts and scales consisting of additive and multiplicative interactions, respectively.
The curves Scale and Shift on the right hand side of Figure 3 demonstrate the impact of these interactions via ablated architectures that use conditional shifts and conditional scales only.
While both types of interactions are important, multiplicative interactions have a much stronger impact.
Context-aware dot product attention.
Self-attention computes the similarity between pixel representations using a dot product between q and k (See: Eq 15).
cAtt applies conditional shifts and scales on q , k and allow modifying this similarity based on contextual information.
The curve cAtt, only v on the right of Figure 3 shows that removing this property, by conditioning only on v leads to worse results.
Fixed vs adaptive global representation: cLN aggregates global information with a ﬂexible learnable spatial pooling layer.
We experimented with a ﬁxed mean pooling layer forcing all the cLN layers to use the same global representation with the same per-pixel weight.
The curve cLN, mean pool on the right of Figure 3 shows that enforcing this constraint causes inferior performance as compared to even having no cLN.
This indicates that different aggregations of global representations are important for different cLN layers.
5. O THER ABLATIONS
Auxiliary Parallel Model.
We study the effect of the hyperparameter λ , which controls the contribution of the auxiliary parallel prediction model described in Section 4.2.
For a given λ , we now optimize ˆ p c ( λ ) =
(1 − λ ) log p c ( . )
+ λ log (cid:101)
p c ( . ) instead of just log p c ( . ) .
Note that (cid:101) p c ( . ) , models each pixel independently, which is more difﬁcult than modelling each pixel conditioned on previous pixels given by p c ( . ) .
Hence, employing ˆ p c ( λ ) as a holdout metric, would just lead to a trivial soluion at λ = 0 .
Instead, the FID of the generated coarse 64x64 samples provides a reliable way to ﬁnd an optimal value of λ .
In Figure 4, at λ = 0 .
01 , our model converges to a better FID faster with a marginal but consistent ﬁnal improvement.
At higher values the performance deteriorates quickly.
Upsamplers.
Upsampling coarse colored, low-resolution images to a higher resolution is much simpler.
Given ground truth 64 × 64 coarse images, the ColTran upsamplers map these to ﬁne grained 256 × 256 images without any visible artifacts and FID of 16.4.
For comparison, the FID between two random sets of 5000 samples from our holdout set is 15.5.
It is further extremely important to provide the grayscale image as input to each of the individual upsamplers, without which the generated images appear highly smoothed out and the FID drops to 27.0.
We also trained a single upsampler for both color and resolution.
The FID in this case drops marginally to 16.6.
5. F RECHET I NCEPTION D ISTANCE
We compute FID using colorizations of 5000 grayscale images of resolution 256 × 256 from the
ImageNet validation set as done in (Ardizzone et al., 2019).
To compute the FID, we ensure that there is no overlap between the grayscale images that condition ColTran and those in the ground-truth distribution.
In addition to ColTran, we report two additional results ColTran-S and ColTran-Baseline .
ColTran-B refers to the baseline Axial Transformer that conditions via addition at the input.
PixColor samples smaller 28 × 28 colored images autoregressively as compared to ColTran’s 64 × 64.
As a control experiment, we train an autoregressive model on resolution 28 × 28 ( ColTran-S ) to disentangle architectural choices and the inherent stochasticity of modelling higher resolution images.
ColTran-S and ColTran-B obtains FID scores of 21.9 and 21.6 that signiﬁcantly improve over the previous best FID of 24.32.
Finally, ColTran achieves the best FID score of 19.71.
All results are presented in Table 2 left. [Table]
Correlation between FID and Log-likelihood.
For each architectural variant, Figure 4 right illustrates the correlation between the log-likelihood and FID after 150K training steps.
There is a moderately positive correlation of 0.57 between the log-likelihood and FID.
Importantly, even an absolute improvement on the order of 0.01 - 0.02 can improve FID signiﬁcantly.
This suggests that designing architectures that achieve better log-likelihood values is likely to lead to improved FID scores and colorization ﬁdelity.
5. Q UALITATIVE E VALUATION
Human Evaluation.
For our qualitative assessment, we follow the protocol used in PixColor (Guadarrama et al., 2017).
ColTran colorizes 500 grayscale images, with 3 different colorizations per image, denoted as seeds.
Human raters assess the quality of these colorizations with a two alternative-forced choice (2AFC) test.
We display both the ground-truth and recolorized image sequentially for one second in random order.
The raters are then asked to identify the image with fake colors.
For each seed, we report the mean fooling rate over 500 colorizations and 5 different raters.
For the oracle methods, we use the human rating to pick the best-of-three colorizations.
ColTran’s best seed achieves a fooling rate of 42.3 % compared to the 35.4 % of PixColor’s best seed.
ColTran
Oracle achieves a fooling rate of 62 %, indicating that human raters prefer ColTran’s best-of-three colorizations over the ground truth image itself.
Visualizing uncertainty.
The autoregressive core model of ColTran should be highly uncertain at object boundaries when colors change.
Figure 5 illustrates the per-pixel, maximum predicted probability over 512 colors as a proxy for uncertainty.
We observe that the model is indeed highly uncertain at edges and within more complicated textures.
C ONCLUSION
We presented the Colorization Transformer (ColTran), an architecture that entirely relies on selfattention for image colorization.
We introduce conditional transformer layers, a novel building block for conditional, generative models based on self-attention.
Our ablations show the superiority of employing this mechanism over a number of different baselines.
Finally, we demonstrate that ColTran can generate diverse, high-ﬁdelity colorizations on ImageNet, which are largely indistinguishable from the ground-truth even for human raters. ACKNOWLEDGEMENTS
We would like to thank Mohammad Norouzi, Rianne van den Berg, Mostafa Dehghani or their useful comments on the draft and Avital Oliver for assistance in the Mechanical Turk setup.
A C ODE , CHECKPOINTS AND TENSORBOARD FILES
Our implementation is open-sourced in the google-research framework at https://github.com/googleresearch/google-research/tree/master/coltran. Our full set of hyperparameters are available here. We provide pre-trained checkpoints of the colorizer and upsamplers on ImageNet at https://console.cloud.google.com/storage/browser/gresearch/coltran.
Our colorizer and spatial upsampler for thesr checkpoints were trained longer for 600K and 300K steps which gave us a slightly improved FID score of ≈ 19 .
Finally, reference tensorboard ﬁles for our training runs are available at colorizer tensorboard, color upsampler tensorboard and spatial upsampler tensorboard.
B E XPONENTIAL M OVING A VERAGE
We found using an exponential moving average (EMA) of our checkpoints, extremely crucial to generate high quality samples.
In Figure 6, we display the FID as a function of training steps, with and without EMA.
On applying EMA, our FID score improves steadily over time.
C N UMBER OF PARAMETERS AND INFERENCE SPEED
Inference speed.
ColTran core can sample a batch of 20 64x64 grayscale images in around 3.5minutes on a P100 GPU vs PixColor that takes 10 minutes to colorize 28x28 grayscale images on a
K40 GPU.
Sampling 28x28 colorizations takes around 30 seconds.
The upsampler networks take in the order of milliseconds.
Further, in our naive implementation, we recompute the activations, c U zs , c U zb , c U fs , c U fb in Table 1 to generate every pixel in the inner decoder.
Instead, we can compute these activations once per-grayscale image in the encoder and once per-row in the outer decoder and reuse them.
This is likely to speed up sampling even more and we leave this engineering optimization for future work.
Number of parameters.
ColTran has a total of ColTran core (46M) + Color Upsampler (14M) +
Spatial Upsampler (14M) = 74M parameters.
In comparison, PixColor has Conditioning network (44M) + Colorizer network (11M) + Reﬁnement Network (28M) = 83M parameters.
D L OWER COMPUTE REGIME
We retrained the autoregressive colorizer and color upsampler on 4 TPUv2 chips (the lowest conﬁguration) with a reduced-batch size of 56 and 192 each.
For the spatial upsampler, we found that a batch-size of 8 was sub-optimal and lead to a large deterioration in loss.
We thus used a smaller spatial upsampler with 2 axial attention blocks with a batch-size of 16 and trained it also on 4 TPUv chips.
The FID drops from 19.71 to 20.9 which is still signiﬁcantly better than the other models in 2.
We note that in this experiment, we use only 12 TPUv2 chips in total while PixColor (Guadarrama et al., 2017) uses a total of 16 GPUs.
E I MPROVED FID WITH T OP -K SAMPLING
We can improve colorization ﬁdelity and remove artifacts due to unnatural colors via Top-K sampling at the cost of reduced colorization diversity.
In this setting, for a given pixel ColTran generates a color from the top-K colors (instead of 512 colors) as determined by the predicted probabilities.
Our results in Figure 6 K = 4 and K = 8 demonstrate a performance improvement over the baseline
ColTran model with K =
F A DDITIONAL ABLATIONS :
Additional ablations of our conditional transformer layers are in Figure 7 which did not help.
• Conditional transformer layers based on Gated layers (Oord et al., 2016) ( Gated ) •
A global conditioning layer instead of pointwise conditioning in cAtt and cLN.
cAtt + cMLP, global
G A UTOREGRESSIVE MODELS
Autoregressive models are a family of probabilistic methods that model joint distribution of data
P ( x ) or a sequence of symbols ( x 1 , x 2 , . . .
x n ) as a product of conditionals (cid:81) Ni =1 P ( x i | x <i ) .
During training, the input to autoregressive models are the entire sequence of ground-truth symbols.
Masking ensures that the contribution of all "future" symbols in the sequence are zeroed out.
The outputs of the autoregressive model are the corresponding conditional distributions.
P ( x i | x <i ) .
Optimizing the parameters of the autoregressive model proceeds by a standard log-likelihood objective.
Generation happens sequentially, symbol-by-symbol.
Once a symbol x i is generated, the entire sequence ( x 1 , x 2 , . . .
x i ) are fed to the autoregressive model to generate x i +1 .
In the case of autoregressive image generation symbols typically correspond to the 3 RGB pixelchannel.
These are generated sequentially in raster-scan order, channel by channel and pixel by pixel.
H R OW /C OLUMN S ELF - ATTENTION
In the following we describe row self-attention, that is, we omit the height dimension as all operations are performed in parallel for each column.
Given the representation of a single row within of an image x i, · ∈ R W × D , row-wise self-attention block is applied as follows:
[Equation]
LN refers to the application of layer normalization (Ba et al., 2016).
Finally, we apply residual connections and a feed-forward neural network with a single hidden layer and ReLU activation ( MLP ) after each self-attention block as it is common practice in transformers.
[Equation]
Column-wise self-attention over x · ,j ∈ R H
× D works analogously.
I O UT OF DOMAIN COLORIZATIONS
We use our trained colorization model on ImageNet to colorize high-resolution grayscale images from LSUN 256 × 256 (Yu et al., 2015) and low-resolution grayscale images from Celeb-A
(Liu et al., 2015) 64 × 64 .
Note that these models were trained only on ImageNet and not ﬁnetuned on
Celeb-A or LSUN.
J N UMBER OF AXIAL ATTENTION BLOCKS
We did a very small hyperparameter sweep using the baseline axial transformer (no conditional layers) with the following conﬁgurations:
• hidden size = 512, number of blocks =• hidden size = 1024, number of blocks =• hidden size = 512, number of blocks =
Once we found the optimal conﬁguration, we ﬁxed this for all future architecture design.
K A NALYSIS OF MT URK RATINGS
We analyzed our samples on the basis of the MTurk ratings in Figure 11.
To the left, we show images, where all the samples have a fool rate > 60 %.
Our model is able to show diversity in color for both high-level structure and low-level details.
In the center, we display samples that have a high variance in MTurk ratings, with a difference of 80 % between the best and the worst sample.
All of these are complex objects, that our model is able to colorize reasonably well given multiple attempts.
To the right of Figure 11, we show failure cases where all samples have a fool rate of 0 %,
For these cases, our model is unable to colorize highly complex structure, that would arguably be difﬁcult even for a human.
L M ORE PROBABILITY MAPS
We display additional probability maps to visualize uncertainty as done in 5.5.
M M ORE SAMPLES
We display a wide-diversity of colorizations from ColTran that were not cherry-picked.