Paper: Image as First-Order Norm+Linear Autoregression: Unveiling Mathematical Invariance (ICLR-2024)

Reviewer comment (Presentation):
Some details of the figures are not clear. For example, what do red and blue arrows in z(x, y) mean? How do they affect the reconstruction? Does a single vector q mean the feature of the whole image or the specific pixel? In the figure, the a high frequency in the reconstruction loss.

Author response:
**Q1: Some details of the figures are not clear.**

**For example, what do red and blue arrows in z(x, y) mean?**
The red and blue arrows in $\mathbfit{z}(x, y)$ represent the two autoregression paths of FINOLA for any given position $(x, y)$ from the center where the global pooled vector $\mathbfit{q}$ is placed. Each arrow signifies a FINOLA step, and the color denotes the order of regression. The red path involves horizontal regression followed by vertical regression, while the green path involves vertical regression followed by horizontal regression. The results of these two paths are averaged to obtain the final feature vector.

**How do they affect the reconstruction?**
It's crucial to note that we repeat this "two-path" autoregression until the entire feature map is generated. Subsequently, the feature map goes through the decoder to reconstruct the images.

**Does a single vector q mean the feature of the whole image or the specific pixel?**
*"A single vector q"* indeed denotes the feature of the *entire image* after pooling.

**In the figure, the a high frequency in the reconstruction loss.**
As for the concern about the loss of high-frequency details in the reconstructed images, we acknowledge this limitation, and it is discussed in detail in Appendix B. We attribute this limitation to the choice of the $L_2$ loss function. Our intentional use of a simple loss function is to highlight that FINOLA is not heavily dependent on complex reconstruction loss functions. To improve image quality, an alternative approach could involve exploring the use of a combination of perceptual loss and an adversarial training procedure with a patch-based discriminator, as outlined in the VQGAN framework.

Reviewer comment (Experiments):
The authors mainly compare VQGAN (Esser et al. (2021)) and stable diffusion (Rombach et al. (2021)) in image reconstruction. It would be better to compare recent image reconstruction methods. In addition, could you compare convolutional U-Net for image reconstruction?

Author response:
**Q3: The authors mainly compare VQGAN (Esser et al. (2021)) and stable diffusion (Rombach et al. (2021)) in image reconstruction. It would be better to compare recent image reconstruction methods. In addition, could you compare convolutional U-Net for image reconstruction?**

***Comparison with convolutional U-Net:*** The table below provides a comparative analysis between our method and convolutional U-Net for image reconstruction (measured by PSNR). Both approaches share the same Mobile-Former encoder and have identical bottleneck dimensions (C=1024 or 3072). In our method, FINOLA is employed to generate a 16x16 feature map, followed by a convolutional decoder to reconstruct an image with a size of 256x256. On the other hand, U-Net has a deeper decoder that utilizes convolution and upsampling for image reconstruction. Both methods boast a similar model size. The results showcase the superior performance of FINOLA over convolutional U-Net, indicating that a single layer FINOLA is more effective than multi-layer convolution and upsampling.

|Method|C=1024|C=3072|
|---|---|---|
| Convolutional U-Net | 23.4 | 25.1 |
| **FINOLA (our)** | **23.7** | **25.8** |

***Comparison with ViT-VQGAN:*** In the following comparison, Masked FINOLA is evaluated against ViT-VQGAN (*"Vector-Quantized Image Modeling with Improved VQGAN, Yu et al., ICLR 2022"*) using ImageNet linear probing. FINOLA achieves higher accuracy with a much more compact encoder.

|Method|Params|Top-1|
|---|---|---|
| ViT-VQGAN | 650M |  65.1 |
| **FINOLA (our)** | **28M** | **66.4** |

For a detailed comparison in image reconstruction, we will address this in the subsequent question (**Q4**).

Reviewer comment (Experiments):
When validating the norm+linear approach, the authors repeated q by W × H times. In the original setting, the authors learned two matrices A and B, with more parameters. Thus, the comparison seems a bit unfair.

Author response:
The objective of this comparison is to demonstrate the superiority of using norm+linear over its *absence*. Repetition, requiring zero additional effort, naturally stands out as a straightforward option. In this context, our intention is to specifically exclude the scenario where the zero-effort repetition achieves performance comparable to our more complex norm+linear approach. Such a scenario would contradict the essential need for norm+linear. The poor performance of repetition serves as evidence that norm+linear cannot be easily replaced by a simpler and more straightforward solution, highlighting the necessity of our chosen approach, even in the presence of additional parameters when learning matrices $\mathbfit{A}$ and $\mathbfit{B}$.

Reviewer comment (Experiments):
In comparable performance with Stable Diffusion, the input of the paper method is an image of the same size as the output, while the input of Stable Diffusion is a lower resolution image. The tasks of the two are different, so it may not be appropriate to compare them together.

Author response:
Let's clarify the comparison with Stable Diffusion. Our comparison is specifically with the ***first stage*** of Stable Diffusion, known as *perceptual image compression*, which involves training an autoencoder to reconstruct images. We explicitly refrain from comparing with the second stage of Stable Diffusion, which includes generation of super-resolution from lower-resolution images. The comparison is thus limited to the common task of image reconstruction and does not involve tasks related to super-resolution from lower-resolution inputs.

Reviewer comment (Novelty):
There are previous regression-based generative models (e.g., refer to related works) in feature space, and what are the major difference and advantage of this approach compared with these models? Is it possible to compare with them for the generation quality and computational speed?

Author response:
First and foremost, it's essential to clarify that FINOLA does *not* function as a generative model. Our primary objective is not image generation but rather the revelation of mathematical properties within the latent space. Image reconstruction, one of our tasks (alongside self-supervised learning), serves as a means to validate FINOLA — a first-order norm+linear autoregressive model — as a fundamental mathematical property shared by all images.

It's crucial to note that FINOLA is both functionally and technically distinct from autoregression-based generative models. The latter typically consists of two stages: the first stage involves learning an autoencoder and vector quantization in the latent space, while the second stage focuses on learning autoregression in the quantized latent space. In contrast, FINOLA operates as a single-stage model, seamlessly learning the encoder, autoregression, and decoder end-to-end, with autoregression performed on a continuous (non-quantized) latent space.

Given that FINOLA doesn't fall into the category of generative models, it isn't suitable for direct comparison in terms of generation quality and computational speed with models designed for such purposes.

Reviewer comment (Theory):
Are the matrix A and B shared for all different images and pixels in the feature space? If it is, why learned constant A, and B can deduce good feature regression?

Author response:
Yes, the matrices $\mathbfit{A}$ and $\mathbfit{B}$ are shared for all images and pixels in the feature space.

The success of using constant $\mathbfit{A}$ and $\mathbfit{B}$ for achieving good feature regression can be attributed to three primary factors. Firstly, images inherently exhibit significant spatial redundancy, leading to strong correlations between consecutive positions within the feature map. Secondly, our hypothesis regarding the existence of a feature map where spatial derivatives correlate with feature values appears to be supported. Lastly, the effectiveness of deep learning allows us to validate this hypothesis by seamlessly learning the encoder, FINOLA, and decoder components in an end-to-end manner.

Reviewer comment (Experiments):
The paper states that this work does not aim to achieve SoTA results, however, comparisons with SoTA regression models or other variants, e.g., using nonlinear regression instead of linear transform, should be able to better give insights to audience.

Author response:
Excellent point! We have conducted a comparison between our norm+linear approach and non-linear regression, which involves replacing the linear transform with two MLP layers incorporating GELU activation in between. This evaluation was performed for both 1024 and 3072 channels. Surprisingly, our norm+linear, despite having fewer parameters, closely trails the performance of the two-layer MLP. This result underscores the effectiveness of norm+linear in modeling the transition between pixels in the feature map.

|Regression|C=1024|C=3072|
|---|---|---|
|Non-Linear (2 layer MLP)|23.8|25.8|
|**Linear (our)**|23.7|25.8|

We further benchmark our method against an alternative non-linear solution: employing a multi-layer convolutional approach to upsample the vector $\mathbfit{q}$ to a 16x16 feature map. This method undergoes four upsample stages, each incorporating three 3x3 convolution layers. The results below  demonstrate the superior performance of a single layer FINOLA over the multi-layer convolutional baseline.

|Regression|C=1024|C=3072|
|---|---|---|
| convolution | 23.4 | 25.1 |
| **FINOLA (our)** | **23.7** | **25.8** |

Reviewer comment (Evaluation):
The regression model should gradually regression dense feature maps? How about the computational overhead/speed for training and inference using this model?

Author response:
***Autoregressive reconstruction:*** The training of the autoregressive model involves regressing dense feature maps, and the computational requirements increase with the size of the feature map. For instance, training on a 16x16 feature map with 3072 latent channels for 100 epochs on ImageNet takes approximately 8 days with 8 V100 GPUs, and extending to a larger feature map like 64x64 increases the training time to 18 days using the same GPU setup.

The complete inference pipeline, encompassing encoding, FINOLA, and decoding, for a 256x256 image on a MacBook Air with an Apple M2 CPU, takes approximately 1.2 seconds.

***Self-supervised learning:*** Training a self-supervised learning model (Masked-FINOLA-B) for 800 epochs on ImageNet demands around 10 days with 8 V100 GPUs and a batch size of 128 per GPU.

It's worth noting that following self-supervised learning, only the encoder is typically used for inference in downstream tasks. Our method utilizes Mobile-Former as an encoder, which is a lightweight architecture efficiently combining MobileNet and Transformer.  The original paper *"Mobile-Former: Bridging MobileNet and Transformer, CVPR 2022"* shows detailed information on runtime performance (comparable with MobileNet and ShuffleNet).

Reviewer comment (Theory):
What is the meaning of the mathematical invariance in this paper for the proposed model?

Author response:
The term "mathematical invariance" refers to the ***consistent relationship shared by all images between feature values and their spatial derivatives on the feature map***. Mathematically, for any given image, the feature values $\mathbfit{z}(x, y) \in \mathbb{R}^C$ at any position $(x, y)$ on the feature map determine its spatial derivatives as follows:

$\mathbfit{z}(x + 1, y) - \mathbfit{z}(x, y) = \mathbfit{A}\mathbfit{z}_n(x, y)$,

$\mathbfit{z}(x, y + 1) - \mathbfit{z}(x, y) = \mathbfit{B}\mathbfit{z}_n(x, y)$,

where $\mathbfit{z}_n(x, y) = \frac{\mathbfit{z}(x, y) - \mu_z}{\sigma_z}$.

It's essential to note that matrices $\mathbfit{A}$ and $\mathbfit{B}$ exhibit ***invariance*** across different images and pixels on the feature map, emphasizing the consistent mathematical relationship that holds universally.

Reviewer comment (Theory):
Justifications for the novel design choice (but seems heuristic) compared to the regular AR model, including predefined assignments in Block-wise Masked FINOLA, predicting three points in Block-wise Masked FINOLA, and why first order.

Author response:
Thanks for the suggestion. Below we provide justification for each item.

***Predefined assignments and predicting three points in Block-wise Masked FINOLA***

This design is meticulously crafted to achieve several key objectives: (a) Full coverage: ensuring that all masked positions receive attention, (b) Balance: distributing the usage of unmasked positions equally, and (c) Efficiency: utilizing each unmasked position to predict three masked positions, especially under the condition where 75% of positions are masked. Additionally, this design can be implemented in a block-wise manner to further enhance efficiency, a crucial consideration for the extended training durations involved in self-supervised learning.

***Why first order?***

The selection of first-order autoregression is driven by its inherent simplicity and local property, where the derivatives along the $x$ and $y$ axes are exclusively determined by the current position. This strategic choice not only simplifies the computational process but also aims to unveil valuable mathematical insights that could be shared across different images, provided the approach proves effective."

Reviewer comment (Theory):
Lack of interpretation of the derived PDE.

Author response:
The derived partial differential equations (PDEs) in Eq. 4 provide a theoretical extension of FINOLA from a discrete grid to continuous coordinates, establishing a conceptual bridge between the discrete representations of neural networks and the continuous nature of human vision. Although our digital images and feature maps are inherently discrete, the human vision system perceives continuous signals.

This theoretical extension allows us to hypothesize about the mathematical properties of continuous feature representations in the human vision system. We posit that the proposed PDEs in Eq. 4, derived from their discrete counterparts in Eq. 1, capture the patterns inherent in continuous features within human vision.

This is related to **Q2** below.

Reviewer comment (Theory):
(page 2) The authors argued that “the coefficient matrices $A$ and $B$ capture the relationship between each position“. How do the coefficient matrices directly know the spatial relationship? They do not take any of the spatial information. The reviewer believes there are some pieces of missing information: (1) coefficient matrices $A$ and $B$ capture channel-wise relationship, (2) the pattern and correlation encoded in the channel are related to the positional information, and (3) the coefficient matrices $A$ and $B$ (indirectly) capture the relationship between each position. The reviewer notes this may be trivial for some readers but not for them.

Author response:
The coefficient matrices $\mathbfit{A}$ and $\mathbfit{B}$ encode the relationship between a ***feature vector*** $\mathbfit{z}(x, y) \in \mathbb{R}^C$ and and ***its spatial derivatives***, namely $\mathbfit{z}(x+1, y)-\mathbfit{z}(x, y)$ and $\mathbfit{z}(x, y+1)-\mathbfit{z}(x, y)$. This relationship is mathematically expressed as:

$\mathbfit{z}(x + 1, y) - \mathbfit{z}(x, y) = \mathbfit{A}\mathbfit{z}_n(x, y)$,

$\mathbfit{z}(x, y + 1) - \mathbfit{z}(x, y) = \mathbfit{B}\mathbfit{z}_n(x, y)$,

where $\mathbfit{z}_n(x, y) = \frac{\mathbfit{z}(x, y) - \mu_z}{\sigma_z}$.

It's essential to note that matrices $\mathbfit{A}$ and $\mathbfit{B}$ exhibit invariance across different images and positions $(x,y)$ within the feature map once they are learned from the data.

***1.1 Coefficient matrices $\mathbfit{A}$ and $\mathbfit{B}$ capture channel-wise relationship.***

Indeed, the difference of the channel vector between consecutive positions, such as $\mathbfit{z}(x + 1, y) - \mathbfit{z}(x, y)$, represents a linear combination of normalized channel values at the current position $\mathbfit{z}_n(x, y) \in \mathbb{R}^C$.

***1.2: The pattern and correlation encoded in the channel are related to the positional information.***

The correlation does *NOT* pertain to the positional information; instead, it is associated with the derivatives (or rate of change) along the $x$ and $y$ axes, often referred to as *"spatial derivatives"*.

***1.3: the coefficient matrices $\mathbfit{A}$ and $\mathbfit{B}$ (indirectly) capture the relationship between each position. Maybe it is trivial for some readers but it is not for me.***

As discussed earlier, the matrices $\mathbfit{A}$ and $\mathbfit{B}$ explicitly capture the relationship between the feature at a given position $\mathbfit{z}(x, y)$ and its derivatives along the $x$ and $y$ axes. Consequently, this encoding implicitly extends to capturing the relationship between consecutive positions, such as from $\mathbfit{z}(x, y)$ to $\mathbfit{z}(x + 1, y)$.

Reviewer comment (Reproducibility):
The specific method of the block-wise Masked FINOLA is unclear. For example, if we see the Corner case, let’s say an input coordinate is (3,2) and it is used for predicting {(11,10), (11,2), (3,10)}. The reviewer guesses that for obtaining (11,2), for instance, the function $\phi^8(z(3,2))$ is applied, which means (4,2) is predicted first and it is used as an input for predicting (5,2) and so forth. The question is, considering that we already have a ground truth within (0,2),(1,2) …, (7,2), why not use the ground truth information?

Author response:
Great observation! The current design of block-wise Masked FINOLA intentionally avoids utilizing available ground truth along the predicting path (e.g., excluding (0,2), (1,2), ..., (7,2)). This decision stems from careful considerations related to both *training speed* and the *difficulty of the prediction task*.

***Training speed:*** Employing multiple ground truths in the unmasked block to predict every masked position can be slow. In our current solution, each masked position is predicted from a single unmasked position, enhancing efficiency. Furthermore, the prediction process can be implemented in parallel at the group level. For instance, applying horizontal FINOLA eight times on the whole block (0:8, 0:8) to predict the block on the right (8:16, 0:8) is possible because all unmasked positions within a group (0:8, 0:8) share offsets towards their assigned masked positions. This parallelism is crucial for self-supervised pre-training, which requires an extended training period (800 or 1600 epochs). However, this advantage is not applicable when involving multiple ground truths to predict every masked position.

***Difficulty of prediction:*** Using closer ground truth (e.g., predicting (11,2) from (7,2) rather than from (3,2)) reduces the prediction challenge, potentially compromising representation learning. We conducted an experiment by constraining the prediction offset to less than 4. For instance, if the unmasked block is (0:8, 0:8), the final feature map is (0:12, 0:12) rather than (0:16, 0:16). Specifically, we used (4:8, 0:8) to predict (8:12, 0:8), (0:8, 4:8) to predict (0:8, 8:12), and (4:8, 4:8) to predict (8:12, 8:12). The linear probing performance dropped from 66.4 to 63.9.

The design of block-wise Masked FINOLA carefully balances considerations of training speed and pre-training difficulty. We appreciate your feedback and will enhance the clarity of these details in the final draft.

Reviewer comment (Reproducibility):
(Table 4 (a)) Even though Stable Diffusion is a Generative Model, the reviewer believes it should have better PSNR than FINOLA for the image reconstruction task. How was the image reconstruction task implemented for Stable Diffusion?

Author response:
It's important to note that even though Stable Diffusion is a generative model, the comparison in this context was limited to the performance of the ***first stage***, referred to as *perceptual image compression*. In this stage, an autoencoder is trained to reconstruct images. The performance metrics for Stable Diffusion were obtained from the original paper titled *"High-Resolution Image Synthesis with Latent Diffusion Models, CVPR 2022,"* specifically from Table 8 in Appendix D. The latent size chosen for this comparison was 16x16 ($f=16$).

Reviewer comment (Experiments):
How fast is the parallel implementation (Fig. 3) compared to the regular AR setting?

Author response:
The table below presents a comparison between our parallel implementation and the regular autoregressive (AR) setting in terms of inference runtime. The runtime evaluation considers the complete inference pipeline, including encoding, autoregression, and decoding, for a 256x256 image on a MacBook Air with an Apple M2 CPU. The comparison is conducted for two feature map sizes, 16x16 and 64x64.

|AR|16x16|64x64|
|---|---|---|
|Regular AR|1.7s|7.7s|
|**Parallel (our)**|**1.2s**|**2.6s**|

The parallel implementation demonstrates superior speed compared to the regular AR setting. It achieves a 30% speedup at a resolution of 16x16 and a threefold increase in speed at a higher resolution of 64x64.

Reviewer comment (Evaluation):
Is this method fast enough to use as a feature extractor for the downstream task?

Author response:
Yes, our method utilizes Mobile-Former as an encoder, which is a lightweight architecture efficiently combining MobileNet and Transformer.  The original paper "*Mobile-Former: Bridging MobileNet and Transformer, CVPR 2022*" shows detailed information on runtime performance (comparable with MobileNet and ShuffleNet). This paper demonstrates its effective performance for downstream tasks like image classification and object detection, after FINOLA pre-training.