Colorization Transformer
Authors: Manoj Kumar
Dirk Weissenborn
Nal Kalchbrenner
Venue: ICLR 2021 Poster (2021)

Abstract: 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

Published as a conference paper at ICLR 2021
COLORIZATION TRANSFORMER
Manoj Kumar, Dirk Weissenborn & Nal Kalchbrenner Google Research, Brain Team
(mechcoder, diwe, nalk) @google.com
ABSTRACT
We present the Colorization Transformer, a novel approach for diverse high fidelity image colorization based self-attention. Given grayscale image, the coloriza- tion 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-ard on colorising ImageNet based on FID results and based on human evaluation in a Mechanical Turk test. Remarkably, in more than 60% 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 ur!
INTRODUCTION
Figure 1: Samples of our model showing diverse, figeliicyiity colorizations.
Image colorization is challenging, inherently stochastic task that requires 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 (Tafftarss et al. 2014). videos (Geshwind
1986) and black-and-whin white cartoons (Sýkora al. 2004 Qu al. 2006 Cinarel Zhang 2017). Colorization also has important technical uses as a way to learn meaninglul representations without explicit supervision (Zhang al. 2016:] Larsson et al., 2016) Vondrick et al., 2018) or as an unsuper- vised data augmentation technique, whereby diverse comantics-preserviie images are produced with a colorization model trained on a potentially much larger set of unlabelled images.
colorizations of labelled
The current state-of-the--rt automated colorization are neural generative approaches based on ligllikel estimation (Guadarrama al., 2017;; Royer et al. 2017 Ardizzone et al. 2019). Probabilistic models are natural fit for the one-to me-to--may task 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 al. 2019h Wang al. 2020). The main

---

Published as a conference paper at ICLR 2021
advantages of axial self-attention blocks are the ability to capture global receptive field with only two layers and DD instead of O(D²) complexity. They can be implemented efficiently using 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-ar autoregressive image generation model that does not require custom kernels (Child al., 2019). While Axial Transformers support conditioning by biasing the input, we find that directly conditioning the transformer layers can improve results significantly. 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. (017) and as an effect this results in improved colorization fidelity. Finally, we employ fast parallel deterministic upsampling models to super-resolve the coarsely colorized image into the final high resolution output. In summary, our main contributions are:
matrix-multiplicatee
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 Improved FID scores demonstrate the usefulness of this auxiliary model.
We establish new state-of-the-art image colorization outperforming prior methods by large margin on FID scores and a 2-Altemattvv 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.
2 RELATED WORK
Colorization methods have initially relied on human-in-the-loop approaches to provide hints in the
form of scribbles (Levin et al. 2004 Troni et al. [2055: Huang et al.] [2055] Yatziv & Sapiro 2006
2006 Qu al..
Luan et al., [007]] Tsaftaris et al. [014]] 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 Troni et al.) 2005 Pitié et al. 2007
2009 Gupta et al. 2012 Xiao et al. 2020). Exemplar based techniques have been recently extended to video as well (Zhang al. (019a). 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 (Tizuka et al., 2016), Pix2Pix (Isola et al. 2017) and DC (Cheng al.) 2015 Dahl 2016) involve variations of CNNs to model per-pixel color information conditioned on the intensity.
Morimoto et al.
Generative colorization models typically extend unconditional image generation models to incorporate conditioning information from a grayscale image. Specifically, cINN (Ardizzone et al., 2019) use
conditional normalizing flows (Dinh et al. 2014), VAE-MDN (Deshpande et al. 2017 2015 and
SCC-DC (Messaoud et al. (018)) use conditional VAEs (Kingma & Welling (2013). 3), and cGAN Cao
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.] (017) with PixColor obtaining slightly better results than PIC due to 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 completely unsupervised fashion, while the conditioning network in PixColor requires pre- preitrannng with an object detection network that provides substantial semantic information. PixColor relies on PixeICNN (Oord et al. 2016) that requires large depth to model interactions between all pixels. ColTran relies on Axial Transformer (Ho et al. (019b) 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, train
2

---

Published as a conference paper at ICLR 2021
autoregressive model on single coarse channel and separate color upsampling network that improves fidelity (See: (.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 al.
2019a Weissenborn et al., 2019) and GANs (Zhang et al. 2019b)
3
BACKGROUND: AXIAL TRANSFORMER
3.1 Row AND COLUMN SELF-ATTENTION
Self offfatteenion (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 modeling innovations that don't 't require running global self-attention
between all pixels. Following
the work of (Ho et al. (019b) 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.2 AXIAL TRANSFORMER
Ths Axial Transformer (Ho et al. 2019b) is autoregressive model that applies (masked) row- and column self-attention operations in a way that efficiently summarizes all past information xi,<j and x to model distribution over pixel x,, position i. Causal masking is employed by setting all Am, = 0 where n > m during self attentieeoo (see Eq. 15).
Outer decoder. The outer decoder computes state S, over all previous rows x by applying N layers of full row self-attention followed by masked column self-attention. (Eq[2). s, is shifted down by single row, such that the output context o, position i,j only contains information about pixels Xei. from prior rows. (Eq
Embeddings((x)
(1)
S, MaskedCColum((oo(()))
X N
(2)
ShiftDDon(s.))
Inner decoder. The embeddings to the inner decoder are shifted right by a single column to mask the current pixel x.j. The context o from the outer decoder conditions the inner decoder by biasing the shifted embeddings. It then computes a final state h. by applying N layers of masked row-wise self-attention to infuse additional information from prior pixels of the same row x, (Eg 4h. h comprises information about all past pixels X and x; A dense layer projects into distribution p(x:z) over the pixel at position (i, J) conditioned on all previous pixels xi,xj and x<i...
z = o + ShiftRight(e)
(4)
h MaskedRow(z)
N
(5)
p(x,)) = Dense(h)
(6)
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 conditioning context с for the "current channel". The context conditions the outer and inner decoder by biasing the inputs in I and Eq respectively.
3

---

Published as a conference paper at ICLR 2021
Pelar (Lr"///lad
Sample
Softmax
Par(irla a (
Arggax
Linez
Inner Decoder
Softmax
Lancar
Pe(24lel (
Couditional Manked
Bow Alfcatics
Argmax
Per (
Colina Aftenton
Sotteax
Shift right
Latear
Row Amonioo
N
Shift
Colara Attentons
Conditional Mookal
Colama Attention
Eniieskingg
Embhelánzz
Amtntion
Conditoocal Roo
N
Grayscals eacnder
Outer Decoder
25
Enbolinns
Firbeslings
Erihaldanss
Friheadinss
intepodation
ColTran Core
ColTran Upsamplers
Figure Depiction ColTran. consists 3 individual models: an autoregressive colorizer (left), a color upsampler (middle) and a spatial upsampler (right). Each model optimized independently. The autoregressive colorizer (ColTran core) is instantiation of Axial Transformer (Sec. 3.2 Ho et al. (2019b)) with conditional transfomer layers and auxiliary parallel head proposed in this work (Sec. 4.1 During training. the ground- truth coarse low resolution image is both the input the decoder and the target. Masked layers ensure that the conditional distributions for each pixel depends solely on previous ground-truth pixels. (See Appendix G for recap on autoregressive models). ColTran upsamplers are stacked row/column attention layers that deterministically upsample color and space in parallel. Each attention block (in green) is residual and consists of the following operations: layer-norm multihead self-attention MLP.
Sampling. The Axial Transformer natively supports semi-parallel sampling that avoids re- evaluation of the entire network generate each pixel of RGB image. The encoder run once per-channel, the outer decoder is 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 and the encoder outer decoder condition the inner decoder (Eq 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. (I) to 33) 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.
4 PROPOSED ARCHITECTURE
Image colorization is the task transforming a grayscale image 29 E Hx into a colored image x E R The task is inherently stochastic; for a given grayscale image x" there exists a conditional distribution over x. p(=[="). Instead predicting & directly from 29, we instead sequentially predict two intermediate low resolution images x and si with different color depth first. Besides simplifying the task of high-resolution image colorization into simpler tasks, the smaller resolution allows for training larger models.
We obtain == spatially downsampled representation of v. by standard area interpolation. s1cs is a bit per-channel representation of = that is, each color channel has only intensities. Thus, there p are 8 = 512 coarse colors per pixel which are predicted directly as a single "color" channel. We rewrite the conditional likelihood p(=[#9) to incorporate the intermediate representations as follows:
p(a ( p(x|=9)-
p(x|=9) pea , [z, x²) = p(æ"-0t " x 2
= p( .r)) ) (o Jarte
.2)) [
(8)
ColTran core (Section 4.1 a parallel color upsampler and a parallel spatial upsampler (Section 4.3) model p(a: (xo)) p eics nr) and p(at respectively. In the subsections below, descrihe
4

---

Published as a conference paper at ICLR 2021
Component
Unconditional
Conditional
y = Softmax
9.k. VD
Vv
Self-Attention
y Softmax
VD
where 1/2 k, q.
(cU?) 6 z + (cU))
h ReLU(xL's + b1) b2
MLP
ReLUU(XU + bi) bb
= (cti)) o h (cUg)
8. Norm(x) Yo
Layer Norm
y = 3 Norm(x) +
where Vp 3cr % E xxxW è RRWxD
== (u -)U ue NEW
Table 1:
We contrast the different components of unconditional self-attention with self-attention conditioned on context E MxXx Learnable parameters specific conditioning denoted by and U. PxD
these individual components in detail. From now we will refer to all low resolutions as M x N and high resolution H W. An illustration of the overall architecture is shown Figure
4.1 COLTRAN CORE
In this section, we describe ColTTaa core, a conditional variant of the Axial Transformer (Ho et al. 2019b) for low resolution coarse colorization. ColTran Core models a distribution po (x9) over 512 coarse colors for every pixel, conditioned
low resolution grayscale image in addition to the colors from previously predicted pixels as per raster order (Eq.
M N
pe(x
= IIII
(a 2.
1x9
SIGI
(9)
=1
Given context representation E R MXNDD propose conditional transformer layers in Table Conditional transformer layers have conditional versions of all components within the standard attention block (see Appendix H Eqs. 18)
Conditional Self-Attention. For every layer in the decoder, we apply six 1 x convolutions to c to obtain three scale and shift vectors which apply element-wise q. k and of the self-attention operation (Appendix 3.1. respectively.
Conditional MIP. 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 as for self-attention.
Conditional Layer Norm. Layer normalization (Ba et al. 2016) globally scales and shifts given normalized input using learnable vectors 3. 7. Instead, we predict 3. and Ye as a function of c. We first aggregate с into global I-D representation с E RR via a learnable, spatial pooling layer. Spatial pooling is initialized as mean pooling layer. Similar 1-D conditional normalization layers (Perez et al., 2017 De Vries et al.] 2017 Dumoulin et al.) 2016; Huang & Belongie 2017). we then apply linear projection on to predict Se and %e. respectively.
A grayscale encoder consisting of multiple, alternating row and column self-attention layers encodes the grayscale image into the initial conditioning context cº. serves as both context for the conditional layers and additional input the embeddings of the outer decoder. The of the outer decoder's s output and e condition the inner decoder. Figure 2 ililustrates how conditioning is applied in the autoregressive of the ColTran architecture.
Conditioning every layer via multiple components allows stronger gradient signals through the encoder and as 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)
5

---

Published as a conference paper at ICLR 2021
4.2 AUXILIARY PARALLEL MODEL
We additionally found it beneficial to train an auxiliary parallel prediction model that models P. (x))))) directly on top of representations learned by the grayscale encoder which we found beneficial for regularization (Eq. T0p
M IIII Pe Pal ((()))))))))) ( ==
Pe
(10)
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 beneficial, regularizing effect on learning. We apply a linear projection, Uparallel E R Lx512 on top of c (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.3
4.3 COLOR & SPATIAL UPSAMPLING
In order to produce high-fidelity 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 srich E MxNxx 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 sic E RgMxx RMMXNXD where k E {R, G, B} indicates the channel. We upsample each channel individually conditioning only on the respective channel's embedding. The channel em- bedding is summed with the respective grayscale embedding for each pixel and serve as input to the subsequent self-attention layers (encoder). The output the encoder further projected to per [
pixel-channel probability distributions Por (#x [x 2²) E IR M x256 over 256 color intensities for all & E {R, G, B} (Eq.
M II
Per (x
Per (1 (x9 x xLLL)
(11)
Spatial Upsampler. We first naively upsample x² E MxNx3 into blurry, high-resolution RGB image using area interpolation. As above, we then embed each channel of the blurry RGB image and run a per-channell encoder exactly the same way as with the color upsampler. The output of the encoder finally projected to per pixel-channel probability distributions P (xx |x* x) E R x x256 over 256 color intensities for all & (R. G, B). (Eq. 12
pa- (2:g J20, x
(12)
In our experiments, similar to GGuadarrama et al., 2017). we found parallel upsampling to sufficient 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. [2]and 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 great variety of colorizations.
Objective. We train our architecture to minimize the negative log-likelihood (Eq. 13) of the data. Pe/Des PSt. Per are maximized independently and is a hyperparameter that controls the relative contribution of Pe and Pe
L= (1 A) log + A Alog Pe + log Per log Per
(13)

---

Published as a conference paper at ICLR 2021
Velidation
Validation
Validation loss
0.70
0.65
0.65
Coffran
Coffran
caltran e 2x
CLN
(SC, enys
caltran
0.65
Caffran
0.63
0.63
Store
Loli ran 2X
0.54
0.62
0.52
0.80
25000
75000
125000
25000
75000
125000
0.50
25000
75000
125000
Training steps
Training steps
Training steps
Figure 3: Per pixel log- ligeliikld of coarse colored 64 x 64 images over the validation set as a function of training steps. We ablate the various components of the ColTran core in each plot. Left: ColTran with Conditional Transformer Layers vs baseline Axial Transformer which conditions via addition (ColTran-B) B). ColTran-B 2x and ColTran-B 4x refer to wider baselines with increased model capacity. Center: Removing each conditional sub comp-oommnnnt one at time (no cLN, no cMLP and no cArr). Right: Conditional shifts only (Shift), Conditional scales only (Scale), removal of kq conditioning in cAt (cAtt, only v) and fixed mean pooling in cLN (cLN, mean pool). See Section 5.2 for more details.
EXPERIMENTS
5.1 TRAINING AND EVALUATION
We evaluate ColTTaa on colorizing 256x256 grayscale images from the ImageNet dataset (Rus. sakovssy et al., 2015). We train the ColTran core, color and spatial upsamplers independently on 16 TPUv2 chips with batch-size of 224, 768 and 32 for 600K, 450K and 300K steps respectively. We use 4 axial attention blocks in each component of our architecture, with a hidden size of 512 and 4 heads. We use RMSprop (Tieleman & Hinton 2012) with fixed learning rate of 3e 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 the grayscale images from this holdout set. We use the public validation set to display qualitative results and report final numbers.
5.2 ABLATIONS OF COLTRAN CORE
The autoregressive core of ColTran models downsampled, coarse-colored images of resolution 64 x 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- ligeliiel over 512 coarse colors the validation set as a function of training steps.
Impact of conditional transformer layers. The left side of Figure illustrates the significant 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 B 4x). Both ColTran-B 2x and ColTran-B 4x have an increased parameter count which makes for 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 cMIP and no cAu in the middle of Figure 3 quantifies our results. While each conditional component improves final 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 the right hand side of architectures that use conditional shifts and conditional scales only. While both types of interactions important, multiplicative interactions have much stronger impact.
the impact of these interactions via ablated
7

---

Published as a conference paper at ICLR 2021
F10
18.5
FID vs. Val dation loss
18.0
0.00
18.0
17.0
0.01
17.5
16.8
17.0
12.0
16.6
16.5
18.5
16.0
15.5
15.5
16.2
15.0
4000
300300
J9,000
120000
15.0
0.0
0.1
0.2
0.3
0.4
0.5
0.009
0.605 6,11 6.015 0.620 6,625
Training stops
A
Validation
Figure 4: Left: FID of generated 64 x 64 coarse samples as a function of training steps for A = 0.01 and À = 0.0. Center: Final FID scores as a function of X. Right: FID as a function of log-likclihood.
Context-aware dot product attention. Self-attention computes the similarity between pixel repre- sentations using a dot product between and k (See: Eq 15). cAtt applies conditional shifts and scales q. k and allow modifying this similarity based contextual information. The curve cArr, 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 flexible learn- able spatial pooling layer. We experimented with fixed 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.3 OTHER ABLATIONS
Auxiliary Parallel Model. We study the effect of the hyperparameter A, which controls the con- tribution of the auxiliary parallel prediction model described in Section 4.2 For given X, we now optimize (A) (1 A) log pe(.) Alg (.)( instead of just log Pe ). Note that Pc).,, models each pixel independently, which more difficult than modelling each pixel conditioned on previous pixels given by Pc(.). Hence, employing (A) holdout metric, would just lead to trivial soluion at A = O. Instead, the FID of the generated coarse 64x64 samples provides a reliable way to find an optimal value of X. In Figure at x 0.01, our model converges better FID faster with marginal but consistent final 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 x 64 coarse images, the ColTran upsamplers map these to fine grained 256 x 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. is further extremely important to provide the grayscale image 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.4 FRECHET INCEPTION DISTANCE
We compute FID using colorizations of 5000 grayscale images of resolution 256 x 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-trutt distribution. addition to ColTran, we report two additional results ColTran-S and ColTran-B. BolTran-B refers to the baseline Axial Transformer that conditions via addition at the input. PixColor samples smaller 28 x 28 colored images autoregressively as compared to ColTran's 64 x 64, As a control experiment, architectural choices and the inherent stochasticity of modelling higher resolution images. ColTran-S and ColTran-B obtains FID scores of 22.06 and 19.98 that significantly improve over the previous best FID of 24.32. Finally, ColTran achieves the best FID score of 19.37. All results are presented in Tef..
train an autoregressive model on resolution 28 x 28 (ColTran-S) to disentangle
8

---

Published as a conference paper at ICLR 2021
Models
AMT Fooling rate
Models
FID
ColTran (Oracle)
62.0 % = 0.99
ColTran
19.37 +1 0.09
ColTran (Seed 1)
40.5 % + 0.81
ColTran-BB
19.98 0.20
ColTran (Seed 2)
42.3 % + 0.76
ColTran-S
22.06 + 0.13
ColTran (Seed 3)
41.7 % = 0.83
PixColor 16
24.32 0.21
PixColor [6] (Oracle)
38.3 % + 0.98
cGAN 3
24.41 + 0.27
PixColor (Seed I)
33.3 % = 1.04
cINN
25.13 1: 0.3
PixColor (Seed 2)
35.4 % = 1.01
VAE-MDN
25.98 1 0.28
PixColor (Sood
33.2 % = 1.03
Ground truth
14.68 + 0.15
CIC [56]
29.2 % = 0.98
Grayscale
30.19 0.1
IRACC 27
30.9 % = 1.02
L.BBC 22
25.8 % + 0.97
Table 2:
We outperform various state-of-the-eee colorization models both FID (left) and human evaluation
(right). We obtain the scores from (Ardizzone et al. 2019) and human evaluation results from Guadarrama et al. (2017). ColTran-B is a baseline Axial Transformer that conditions via addition and ColTran-S is a control experiment where we train Col Crl'Tran core (See: 4.1 on smaller 28 x 28 colored images.
Figure 5: We display the per-pixel, maximum predicted probability over 512 colors proxy for uncertainty.
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 moderately positive correlation of 0.57 between the log- ligliikll and FID. Importantly, even an absolute improvement the order of 0.01 0.02 can improve FID significantly. This suggests that designing architectures that achieve better log-likelihodd values is likely to lead to improved FID scores and colorization fidelity.
5.5 QUALITATIVE EVALUATION
Human Evaluation. For our qualitative assessment, follow the protocol used in PixColor (Guadarrama et al. 2017). ColTran colorizes 500 grayscale images, with 3 different colorizations per image, denoted seeds. Human raters assess the quality of these colorizations with two cirmnative-foreed choice (2AFC) test. We display both the ground-truie 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 fooling rate of 42.3 % compared to the 35.4 % of PixColor's ss best seed. Col'ran Oracle achieves fooling rate 62 %, indicating that human raters prefer Col'ran'ss 's best-of-thre colorizations over the ground truth image itself.
Visualizing uncertainty. The autoregressive core model of ColTran should be highly uncertain object boundaries when colors change. Figure illustrates the per-pixel, maximum predicted probability over 512 colors as a proxy for uncertainty. We observe that the model is indeed highly uncertain edges and within more complicated textures.
CONCLUSION
We presented the Colorization Transformer (ColTran), an architecture that entirely relies on self- attention for image colorization. We introduce conditional transformer layers, novel building block for conditional, generative models based on self-attention. Our ablations show the superiority of number of different baselines. Finally, demonstrate that ColTran
employing this mechanism can generate diverse, sighfiidelity colorizations on ImageNet, which are largely indistinguishable from the ground-truth for human raters.

---

Published as a conference paper at ICLR 2021
REFERENCES
Lynton Ardizzone, Carsten Lüth, Jakob Kruse, Carsten Rother, and Ullrich Köthe. Guided image generation with conditional invertible neural networks. arXiv preprint arXiv: 1907002390 2019.
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprini arXiv:1607.064500 2016.
Yun Cao, Zhiming Zhou, Weinan Zhang, and Yong Yu. Unsupervised diverse colorization via generative adversarial networks, 2017.
Zezhou Cheng, Qingxiong Yang, and Bin Sheng. Deep colorization. Proceedings of the IEEE International Conference Computer Vision, pp. 415 111--23,222 2015.
Rewon Child, Scott Gray, Alec Radford, and llya Sutskever. Generating long sequences with sparse transformers. arXiv preprint
2019.
Yuanzheng Ci, Xinzhu Ma, Zhihui Wang, Haojie Li, and Zhongxuan Luo. User-guided deep anime line art colorization with conditional adversarial networks. In Proceedings of the 26th ACM international conference on Multimedia, pp. 1536-1544, 2018.
Ceyda Cinarel and Byoung-Take Tak Zhang. Into the colorful world of webtoons: Through the lens of neural networks. In 2017 14th IAPR International Conference on Document Analysis and Recognition (ICDAR), volume pp. 35-40. IEEE, 2017.
Ryan Dahl. Automatic colorization, 2016
Harm De Vries, Florian Strub, Jérémie Mary, Hugo Larochelle, Olivier Pietquin, and Aaron C Courville. Modulating early visual processing by language. In Advances in Neural Information Processing Systems, pp. 6594-6604, 2017.
Aditya Deshpande, Jason Rock, and David Forsyth. Learning large-scale automatic image colorization. In Proceedings of the IEEE International Conference Computer Vision, pp. 567-575, 2015.
Aditya Deshpande, Jiajun Lu, Mao-Chuang Yeh, Min Jin Chong, and David Forsyth. Learning diverse image colorization. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pp. 6837-6845, 2017.
Laurent Dinh, David Knueger, and Yoshua Bengio. Nice: Non-linear independent components estimation. arXiv preprint
16, 2014.
Vincent Dumoulin, Jonathon Shlens, and Manjunath Kudlur. A learned representation for artistic style. arXiv preprint orXiv:1610.0762. 2016.
David M Geshwind. Method for colorizing black and white footage, August 19 1986. US Patent
66666665
Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial networks. arXiv preprint arXiv:1406.2661, 2014.
Sergio Guadarrama, Ryan Dahl, David Bieber, Mohammad Norouzi, Jonathon Shlens, and Kevin Murphy. Pixcolor: Pixel recursive colorization. arXiv preprint arXiv: 1705::7705.07208 2017.
Raj Kumar Gupta, Alex Yong-Sang Chia, Deepu Rajan, Ee Sin Ng. and Huang Zhiyong. Image colorization using similar images. In Proceedings of the 20th ACM international conference on Multimedia, pp. 369-378, 2012.
Jonathan Ho, Xi Chen, Aravind Srinivas, Yan Duan, and Pieter Abbeel. Flow++: Improving flow- based generative models with variational dequantization and architecture design. arXiv preprint arXiv:1902.00225, 00205 2019a.
Jonathan Ho, Nal Kalchbrenner, Dirk Weissenborn. and Tim Salimans. Axial attention in multidi- mensional transformers. arXiv preprint 1912: 2019b.
10

---

Published as a conference paper at ICLR 2021
Xun Huang and Serge Belongie. Arbitrary style transfer in real-time with adaptive instance normal- ization. In Proceedings of the IEEE International Conference on Computer Vision, pp. 1501-1510, 2017.
Yi-Chin Huang, Yi-Shin Tung, Jun-Chene Chen, Sung- Wen Wang, and Ja-Ling Wu. An adaptive edge detection based colorization algorithm and its applications. In Proceedings of the 13th anmual ACM international conference on Mullimedia, pp. 351-354, 2005.
Satoshi lizuka, Edgar Simo-Serra and Hiroshi Ishikawa. Let there be color! joint end-to-end learning of global and local image priors for automatic image colorization with simultaneous classification. ACM Transactions on Graphics (ToG), 35(4):1-11, 2016.
Revital Ironi, Daniel Cohen-Or, and Dani Lischinski. Colorization by example. In Rendering Techniques, pp. 201-210. Citeseer, 2005.
Phillip Isola, Jun-Yan Zhu, Tinghui Zhou, and Alexei A Efros. Image-to-image translation with conditional adversarial networks. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 1 1125-1134, 2017.
Norman P. Jouppi, Cliff Young, Nishant Patil, David Patterson, Gaurav Agrawal, Raminder Bajwa, Sarah Bates, Suresh Bhatia, Nan Boden, Al Borchers, Rick Boyle, Pierre luc Cantin, Clifford Chao, Chris Clark, Jeremy Coriell, Mike Daley, Matt Dau, Jeffrey Dean, Ben Gelb, Tara Vazir Ghaemmaghami, Rajendra Gottipati, William Gulland, Robert Hagmann, C. Richard Ho. Doug Hogberg, John Hu, Robert Hundt, Dan Hurt, Julian Ibarz, Aaron Jaffey, Alek Jaworski, Alexander Kaplan, Harshit Khaitan, Andy Koch, Naveen Kumar, Steve Lacy, James Laudon, James Law, Diemthu Le. Chris Leary, Zhuyuan Liu, Kyle Lucke, Alan Lundin, Gordon MacKean, Adriana Maggiore, Maire Mahony, Kieran Miller, Rahul Nagarajan, Ravi Narayanaswami, Ray Ni, Kathy Nix, Thomas Norrie, Mark Omernick, Narayana Penukonda, Andy Phelps, Jonathan Ross, Matt Ross, Amir Salek, Emad Samadiani, Chris Severn, Gregory Sizikov, Matthew Snelham, Jed Souter, Dan Steinberg, Andy Swing, Mercedes Tan, Gregory Thorson, Bo Tian, Horia Toma, Erick Tuttle, Vijay Vasudevan, Richard Walter, Walter Wang, Eric Wilcox, and Doe Hyun Yoon. In-datacenter performance analysis of tensor processing unit, 2017.
Diederik P Kingma and Max Welling. arXiv:1312.61 14, 2013.
Auto-encoding variational bayes.
arXiv preprint
Gustav Larsson, Michael Maire, and Gregory Shakhnarovich. Learning representations for automatic colorization. In European conference on computer vision, pp. 577-593. Springer, 2016.
Anat Levin, Dani Lischinski, and Yair Weiss. Colorization using optimization. In ACM SIGGRAPH 2004 Papers, pp. 689-694. 2004.
Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang. Deep learning face attributes in the wild. In Proceedings of the IEEE international conference computer vision, pp. 3730-3738, 2015.
Qing Luan, Fang Wen, Daniel Cohen-Or, Lin Liang, Ying- Ying-Qing Xu, and Heung- Yeung Shum. Natural image colorization. In Proceedings of the 18th Eurographics conference on Rendering Techniques, pp. 309-320, 2007.
Jacob Menick and Nal Kalchbrenner. Generating high fidelity images with subscale pixel networks and multidimensional upscaling. arXiv preprint arXiv:18/2.0/66 18/2.0160,, 2018.
Safa Messaoud, David Forsyth, and Alexander G. Schwing. Structural consistency and controllability for diverse colorization. In Proceedings of the European Conference on Computer Vision (ECCV), September 2018
Yuji Morimoto, Yuichi Taguchi, and Takeshi Naemura. Automatic colorization of grayscale images using multiple images on the web. In SIGGRAPH 2009: Talks, pp. 1-1, 2009.
Aaron van den Oord, Nal Kalchbrenner, and Koray Kavukcuoglu. Pixel recurrent neural networks. arXiv preprint 1601 006559, 2016.
Niki Parmar, Ashish Vaswani, Jakob Uszkoreit, Lukasz Kaiser, Noam Shazeer, Alexander Ku, and Dustin Tran. Image transformer. arXiv preprint arXiv: 1=X2.0/802.0075/ 2018.
11

---

Published as a conference paper at ICLR 2021
Ethan Perez, Florian Strub, Harm De Vries, Vincent Dumoulin, and Aaron Courville. Film: Visual reasoning with a general conditioning layer. arXiv preprint arXiv: 2017.
François Pitié, Anil C Kokaram, and Rozenn Dahyot. Automated colour grading using colour distribution transfer. Computer Vision and Image Understanding, (07(1-2):123-133) 2007
Yingge Qu, Tien-Tsin Wong, and Pheng-Ann Heng. Manga colorization. ACM Transactions on Graphics (TOG), 25())1214-11220, 2006.
Erik Reinhard, Michael Adhikhmin, Bruce Gooch, and Peter Shirley. Color transfer between images.
IEEE Computer graphics and applications, 21(5):3441, 2001.
Amelie Royer, Alexander Kolesnikov, and Christoph Lampert. Probabilistic image colorization. arXiv preprini arXiv: 770504428, 2017.
Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael Bernstein, et al. Imagenet large scale visual recognition challenge. International journal of computer vision, 115(3):21 2015.
Daniel Sýkora, Jan Buriánek, and Jifí Zára. Unsupervised colorization black-and and-white cartoons. In Proceedings of the 3rd international symposium Non-photorealistic animation and rendering, pp. 121-127, 2004.
Yu-Wing Tai, Jiaya Jia, and Chi-Keung Tang. Local color transfer via probabilistic segmentation by 2005 IEEE Computer Society Conference on Computer Vision and Pattern Recognition (CVPR 05), volume 1, pp. 747-754. IEEE, 2005.
popectationmmaximizztioo
Tijmen Tieleman and Geoffrey Hinton. Lecture 6.5-rmspoop: Divide the gradient by running average of its recent magnitude. COURSERA: Neural networks for machine learning. 4(2):26-31, 2012.
Sotirios A Tsaftaris, Francesca Casadio, Jean-Louis Andral, and Aggelos K Katsaggelos. A novel visualization tool for art history and conservation: Automated colorization of black and white archival photographs of works of art. Studies in conservation, 59(3):125-1133 2014.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in neural information processing systems, pp. 5998-6008, 2017.
Carl Vondrick, Abhinav Shrivastava, Alireza Fathi, Sergio Guadarrama, and Kevin Murphy, Tracking emerges by colorizing videos. In Proceedings of the European conference on computer vision (ECCV), pp. 391-408, 2018.
Huiyu Wang, Yukun Zhu, Bradley Green, Hartwig Adam, Alan Yuille, and Liang-Chieh Chen. Axial- deeplab: Stand-alone axial-attention for panoptic segmentation. arXiv preprint arXiv:2003.07853, 2020.
Dirk Weissenborn, Oscar Täckström, and Jakob Uszkoreit. Scaling autoregressive video models. arXiv preprint arXiv:1906.000 02634, 2019.
Tomihisa Welsh, Michael Ashikhmin, and Klaus Mueller. Transferring color to greyscale images. In Proceedings of the 29th annual conference on Computer graphics and interactive techniques, pp. 277-280, 2002.
Chufeng Xiao, Chu Han, Zhuming Zhang, Jing Qin, Tien-Tsin Wong, Guoqiang Han, and Shengfeng He. Example-based colourization via dense encoding pyramids. In Computer Graphics Forum, volume 39, pp. 20-33. Wiley Online Library, 2020.
Liron Yatziv and Guillermo Sapiro. Fast image and video colorization using chrominance blending. IEEE transactions image processing. 150-):12001129, 2006.
Fisher Yu, Ari Seff, Yinda Zhang, Shuran Song, Thomas Funkhouser, and Jianxiong Xiao. Lsun: Construction of a large-scale image dataset using deep learning with humans in the loop. arXiv preprint arXiv: / 500. 00366. 2015.
12

---

Published as a conference paper at ICLR 2021
Bo Zhang, Mingming He, Jing Liao, Pedro V Sander, Lu Yuan, Amine Bermak, and Dong Chen. Deep exemplar- based video colorization. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pp. 8052-8061, 2019a.
Han Zhang, Ian Goodfellow, Dimitris Metaxas, and Augustus Odena. Self-attention generative adversarial networks. In International Conference on Machine Learning, pp. 7354-7363. PMLR, 2019b.
Richard Zhang, Phillip Isola, and Alexei A Efros. Colorful image colorization. In European conference on computer vision, pp. 649-666. Springer, 2016.
Richard Zhang, Jun-Yan Zhu, Phillip Isola, Xinyang Geng, Angela Lin, Tianhe Yu, and Alexei A Efros. Real-time user-guided image colorization with learned deep priors. arXiv preprint 17Xiv:1705.02999, 2017. 2017.
13

---

Published as a conference paper at ICLR 2021
Exponential Moving Averaga: Efrect
19.5
2000 12000 now 13000 mran TecES mmans 2020
16
72
128
256
512
Training steps
K
Figure 6: Left: FID training steps, with and without polyak averaging. Right: The effect of in top-K sampling on FID. See Appendix B
and
ACKNOWLEDGEMENTS
We would like to thank Mohammad Norouzi, Rianne van den Berg, Mostafa Dehghani for their useful comments on the draft and Avital Oliver for assistance the Mechanical Turk setup.
CHANGELOG
v2: Dataset Sharding fix across multiple TPU workers. This changed the FID scores of Col'rran, ColTran-B and ColTran-S from their vl values of 19.71, 21.6 and 21.9 their v2 values of 19.37, 19.98 and 22.06 respecitvely.
A CODE, CHECKPOINTS AND TENSORBOARD FILES
Our implementation is open- sourced in the google-research framework at htttts///iihub..om/ooooo with zip compressed version here, Our full set of
hyperparameters are available here
We provide pre-trained checkpoints of the colorizer and upsamplers ImageNet at Finally, reference tensorboard available at ccolorizer tensorboard] color upsampler tensorboard and spatial upsampler tensorboard
https://consolc.clous cloud. www.:l/ooooooooooocooooooooo files for our training googllc.oom//torggg//ooaeo woo.://tonnggooowe//rrreeeeecece
B EXPONENTIAL MOVING AVERAGE
We found using exponential moving average (EMA) of 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 NUMBER OF PARAMETERS AND INFERENCE SPEED
Inference speed. ColTran core can sample batch of 20 64x64 grayscale images in around 3.5 -5 minutes on a P100 GPU vs PixColor that takes 10 minutes to colorize 28x28 grayscale images on 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, cU; cUj cU! cUg in Table D 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 likely to speed up sampling even more and we leave this engineering optimization for future work.
Number of parameters. ColTran has total of ColTran core (46M) + Color Upsampler (14M) + Spatial Upsampler (14M) = 74M parameters. comparison, PixColor has Conditioning network (44M) Colorizer network (11M)+ Refinement Network (28M) 83M parameters.
14

---

Published as a conference paper at ICLR 2021
Validation
0.65
0.63
0.62
25000
75000
125000
steps
Figure Ablated models. Gated: Gated conditioning layers as done in (Oord et al. 2016) and cArr cMLP, global: Global conditioning instead of pointwise conditioning in cAtl and cLN.
D LOWER COMPUTE REGIME
We retrained the autoregressive colorizer and color upsampler on 4 TPUv2 chips (the lowest con- figuration) with reduced-batch size of 56 and 192 each. For the spatial upsampler, we found that a batch-size was sub-optimal and lead spatial upsampler with axial attention blocks with batch-size of 16 and trained it also 4 TPUv2 chips. The FID drops from 19.71 to 20.9 which is still significantly better than the other models in 2 We note that in this experiment, et al. 2017) uses a total of 16 GPUs.
large deterioration in loss. We thus used a smaller
only 12 TPUv2 chips in total while PixColor (Guadarrama
E
IMPROVED FID WITH TOP-K SAMPLING
We can improve colorization fidelity 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) determined by the predicted probabilities. Our results in Figure K = 4 and = demonstrate a performance improvement over the baseline ColTran model with K 512
F ADDITIONAL ABLATIONS:
Additional ablations of our conditional transformer layers are Figure which did not help.
Conditional transformer layers based on Gated layers (Orrd et al. 2016) (Gated)
A global conditioning layer instead of pointwise conditioning in cAtt and cLN. cAtr + cMLP, global
G AUTOREGRESSIVE MODELS
Autoregressive models are a family of probabilistic methods that model joint distribution of data P(x) or a sequence of symbols X2,,,2,.... I,)) as a product of conditionals II," P(a,(x<)). During training, the input to autoregressive models are the entire sequence of ground-truth symbols. Masking ensures that the contribution all "future" symbols in the sequence are zeroed out. The outputs of the autoregressive model are the corresponding conditional distributions. ((xr)))))) Optimizing the parameters of the autoregressive model proceeds by a standard log- likeeiiodood objective.
Generation happens sequentially, symbol-by byymbbol.. Once symbol x; is generated, the entire sequence (I1, .2,..... xx)) are fed to the autoregressive model to generate 2;-1-
In the case of autoregressive image generation symbols typically correspond to the 3 RGB pixel- channel. These are generated sequentially in raster-scan order, channel by channel and pixel by pixel.
15

---

Published as a conference paper at ICLR 2021
E1
Figure 8: We train colorization model ImageNet and display high resolution colorizations from LSUN
H ROW/COLUMN SELF-ATTENTOON
In the following describe self-attention, that is, omit the height dimension all operations are performed in parallel for each column. Given the representation of a single row within of an image x,, R W xD row-wise self-attention block is applied follows:
qq. k. vl Lught
Ugko E E Dx3Dn
(14)
A softmax (qk /VDD)
A AERW
(15)
SA(x..) Av
(16)
MSA(x;.) (SA,
SSA,(x,))))) Vout
Unut E &D
(17)
LN refers to the application of layer normalization (Ba et al. 2016). Finally, 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 common practice in transformers.
MLP(LIN(x.))
MSA(x; X
(18)
Column-wise self-attention over x. E H works analogously.
I OUT OF DOMAIN COLORIZATIONS
We use our trained colorization model on ImageNet to colorize high-resolution grayscale images from LSU 256 x 256 YYu et al. 2015) and low-resolution grayscale images from Celeb-A (Liu
16

---

Published as a conference paper at ICLR 2021
Figure 9: We train colorization model ImageNet and display low resolution colorizations from Celeb-A
17

---

Published as a conference paper at ICLR 2021
MUS SUAA
Fodd 20.22 Red 10:4
Had BE.A Fand raSo (AAU
Figure 10: Top: Colorizations Bottom: Ground truth. From left to right, our colorizations have a progressively higher fooling rate.
et al. 2015) 64 x 64. Note that these models were trained only on ImageNet and not finetuned on Celeb-A LSUN.
NUMBER OF AXIAL ATTENTION BLOCKS
We did very small hyperparameter sweep using the baseline axial transformer (no conditional layers) with the following configurations:
hidden size 512, number blocks
hidden size 1024, number of blocks 2
hidden size 512, number of blocks 2
Once we found the optimal configuration, we fixed this for all future architecture design.
K ANALYSIS OF MTURK RATINGS
Figure 11: In each column, we display the ground truth followed by 3 samples. Left:: Diverse and real. Center: Realism improves from left to right. Right: Failure cases
Figure 12: We display the per-pixel, maximum predicted probability over 512 colors as a proxy for uncertainty.
We analyzed our samples on the basis of the MTurr ratings in Figure IT the left, we show images, where all the samples have 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 [I] show failure cases where all samples have fool rate of 0 %. For these cases, our model is unable to colorize highly complex structure, that would arguably difficult even for human.
18

---

Published as a conference paper at ICLR 2021
L MORE PROBABILITY MAPS
We display additional probability maps to visualize uncertainty as done in 5.3
M MORE SAMPLES
We display a wide-diversity of colorizations from ColTran that were not
19

---

Published as a conference paper at ICLR 2021
20

---

Published as a conference paper at ICLR 2021
HI iC
HI a
MI 38 N
21

---

Published as a conference paper at ICLR 2021
WAr csty Esttiii
Wiew emuifaitiaii
www.fatt de
22

---

Published as a conference paper at ICLR 2021
3
23

---

Published as a conference paper at ICLR 2021
AURO
Ellle
24


Reviews:
[{"review_id": "v6Tpi142mBb", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "--- Update ---\nThe authors have addressed several concerns that I had regarding the work.  While this is largely an application of a previous method, they have made some application-specific decisions in order to achieve the significant boost in performance on colorization that they saw.  While the metrics for this task are much improved, there are still some things to be desired on the qualitative results (i.e. the diversity of results tends to be in blocks, as opposed to high within-image color variability).  Nevertheless, I think the improvement from this approach my guide future work in this area.  Given the author's responses and changes made, I have amended my recommendation accordingly. \n\n__1.  Summary__\nThe authors propose a method for image colorization based on self-attention largely following the architecture of the Axial Transformer (Ho et al., 2019b).  This approach outperforms several SOTA colorization models on FID and human evaluation.  \n\n\n__2a. Strong Points__\nThe motivation for this work is clear.  Image colorization has many applications and while past approaches have significantly advanced in the past few years, there is certainly much left to be explored in this space.\n\nThe recap/explanation of the Axial Transformer is clear and concise.  My concern (see below) is not with the articulation of this section, but more on the reliance of an approach that hasn’t been accepted via peer review.  \n\nThe performance of this method using both FID as well as using human evaluators is compelling.  \n\nBreaking the problem of colorization into two intermediate low resolution images is a nice approach for enabling larger models.  One question would be how well a single model would perform if smaller images were all that was required.  \n\nThe ablation studies show how different components impact the performance.  \n\n\n__2b. Weak Points__\nAll three modules of this approach are based on method of (Ho et al., 2019b), which is available on arxiv, but was rejected from ICLR 2020.  The current work is focused on the application of that method. This makes for a bit of a tricky situation.  The description of the Axial Transformer is given in section 4, but it is only textual and refers the readers back to the pre-print for more detail.  Since this is the central method of the current work, at a minimum I think it requires more explanation/justification as opposed to pointing to a work that has not been accepted via peer review.  \n\nWhile the language of the paper is fine, the overall flow of the paper is lacking a bit of narrative.  Overall I found myself having to jump around to find the definition and explanation of important things.  Particularly within the description of the model, it would be good to add some language to help the sections flow- currently they feel very independent.  Alternately, if maybe help if in the beginning part of the model, the different model components (fc, fs, etc.) are named there.  Related, the Architecture Section feels out of place after the Model description.  There are references to the attention layers in the model description which are not explored until the Architecture section.  Perhaps it makes sense to put the Architecture section first because it’s addressing layers/mechanisms that span all aspects of the model.  Or perhaps combining the two sections?  Right now it feels like there are two methods sections.\n\nSome of the text around Eqn(7) seems to be missing because the sentence structure doesn't make sense.\n\nIt’s not clear what some of the labels in Figure 3 mean.  You have to go into the text to find out what MLP 4x means, for example, and then when you find it in section 5.2, you have to go back to section 4.3 to actually understand what it means.  \n\nThe ablation studies feel like they’re done in relative isolation.  It would be useful to know, for example, how the lower performance of using the standard Axial Transformer vs. the conditional Axial transformer impacts the final results, not just that portion.  The section “Conditioning Details” in 5.2 just feels like a results dump.  It’s unclear what motivates those particular ablation choices and what those results tell the reader more generally about this approach.  Some kind of context or discussion would be useful.  In general, this section feels like it’s being included just to show that ablation studies were performed without providing any greater understanding as to the approach (to potentially motivate future work or other examples, for instance).  The descriptions are also very terse.  If these experiments add meaningful insight to this approach, then they belong in the main text with additional explanation and discussion.  If they are merely a justification that this approach works, then I would suggest moving most of this section to the appendix and using the space to give better explanation of the methods and results which are central to the application.  \n\nSome of the models which the current method is compared to (Table 2) are not referenced to the best of my knowledge.  What does \"CNN\" mean in this case?  Do all of these methods use a combined spatial and color upsampling method?  If not, how were they implemented?  This is actually a pretty significant issue as it limits the reproducibility of the comparative experiments.  \n\n\n__3. Recommendation__\nReject.  While the results are compelling, the work largely relies on a method which has not been accepted via peer review.  That in and of itself does not warrant rejection, but I believe it contributed to some of the difficulties in explaining the approach, the motivation behind the approach, the results of the ablation studies, etc., which make the paper extremely difficult to follow, likely difficult to build upon, and potentially difficult to reproduce.\n\n\n__4. Recommendation Explanation__\nI would argue that the main goal of this paper is to show a novel application of the Axial Transformer approach of Ho et al 2019b and this is done by adapting that method to the task of Image Colorization.  I would argue the focus is around applying that method, not exclusively doing better Image Colorization, because there is no discussion around how this advances our understanding of image colorization broadly.  Nevertheless, that (showing the usefulness of an approach to a new task) is a valid objective, but because Ho et al 2019b has not been formally accepted, it also somewhat then requires this work to explain and justify approaches of that work.  I believe that challenge has a lot to do with some of the difficulties in the paper around the methods and experiment explanation.  \n\nWhile the (within sentence) language is clear, the overall flow of the paper is  difficult to follow.  It feels like the authors were strongly up-against the page limit, so important explanation and discussion was omitted or made very terse.  For example, the ablation studies, while thorough, sort of feel dumped there.  There's no discussion as to why those and not other experiments were run and what the results of those experiments tell us more broadly.  Similarly, the model and architecture section seem like they should be more intertwined.  As another example, some of the methods in Table 2 are not referenced anywhere and it's not clear how they were used in this context (did they start with a low res image, or high-res image).  That calls the reproducibility of the comparison studies into question.\n\n\n__5.  Questions__\nOverall it seems like every generated image has a red, green, and blue variant.  Were they sampled in a particular manner to guarantee this?  Obviously it is possible to draw other samples, but do they all largely fall into one of these three coarse categories?  When the performance is poor for a given sample, it usually because entire swaths of the image are being painted in with a very non-natural color (like someone’s face being green, or the entire picture having a blue-ish exposure).  Can you speak to this and other common “mistakes” that are observed?  How do these compare with some of the other methods you compared yours against?  Are there simply fewer “mistakes” (i.e. non-natural images), or are the types of imperfections created by this approach different that would warrant different use-cases?\n\nIt seems like a lot of compute (16 TPUv2) was used and the batch size was relatively large.  Is the large batch size necessary for obtaining these results, or could a smaller amount of compute and smaller batch size be used?\n\nWhy does training baselines with 2x and 4x wider MLP dimensions make “a fair comparison”?  Is “Baseline” in Figure 3, x1 (standard) MLP but no conditioning?  Why would x1 be better than x4, but worse than x2?\n\nThe caption of Figure 2 feels a bit imbalanced.  ColTran core is called out specifically, but then the ColTran Upsamplers are not referenced.  Is the “Axial Transformer” just the right branch of the ColTran Core (which the figure seems to suggest) or the entire ColTran core, as the caption seems to suggest.\n\nOn pg. 3 “ColTran Core” it is stated that “we also train a parallel prediction head which we found beneficial for regularization”.  I think it would be useful to given additional explanation here as it’s a fairly significant architectural choice.  If results of not including this head exist, perhaps it would be useful to show this in the appendix.  Otherwise a brief explanation as to why this additional head aids the regularization would be useful.  Since this is an instantiation of the Axial Transformer, is this prediction head added to that approach for this particular task, or is this already a part of the standard Axial Transformer (and therefore maintained here for consistency)?  Ah, this is explored further in section 5.3…. It would be helpful to the reader to reference this section when you introduce the prediction head (i.e. that the impact will be explored in section 5.3).\n\nIn 4.2 it says they “adapt the Axial Transformer model for colorization”.  Can you elaborate on the adaptation?  It’s not clear (without looking up that reference), what belongs to the original approach vs. what was added/changed here for this specific task.\n\nIt feels odd to mention the number of axial attention blocks in the training section as opposed to the model or architecture.  This is a fundamental architectural choice, is it not?  \n\nWhy are the set of models compared via FID and Human Evaluation different?\n\n\n__6.  Feedback__\nThe demonstrated colorization scores and output are compelling, however, I believe the structure of text is very detrimental.  I think it would potentially be feasible to fully rework the text to make it more readable and reproducible and therefore a solid publication because the result is compelling, but as it stands, there is substantial rewritting which would need to be done in my opinion.\n\nIn “Model: ColTran Core” fc is described as a conditional, auto-regressive axial transformer.  While the definition of pc and pc~ are stated thereafter, there is not any further description as to what this means and/or a citation.  The Ho et al. citation is provided in the Figure 2 caption.  At a minimum that citation should be given here as well, but it would be good to give a textual description as to what an “a conditional auto-regressive axial transformer” is since it is not a commonly used architecture.  \n\nThe second paragraph under ColTran Upsamplers (In our experiments…) is slightly confusing.  It seems to suggest that parallel upsampling is sufficient and advantageous for a number of reasons, but that prediction is chosen to reduce color inconsistencies.  Then it seems to go back to again say that Parallel upsampling has a huge advantage of being fast.  This is perhaps also confusing because there is a “Sample” label in Figure 2.  The confusion is less about the validity of the approach and more that the language (in conjunction with the figure) is difficult to follow for someone not already familiar with Guadarrama 2017.\n\nWhile not necessary, it would be interesting to see how this approach performs on out of domain images (i.e. not from ImageNet).  \n\nIn 5.5, it’s stayed you follow the protocol used in PixColor.  It would probably be best to additionally include the citation here, or the citation in place of “PixColor” even though that work is cited near the beginning of the paper (when the reader comes to this section, they may be unfamiliar with this approach and would like to go directly to that reference as opposed to having the search for “PixColor” and then go find the reference).  ", "extra_scores": {}, "extra_text": {}}, {"review_id": "lSNRY4-0hgd", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Thank the authors for addressing reviewers' comments extensively. After rebuttal, I agree with the significance of the proposed method in terms of performance improvement in this particular task. However, the technical novelty is still limited. Thus, I increased my rating to 5.\n\nIn this paper, the authors propose an autoregressive image colorization method based on self-attention. The proposed method first infers an initial low-resolution colorization in an autoregressive manner, then upsamples both spatial resolution and color depth. The authors adopt self-attention to encode contextual information of the scene. Experimental results show that each component of the proposed method is effective and the proposed method outperforms an existing autoregressive method.\n\nOverall, it is difficult to understand the contribution of this paper. I think it is because the writing in Sec. 1 and 2 is unclear. Particularly, the writing of introduction needs a significant improvement as the authors reveal too much details of this paper instead of describing the high-level motivation of the proposed method and the technical contribution. The clarity also needs to be improved in the method and experiment sections. (e.g. ColTran Core in Fig. 2 is confusing. It looks complicated, but the writing is too short. what is the ground-truth in the objective? what about Table 2? each baseline is not explained and cited.)\n\nTechnical novelty is incremental. I could not understand the motivation of the proposed network due to the clarity issue, but this paper generally adopts existing methods such as an autoregressive model and self-attention blocks to apply them to an image colorization problem, which limits the novelty of this paper.\n\nEvaluation is weak. PixColor is an old model (in 2017), so recent methods and state-of-the-art methods should be compared. I could not find out what the baseline methods in Table 2 are, but they do not look like state-of-the-art models. Performance gain over the previous autoregressive model using widely-used self-attention blocks is not enough for accepting this paper.", "extra_scores": {}, "extra_text": {}}, {"review_id": "IUQ6RLQFgRR", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Update: I really appreciate the authors' efforts to address my original concerns. I believe that this work is a nice application of transformers to image colorization. The paper is well-written and the performance of proposed transformer architecture is strong. I think that this work is above the threshold of acceptance.\n\n**Strengths**\nThe motivation of the proposed architecture is reasonable. The paper is generally well-written. \n\n**Major comments**\nIt’s better to include some discussion on regularization effects from Eq. (4). Eq. (4) seems to be helpful to capture the overall structure in an image, rather than capturing only local correlation from autoregressive formulation.\n\nFor upsampling, do we really need to make use of an autoregressive model? A stack of transposed convolutions might be working well, because we only need to upsample input/color resolutions. I totally agree that autoregressive formulation does help achieve better results, but it may be possible to achieve similar performance by just using transposed convolutions. \n\nIt’s better to include more details for reproducing results. \n(1) Even if different batch sizes used (224, 768 and 32), learning rates for all experiments are fixed 3e-4?\n(2) How many epochs or steps are required for convergence?\n(3) Figure 6 shows that EMA is extremely important. How about using cosine annealing for a learning rate scheduler? It may help achieve more robust FID scores without EMA.\n(4) Compared to baselines, this approach is extremely slow due to the autoregressive sampling. It’s better to report inference time.\n\nI'm not sure that conditional layer normalization is indeed helpful. \n\n**Minor comments**\nThe x-axis title of figure 4c (“training steps”) seems to be wrong.", "extra_scores": {}, "extra_text": {}}, {"review_id": "Ch646_YhDE1", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Update: Thanks for the additional ablation studies. I would like to keep my original evaluation which is acceptance.\n\n--------------------\n\nThis paper proposes a transformer architecture for image colorization. It uses an axial transformer to process the low-resolution grayscale image, and uses a conditional version of the axial transformer to predict a low-resolution color image autoregressively conditioned on the gray image. It then uses an axial transformer to predict the final high-resolution output pixels.\n\nPros:\n+ The paper is well-written and easy to read. The literature review is comprehensive.\n+ Image colorization is an important problem in computer vision. To my knowledge, this is the first paper that applies Transformer to colorization. It could potentially be very impacted and inspire future work.\n+ Both automatic metric (FID) and human evaluation are used to compare the method with existing approaches. The performance of the proposed method significantly outperforms the previous state of the art. The qualitative examples are very impressive as well.\n+ The paper performs extensive ablation studies (Figure 3) to verify the contribution of different components. \n\nCons:\n- The technical novelty of this paper is a bit limited. It basically applies existing conditioning techniques to the axial transformer and uses it for image colorization.\n- It seems that no cLN (Fig. 3 mid) is better than cLN with mean-pool only (Fig. 3 right), which is a bit counterintuitive. Any possible explanation? Also, is there a reason to use the globally aggregated context for cLN but not for cMLP/cAtt? An ablation study on that would be helpful. Besides, there is an ablation study on shift-only modulation but I am curious about how scale-only modulation performs.\n- It would be nice to show the number of parameters, training/inference speed of the proposed approach, and compare them to the baselines.\n- Please add references to all baseline methods compared in Table 2. I'm able to find the citation of PixColor in other parts of the paper, but cannot find most of the others'.\n\nMinor problems that do not affect my score:\n- P1: determinisitic -> deterministic\n- The aggregated context is denoted as \\hat{c} in Table 1 but as \\bar{c} in section 4.3.\n- It would be better to use the vector format for Figure 3/4, and enlarge Figure 5 a bit.\n\nOverall, I vote for acceptance. The novelty is not huge but I still think it would be a nice paper for ICLR and have impacts on the field given its strong empirical performance.", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper initially received a mixed rating, with two reviewers rate the paper below the bar and two above the bar. The raised concerns include the need for an autoregressive model for upsampling and the effect of batch sizes. These concerns were well-addressed in the rebuttal. Both of the reviewers that originally rated the paper below the bar raise the scores. After consulting the paper, the reviews, and the rebuttal, the AC agrees that the paper has its merits and is happy to accept the paper.



Decision: Accept (Poster)