Venue: ICLR 2021 Conference

Reviews:

final_score: rating: 5: Marginally below acceptance threshold
review_content: 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.

In 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.

Overall, 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.)

Technical 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.

Evaluation 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.
review_title: Lack of clarity and novelty, weak evaluation

final_score: rating: 7: Good paper, accept
review_content: 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.

**Strengths**
The motivation of the proposed architecture is reasonable. The paper is generally well-written. 

**Major comments**
It’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.

For 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. 

It’s better to include more details for reproducing results. 
(1) Even if different batch sizes used (224, 768 and 32), learning rates for all experiments are fixed 3e-4?
(2) How many epochs or steps are required for convergence?
(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.
(4) Compared to baselines, this approach is extremely slow due to the autoregressive sampling. It’s better to report inference time.

I'm not sure that conditional layer normalization is indeed helpful. 

**Minor comments**
The x-axis title of figure 4c (“training steps”) seems to be wrong.
review_title: Reasonable approach / Well-written / Better than baselines

final_score: rating: 6: Marginally above acceptance threshold
review_content: --- Update ---
The 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. 

__1.  Summary__
The 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.  


__2a. Strong Points__
The 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.

The 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.  

The performance of this method using both FID as well as using human evaluators is compelling.  

Breaking 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.  

The ablation studies show how different components impact the performance.  


__2b. Weak Points__
All 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.  

While 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.

Some of the text around Eqn(7) seems to be missing because the sentence structure doesn't make sense.

It’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.  

The 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.  

Some 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.  


__3. Recommendation__
Reject.  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.


__4. Recommendation Explanation__
I 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.  

While 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.


__5.  Questions__
Overall 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?

It 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?

Why 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?

The 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.

On 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).

In 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.

It 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?  

Why are the set of models compared via FID and Human Evaluation different?


__6.  Feedback__
The 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.

In “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.  

The 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.

While not necessary, it would be interesting to see how this approach performs on out of domain images (i.e. not from ImageNet).  

In 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).  
review_title: Image colorization based on self-attention

final_score: rating: 7: Good paper, accept
review_content: Update: Thanks for the additional ablation studies. I would like to keep my original evaluation which is acceptance.

--------------------

This 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.

Pros:
+ The paper is well-written and easy to read. The literature review is comprehensive.
+ 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.
+ 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.
+ The paper performs extensive ablation studies (Figure 3) to verify the contribution of different components. 

Cons:
- 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.
- 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.
- It would be nice to show the number of parameters, training/inference speed of the proposed approach, and compare them to the baselines.
- 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'.

Minor problems that do not affect my score:
- P1: determinisitic -> deterministic
- The aggregated context is denoted as \hat{c} in Table 1 but as \bar{c} in section 4.3.
- It would be better to use the vector format for Figure 3/4, and enlarge Figure 5 a bit.

Overall, 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.
review_title: Extensive experiments and strong performance, novelty is a bit incremental.

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)