Original: Stochastic Window Transformer for Image Restoration
Revised: Stochastic Window Transformer for Image Restoration

Original: Abstract
Revised: Abstract

Original: Thanks to the powerful representation capabilities, transformers have made impressive progress in image restoration.
Revised: Thanks to the powerful representation capabilities, transformers have made impressive progress in image restoration.

Original: However, existing transformers-based methods do not carefully consider the particularities of image restoration.
Revised: However, existing transformers-based methods do not carefully consider the particularities of image restoration.

Original: In general, image restoration requires that an ideal approach should be translation-invariant to the degradation, i.e., the undesirable degradation should be removed irrespective of its position within the image.
Revised: In general, image restoration requires that an ideal approach should be translation-invariant to the degradation, i.e., the undesirable degradation should be removed irrespective of its position within the image.

Original: Furthermore, the local relationships also play a vital role, which should be faithfully exploited for recovering clean images.
Revised: Furthermore, the local relationships also play a vital role, which should be faithfully exploited for recovering clean images.

Original: Nevertheless, most transformers either adopt a fixed local window strategy ora global attention mechanism, which unfortunately breaks the translation invariance and causes huge loss of local relationships.
Revised: Nevertheless, most transformers either adopt local attention with the fixed local window strategy or global attention, which unfortunately breaks the translation invariance and causes huge loss of local relationships.

Original: To address these issues, we propose an elegant stochastic window strategy for transformers.
Revised: To address these issues, we propose an elegant stochastic window strategy for transformers.

Original: Specifically, we first introduce the window partition with stochastic shift to replace the original fixed window partition for training.
Revised: Specifically, we first introduce the window partition with stochastic shift to replace the original fixed window partition for training.

Original: Then, we design a new layer expectation propagation algorithm to efficiently approximate the expectation of the induced stochastic transformer for testing.
Revised: Then, we design a new layer expectation propagation algorithm to efficiently approximate the expectation of the induced stochastic transformer for testing.

Original: Our stochastic window transformer not only enjoys powerful representation but also maintains the desired property of translation invariance and locality.
Revised: Our stochastic window transformer not only enjoys powerful representation but also maintains the desired property of translation invariance and locality.

Original: Experiments validate the stochastic window strategy consistently improves performance on various image restoration tasks (deraining, denoising and deblurring) by significant margins.
Revised: Experiments validate the stochastic window strategy consistently improves performance on various image restoration tasks (deraining, denoising and deblurring) by significant margins.

Original: The code is available at https://github.com/jiexiaou/Stoformer .
Revised: The code is available at https://github.com/jiexiaou/Stoformer .

Original:  Introduction
Revised:  Introduction

Original: Image restoration aims to recover latent clean images from their noise-polluted counterparts, which lays the foundation for various vision tasks.
Revised: Image restoration aims to recover latent clean images from their noise-polluted counterparts, which lays the foundation for various vision tasks.

Original: Generally, image restoration methods should satisfy a constraint: an ideal approach should remove undesirable degradation irrespective of its position within the image.
Revised: Generally, image restoration methods should satisfy a constraint: an ideal approach should remove undesirable degradation irrespective of its position within the image.

Original: In other words, maintaining the translation invariance is the fundamental requirement for image restoration, which makes convolutional neural networks (CNNs) well-suited for this specific vision task [56, 23, 9, 73, 70, 61, 66, 58, 72, 6, 15].
Revised: In other words, maintaining the translation invariance is the fundamental requirement for image restoration, which makes convolutional neural networks (CNNs) well-suited for this specific vision task [56, 23, 9, 73, 70, 61, 66, 58, 72, 6, 15].

Original: In comparison with multi-layer perceptrons (MLPs), the distinct characteristics of CNNs [26, 22, 17] are their built-in locality and weight sharing, resulting in the desirable property of translation invariance.
Revised: In comparison with multi-layer perceptrons (MLPs), the distinct characteristics of CNNs [26, 22, 17] are their built-in locality and weight sharing, resulting in the desirable property of translation invariance.

Original: These two priors play a vital role in image restoration, i.e., pixels within local regions tend to exhibit strong correlations while translation invariance is one of the desiderata for the potential image restoration method.
Revised: These two priors play a vital role in image restoration, i.e., pixels within local regions tend to exhibit strong correlations while translation invariance is one of the desiderata for the potential image restoration method.

Original: On the other hand, transformer-based image restoration methods [2, 54, 33, 62] are not compatible with the translation invariance and locality.
Revised: On the other hand, transformer-based image restoration methods [2, 54, 33, 62] are not compatible with the translation invariance and locality.

Original: The patterns that transformers adopted for recovering high-resolution images can be summarized as small patch with global attention (e.g., 48 × 48 for IPT [2]) and large patch with local attention (e.g., Uformer [54] and SwinIR [33]).
Revised: The patterns that transformers adopted for recovering high-resolution images can be summarized as small patch with global attention (e.g., 48 × 48 for IPT [2]) and large patch with local attention (e.g., Uformer [54] and SwinIR [33]).

Original: Nevertheless, neither of these two patterns can meet the requirements of translation invariance or locality.
Revised: Nevertheless, neither of these two patterns can meet the requirements of translation invariance or locality.

Original: Here we detail the latter pattern while the former one is included in the supplemental material.
Revised: Here we detail the latter pattern while the former one is included in the supplemental material.

Original: The latter pattern
Revised: The latter pattern

Original: [Figure]
Revised: [Figure]

Original: partitions the feature map into fixed non-overlapped windows and self-attention is restricted inside these local windows.
Revised: partitions the feature map into fixed non-overlapped windows and self-attention is restricted inside these local windows.

Original: However, we argue that the fixed partition strategy will break the translation invariance and further lead to huge loss of local relationships.
Revised: However, we argue that the fixed partition strategy will break the translation invariance and further lead to huge loss of local relationships.

Original: As shown in Figure 1(a), for the fixed window partition strategy, translation invariance can only be ensured among specific partitioned windows (blue windows) 2 rather than general windows with arbitrary shifts (red windows), which disagrees with the translation invariance.
Revised: As shown in Figure 1(a), for the fixed window partition strategy, translation invariance can only be ensured among specific partitioned windows (blue windows) 2 rather than general windows with arbitrary shifts (red windows), which disagrees with the translation invariance.

Original: The reason behind this is that the fixed window partition imposes the artificial prior to windows with specific shifts, which turns out to break the translation invariance.
Revised: The reason behind this is that the fixed window partition imposes the artificial prior to windows with specific shifts, which turns out to break the translation invariance.

Original: Figure 1(b) intuitively illustrates that the fixed window partition cannot capture intact local relationships.
Revised: Figure 1(b) intuitively illustrates that the fixed window partition cannot capture intact local relationships.

Original: Without loss of generality, we consider the Euclidean distance, based on which the neighbor space will be the region encompassed by a circle, in 2 D plane to measure locality.
Revised: Without loss of generality, we consider the Euclidean distance, based on which the neighbor space will be the region encompassed by a circle, in 2 D plane to measure locality.

Original: For a certain token (central red dot), many tokens (green dots) within its neighbor space (encompassed by the red circle with radius r ) cannot participate in the calculation of self-attention, which leads to tremendous loss of local relationships.
Revised: For a certain token (central red dot), many tokens (green dots) within its neighbor space (encompassed by the red circle with radius r ) cannot participate in the calculation of self-attention, which leads to tremendous loss of local relationships.

Original: As the neighbor space expands, i.e., increasing radius r , the loss will be more severe.
Revised: As the neighbor space expands, i.e., increasing radius r , the loss will be more severe.

Original: Figure 1(c) quantitatively depicts the proportion of local relationships that are captured by the fixed window partition over the total relationships existing in the neighbor space.
Revised: Figure 1(c) quantitatively depicts the proportion of local relationships that are captured by the fixed window partition over the total relationships existing in the neighbor space.

Original: Ideally, when the feature map is divided into non-overlapping windows of size s × s , it is expected that the relationships of token pairs with distance less than s can be faithfully captured.
Revised: Ideally, when the feature map is divided into non-overlapping windows of size s × s , it is expected that the relationships of token pairs with distance less than s can be faithfully captured.

Original: However, we can observe that there exists significant loss of local relationships even within the neighbor space whose radius is much smaller than s .
Revised: However, we can observe that there exists significant loss of local relationships even within the neighbor space whose radius is much smaller than s .

Original: A seemingly straightforward solution to this problem is to adopt the sliding window strategy to scan the whole feature map using the local window like CNNs.
Revised: A seemingly straightforward solution to this problem is to adopt the sliding window strategy to scan the whole feature map using the local window like CNNs.

Original: However, this strategy incurs huge overhead in terms of both memory and computing speed.
Revised: However, this strategy incurs huge overhead in terms of both memory and computing speed.

Original: In this work, we propose a novel stochastic window strategy to impart translation invariance to the transformer and make it faithfully model local relationships.
Revised: In this work, we propose a novel stochastic window strategy to impart translation invariance to the transformer and make it faithfully model local relationships.

Original: Specifically, instead of partitioning the feature map into fixed non-overlapped windows, we choose to cover the whole feature map with a stochastically shifted window partition (see Figure 2).
Revised: Specifically, instead of partitioning the feature map into fixed non-overlapped windows, we choose to cover the whole feature map with a stochastically shifted window partition (see Figure 2).

Original: By introducing stochastic shifts, all windows are treated equally so that the translation invariance and locality can be ensured.
Revised: By introducing stochastic shifts, all windows are treated equally so that the translation invariance and locality can be ensured.

Original: As shown in Figure 1(c), our stochastic window strategy is able to capture complete local relationships until the radius of neighbor space reaches the window size.
Revised: As shown in Figure 1(c), our stochastic window strategy is able to capture complete local relationships until the radius of neighbor space reaches the window size.

Original: It is noteworthy that unlike the sliding window strategy, the time and memory overhead consumed for training stochastic window transformer is comparable to the fixed window strategy.
Revised: It is noteworthy that unlike the sliding window strategy, the time and memory overhead consumed for training stochastic window transformer is comparable to the fixed window strategy.

Original: For testing, we further design the layer expectation propagation algorithm to approximately marginalize the stochastic shifts, during which the translation invariance and locality can also be ensured.
Revised: For testing, we further design the layer expectation propagation algorithm to approximately marginalize the stochastic shifts, during which the translation invariance and locality can also be ensured.

Original: In conclusion, the contributions are threefold:
Revised: In conclusion, the contributions are threefold:

Original: • We analyze the phenomena of translation invariance breaking and local relationships loss in existing transformer-based image restoration approaches.
Revised: • We analyze the phenomena of translation invariance breaking and local relationships loss in existing transformer-based image restoration approaches.

Original: To the best of our knowledge, this is the first work to point out the importance of translation invariance for image restoration.
Revised: To the best of our knowledge, this is the first work to point out the importance of translation invariance for image restoration.

Original: These windows are processed using shared weights, e.g., W Q , W K , W V for Q, K, V .
Revised: These windows are processed using shared weights, e.g., W Q , W K , W V for Q, K, V .

Original: • We propose a new stochastic window strategy, which comprises stochastic windows for training and layer expectation propagation for testing, to compensate for the broken translation invariance of transformers and enable them to faithfully model local relationships.
Revised: • We propose a new stochastic window strategy, which comprises stochastic windows for training and layer expectation propagation for testing, to compensate for the broken translation invariance of transformers and enable them to faithfully model local relationships.

Original: • Extensive experiments on various tasks, e.g., deraining, denoising and deblurrring, validate that the restoration performance can be constantly improved by equipping our stochastic window strategy.
Revised: • Extensive experiments on various tasks, e.g., deraining, denoising and deblurrring, validate that the restoration performance can be constantly improved by equipping our stochastic window strategy.

Original:  Related works
Revised:  Related works

Original: Image Restoration.
Revised: Image Restoration.

Original: Image Restoration [67, 11, 5, 36, 34, 38, 63, 40, 29, 28] aims to restore the clean image from its degraded version.
Revised: Image Restoration [67, 11, 5, 36, 34, 38, 63, 40, 29, 28] aims to restore the clean image from its degraded version.

Original: In recent years, remarkable progress against traditional modeldriven methods [4, 53, 16] has been achieved due to the development of deep learning technologies [27], especially CNNs [26, 22, 17].
Revised: In recent years, remarkable progress against traditional modeldriven methods [4, 53, 16] has been achieved due to the development of deep learning technologies [27], especially CNNs [26, 22, 17].

Original: Instead of relying on preset image priors, learning based methods directly learn to project from noisy to clean ones from a large collection of noisy-clean image pairs.
Revised: Instead of relying on preset image priors, learning based methods directly learn to project from noisy to clean ones from a large collection of noisy-clean image pairs.

Original: Numerous representative CNNs have sprung up across various classical image restoration tasks, including image denoising [25, 60, 68, 66, 7, 74, 32], image super-resolution [72, 6, 15, 20, 75], image deblurring [9, 23, 39, 23, 47, 65, 43, 42], image deraining [58, 50, 13, 64, 51, 41, 30, 14], etc.
Revised: Numerous representative CNNs have sprung up across various classical image restoration tasks, including image denoising [25, 60, 68, 66, 7, 74, 32], image super-resolution [72, 6, 15, 20, 75], image deblurring [9, 23, 39, 23, 47, 65, 43, 42], image deraining [58, 50, 13, 64, 51, 41, 30, 14], etc.

Original: In general, the translation invariance (derived from weight sharing) and locality have been hard-coded into the inherent structure of CNNs so that it seems CNNs are well-suited for image restoration tasks.
Revised: In general, the translation invariance (derived from weight sharing) and locality have been hard-coded into the inherent structure of CNNs so that it seems CNNs are well-suited for image restoration tasks.

Original: However, compared with Transformer, CNNs are restricted by their limited flexibility.
Revised: However, compared with Transformer, CNNs are restricted by their limited flexibility.

Original: Vision Transformer.
Revised: Vision Transformer.

Original: Recently, with great success of transformer [49, 8] in the NLP field, Vision
Revised: Recently, with great success of transformer [49, 8] in the NLP field, Vision

Original: Transformers [10, 48, 52, 35, 57, 59, 3] have also been prevalent in vision community.
Revised: Transformers [10, 48, 52, 35, 57, 59, 3] have also been prevalent in vision community.

Original: ViT [10] treated image patches as token sequence and applied the vanilla transformer on it for image classification.
Revised: ViT [10] treated image patches as token sequence and applied the vanilla transformer on it for image classification.

Original: With the goal of bringing in reasonable priors into transformer to improve efficiency, Swin [35] introduced local window based attention and established a hierarchical architecture.
Revised: With the goal of bringing in reasonable priors into transformer to improve efficiency, Swin [35] introduced local window based attention and established a hierarchical architecture.

Original: Inspired by the key insights from high-level vision, a few transformers [2, 33, 54, 62, 55] for low-level vision have arose.
Revised: Inspired by the key insights from high-level vision, a few transformers [2, 33, 54, 62, 55] for low-level vision have arose.

Original: But most of them directly transfer high-level designs (e.g., local attention) into low-level vision without careful consideration of its particularities.
Revised: But most of them directly transfer high-level designs (e.g., local attention) into low-level vision without careful consideration of its particularities.

Original: As we discuss above, unlike CNNs, transformers do not possess the translation invariance and intact locality.
Revised: As we discuss above, unlike CNNs, transformers do not possess the translation invariance and intact locality.

Original: In general, low-level tasks usually require to accomplish pixel-level regression, where more strict locality and translation invariance are expected.
Revised: In general, low-level tasks usually require to accomplish pixel-level regression, where more strict locality and translation invariance are expected.

Original: Besides, the CNNs-style sliding window strategy will incur huge burden in terms of memory and computing speed.
Revised: Besides, the CNNs-style sliding window strategy will incur huge burden in terms of memory and computing speed.

Original: Therefore, an efficient mechanism to remedy the broken translation invariance and locality of transformer is essential for low-level vision.
Revised: Therefore, an efficient mechanism to remedy the broken translation invariance and locality of transformer is essential for low-level vision.

Original:  Stochastic Window Transformer for Image Restoration
Revised:  Stochastic Window Transformer for Image Restoration

Original: Transformers have attained impressive performance for image restoration due to their strong representation ability.
Revised: Transformers have attained impressive performance for image restoration due to their strong representation ability.

Original: Given the quadratic complexity in both computation and memory usage, transformers for image restoration tend to employ the local attention.
Revised: Given the quadratic complexity in both computation and memory usage, transformers for image restoration tend to employ the local attention.

Original: Except for high-efficiency, local attention should have been designed to model local relationships.
Revised: Except for high-efficiency, local attention should have been designed to model local relationships.

Original: However, as we analyzed before, local attention 𝑠 𝑠 breaks the ideal translation invariance and further leads to huge loss of local relationships.
Revised: However, as we analyzed before, local attention 𝑠 𝑠 breaks the ideal translation invariance and further leads to huge loss of local relationships.

Original: As shown in Figure 2, for a feature map, the window partition with arbitrary shift ( ξ h , ξ w ) contains a comparable quantity of local relationships and implies the translation invariance among divided windows.
Revised: As shown in Figure 2, for a feature map, the window partition with arbitrary shift ( ξ h , ξ w ) contains a comparable quantity of local relationships and implies the translation invariance among divided windows.

Original: In other words, all the window partitions are equally informative in terms of translation invariance and locality and should be treated equally.
Revised: In other words, all the window partitions are equally informative in terms of translation invariance and locality and should be treated equally.

Original: Therefore, it is unreasonable for the fixed window strategy to express infinite favoritism towards the certain partition (e.g., (0 , 0) or ( s 2 , s 2 ) ) while simply discarding others.
Revised: Therefore, it is unreasonable for the fixed window strategy to express infinite favoritism towards the certain partition (e.g., (0 , 0) or ( s 2 , s 2 ) ) while simply discarding others.

Original: In this work, we propose the stochastic window strategy, by which all the window partitions are treated fairly.
Revised: In this work, we propose the stochastic window strategy, by which all the window partitions are treated fairly.

Original: The network can be trained based on local attention but with stochastic rather than fixed shift.
Revised: The network can be trained based on local attention but with stochastic rather than fixed shift.

Original: Therefore, transformer with the stochastically shifted window can be trained as efficiently as the fixed window.
Revised: Therefore, transformer with the stochastically shifted window can be trained as efficiently as the fixed window.

Original: At test time, we propose the layer expectation propagation algorithm to approximate the expectation of the introduced stochastic shift, which also helps to maintain the desired property of translation invariance and locality.
Revised: At test time, we propose the layer expectation propagation algorithm to approximate the expectation of the introduced stochastic shift, which also helps to maintain the desired property of translation invariance and locality.

Original: 3.1 Stochastic Window Strategy
Revised: 3.1 Stochastic Window Strategy

Original: Stochastic window strategy aims to impart the translation invariance to transformer and make it faithfully exploit local relationships.
Revised: Stochastic window strategy aims to impart the translation invariance to transformer and make it faithfully exploit local relationships.

Original: Canonical transformer consists of alternate layers of selfattention (SA) and MLP.
Revised: Canonical transformer consists of alternate layers of selfattention (SA) and MLP.

Original: In order to promote efficiency, transformer tends to employ local window based attention and shifted window strategy is utilized to allow inter-window connections.
Revised: In order to promote efficiency, transformer tends to employ local window based attention and shifted window strategy is utilized to allow inter-window connections.

Original: Specifically, the feature map is partitioned into non-overlapped windows and then SA is computed within local windows.
Revised: Specifically, the feature map is partitioned into non-overlapped windows and then SA is computed within local windows.

Original: Suppose the window partition is denoted by Par( · ; s, ξ h , ξ w ) , where size of the local window is s and shift of the whole window partition is ( ξ h , ξ w ) (see Figure 2).
Revised: Suppose the window partition is denoted by Par( · ; s, ξ h , ξ w ) , where size of the local window is s and shift of the whole window partition is ( ξ h , ξ w ) (see Figure 2).

Original: With these notations, the shifted window transformer can be reformulated as:
Revised: With these notations, the shifted window transformer can be reformulated as:

Original: [Equation]
Revised: [Equation]

Original: where z l and x l denote the output feature of SA and MLP for layer l , respectively.
Revised: where z l and x l denote the output feature of SA and MLP for layer l , respectively.

Original: To keep notation uncluttered, we overlook layer normalization in Equation (1).
Revised: To keep notation uncluttered, we overlook layer normalization in Equation (1).

Original: We figure out that all the window partitions are equally informative and preserving only one configuration each layer will lead to the broken translation invariance and tremendous loss of locality.
Revised: We figure out that all the window partitions are equally informative and preserving only one configuration each layer will lead to the broken translation invariance and tremendous loss of locality.

Original: The thought of informative equality among all the window partitions leads us to the stochastic window transformer, which is computed as:
Revised: The thought of informative equality among all the window partitions leads us to the stochastic window transformer, which is computed as:

Original: [Equation]
Revised: [Equation]

Original: where the set ℜ s contains all possible shifts and U ( · ) denotes the uniform distribution.
Revised: where the set ℜ s contains all possible shifts and U ( · ) denotes the uniform distribution.

Original: Given the periodicity of the window partition, ℜ s can be simplified as
Revised: Given the periodicity of the window partition, ℜ s can be simplified as

Original: [Equation]
Revised: [Equation]

Original: where × means the Cartesian product.
Revised: where × means the Cartesian product.

Original: During training, ( ξ lh , ξ lw ) are treated as i.i.d.
Revised: During training, ( ξ lh , ξ lw ) are treated as i.i.d.

Original: random variables and sampled from the uniform distribution U ( ℜ s ) .
Revised: random variables and sampled from the uniform distribution U ( ℜ s ) .

Original: Suppose the total number of SA layer is N , stochastic shifts from different layers { ( ξ lh , ξ lw ) } N − 1 l =0 are deliberately designed to be independent so that faithful locality and translation invariance can be ensured on layer level.
Revised: Suppose the total number of SA layer is N , stochastic shifts from different layers { ( ξ lh , ξ lw ) } N − 1 l =0 are deliberately designed to be independent so that faithful locality and translation invariance can be ensured on layer level.

Original: By this treatment, despite all possible shifts are taken into account for each individual layer, single forward propagation only requires a set of sampled shifts { ( ξ lh , ξ lw ) } N − 1 l =0 so that the stochastic window transformer can be trained efficiently.
Revised: By this treatment, despite all possible shifts are taken into account for each individual layer, single forward propagation only requires a set of sampled shifts { ( ξ lh , ξ lw ) } N − 1 l =0 so that the stochastic window transformer can be trained efficiently.

Original: For testing, from Bayesian perspective, stochastic shifts should be averaged according to their posterior distribution.
Revised: For testing, from Bayesian perspective, stochastic shifts should be averaged according to their posterior distribution.

Original: Hence, the exact inference procedure is:
Revised: Hence, the exact inference procedure is:

Original: [Equation]
Revised: [Equation]

Original: where F ( · ) denotes the function of the overall transformer and U is the uniform distribution.
Revised: where F ( · ) denotes the function of the overall transformer and U is the uniform distribution.

Original: The derivation of Equation (4) follows { ( ξ lh , ξ lw ) } N − 1 l =0 are independent.
Revised: The derivation of Equation (4) follows { ( ξ lh , ξ lw ) } N − 1 l =0 are independent.

Original: Obviously, according to Equation (4), the single exact inference requires forward propagation exponential times so that the computational overhead of exact inference grows exponentially with the depth N , which is prohibitively expensive.
Revised: Obviously, according to Equation (4), the single exact inference requires forward propagation exponential times so that the computational overhead of exact inference grows exponentially with the depth N , which is prohibitively expensive.

Original: Consequently, we turn to seek an approximate inference, which should also guarantee the translation invariance and intact locality, to replace the costly exact inference.
Revised: Consequently, we turn to seek an approximate inference, which should also guarantee the translation invariance and intact locality, to replace the costly exact inference.

Original: Inspired by the approximation in Dropout [46], we propose the general layer expectation propagation algorithm to approach the original exact inference process.
Revised: Inspired by the approximation in Dropout [46], we propose the general layer expectation propagation algorithm to approach the original exact inference process.

Original: The layer expectation propagation is formulated as:
Revised: The layer expectation propagation is formulated as:

Original: [Equation]
Revised: [Equation]

Original: I npu t
Revised: I npu t

Original: E m b e dd i n g
Revised: E m b e dd i n g

Original: S t o B l o c k
Revised: S t o B l o c k

Original: D o w n
Revised: D o w n

Original: S a m p l e
Revised: S a m p l e

Original: S t o B l o c k
Revised: S t o B l o c k

Original: D o w n
Revised: D o w n

Original: S a m p l e
Revised: S a m p l e

Original: S t o B l o c k
Revised: S t o B l o c k

Original: D o w n S a m p l e
Revised: D o w n S a m p l e

Original: S t o B l o c k S t o B l o c k
Revised: S t o B l o c k S t o B l o c k

Original: U p S a m p l e
Revised: U p S a m p l e

Original: S t o B l o c k
Revised: S t o B l o c k

Original: U p S a m p l e
Revised: U p S a m p l e

Original: S t o B l o c k
Revised: S t o B l o c k

Original: U p S a m p l e
Revised: U p S a m p l e

Original: R e c o n s t r u c t
Revised: R e c o n s t r u c t

Original: i o n
Revised: i o n

Original: Encoder
Revised: Encoder

Original: Decoder
Revised: Decoder

Original: Layer Norm Layer Norm
Revised: Layer Norm Layer Norm

Original: StoWin-MSA
Revised: StoWin-MSA

Original: MLP ! "
Revised: MLP ! "

Original: #$ % " ! "
Revised: #$ % " ! "

Original: (a) Architecture (b) StoBlock &×(×) &2×(2×2) &2 +#$ × (2 +#$ ×2 +#$ ) &2 + ×(2 + ×2 + ) &2×(2×2) &×(×)
Revised: (a) Architecture (b) StoBlock &×(×) &2×(2×2) &2 +#$ × (2 +#$ ×2 +#$ ) &2 + ×(2 + ×2 + ) &2×(2×2) &×(×)

Original: LevelLevel-
Revised: LevelLevel-

Original: Level-k
Revised: Level-k

Original: Level-k
Revised: Level-k

Original: LevelLevel- &2 +#$ × (2 +#$ ×2 +#$ )
Revised: LevelLevel- &2 +#$ × (2 +#$ ×2 +#$ )

Original: According to Equation (5), the single approximate inference merely requires one forward propagation of the expected signal, which accelerates inference considerably.
Revised: According to Equation (5), the single approximate inference merely requires one forward propagation of the expected signal, which accelerates inference considerably.

Original: We then examine whether the stochastic window strategy is able to keep the translation invariance and model intact local relationships.
Revised: We then examine whether the stochastic window strategy is able to keep the translation invariance and model intact local relationships.

Original: Translation Invariance.
Revised: Translation Invariance.

Original: During training, the stochastic window strategy treats arbitrarily shifted window equally and eliminates the particularity of certain windows introduced by the fixed window partition (blue windows in Figure 1(b)).
Revised: During training, the stochastic window strategy treats arbitrarily shifted window equally and eliminates the particularity of certain windows introduced by the fixed window partition (blue windows in Figure 1(b)).

Original: Therefore, the stochastic window strategy can maintain the translation invariance during training.
Revised: Therefore, the stochastic window strategy can maintain the translation invariance during training.

Original: For testing, by the layer expectation propagation, each individual token is able to aggregate information from its neighbor space with shared weights like
Revised: For testing, by the layer expectation propagation, each individual token is able to aggregate information from its neighbor space with shared weights like

Original: CNNs.
Revised: CNNs.

Original: Consequently, the translation invariance can also be well maintained.
Revised: Consequently, the translation invariance can also be well maintained.

Original: We also provide experimental evidence to support the translation invariance in Section 4.2.
Revised: We also provide experimental evidence to support the translation invariance in Section 4.2.

Original: Locality.
Revised: Locality.

Original: We have argued that the fixed window partition will lead to huge loss of local relationships since the majority of equally informative window partitions are simply discarded.
Revised: We have argued that the fixed window partition will lead to huge loss of local relationships since the majority of equally informative window partitions are simply discarded.

Original: As shown in Figure 1(b), there exists some token pairs whose distance is small enough but they are not present in the same local window, which results in tremendous loss of locality.
Revised: As shown in Figure 1(b), there exists some token pairs whose distance is small enough but they are not present in the same local window, which results in tremendous loss of locality.

Original: In contrast, with the stochastic window strategy, the window partition is randomly shifted, which will guarantee that any token pair will attend in the same window as long as their distance is smaller than the window size s .
Revised: In contrast, with the stochastic window strategy, the window partition is randomly shifted, which will guarantee that any token pair will attend in the same window as long as their distance is smaller than the window size s .

Original: As shown in Figure 1(c), the stochastic window strategy can capture intact local relationships when the radius of neighbor space is not larger than window size while the fixed window partition suffers from severe loss of local relationships.
Revised: As shown in Figure 1(c), the stochastic window strategy can capture intact local relationships when the radius of neighbor space is not larger than window size while the fixed window partition suffers from severe loss of local relationships.

Original: For testing, the proposed layer expectation propagation also considers all the shifted window partitions so that intact locality can be ensured as well.
Revised: For testing, the proposed layer expectation propagation also considers all the shifted window partitions so that intact locality can be ensured as well.

Original: Implicit Model Ensemble.
Revised: Implicit Model Ensemble.

Original: Compared with the typical fixed window strategy (in particular, the shifted window strategy), we can observe significant performance improvements from extensive experiments.
Revised: Compared with the typical fixed window strategy (in particular, the shifted window strategy), we can observe significant performance improvements from extensive experiments.

Original: In addition to the explanation from the compensation of the lost local information and the broken translation invariance, another explanation for performance improvements is that training with the stochastic window can be seen as training an ensemble of the fixed window transformer implicitly [18].
Revised: In addition to the explanation from the compensation of the lost local information and the broken translation invariance, another explanation for performance improvements is that training with the stochastic window can be seen as training an ensemble of the fixed window transformer implicitly [18].

Original: Each self-attention layer contains w 2 window partition, which results in w 2 N possible network combinations.
Revised: Each self-attention layer contains w 2 window partition, which results in w 2 N possible network combinations.

Original: For each training mini-batch, one of the w 2 N networks is sampled and then updated.
Revised: For each training mini-batch, one of the w 2 N networks is sampled and then updated.

Original: For testing phase, all the networks are implicitly integrated using the proposed layer expectation propagation.
Revised: For testing phase, all the networks are implicitly integrated using the proposed layer expectation propagation.

Original: 3.2 Network Architecture
Revised: 3.2 Network Architecture

Original: The proposed stochastic window transformer provides an effective and efficient way to integrate the well-acknowledged priors, i.e., locality and translation invariance, into transformer.
Revised: The proposed stochastic window transformer provides an effective and efficient way to integrate the well-acknowledged priors, i.e., locality and translation invariance, into transformer.

Original: For image restoration tasks which require pixel-level regression, faithfully exploiting these priors is especially important.
Revised: For image restoration tasks which require pixel-level regression, faithfully exploiting these priors is especially important.

Original: To validate the effectiveness of the proposed stochastic window strategy, we integrate it with the widely-used U-shaped architecture [44, 21].
Revised: To validate the effectiveness of the proposed stochastic window strategy, we integrate it with the widely-used U-shaped architecture [44, 21].

Original: The resulting transformer, named Stoformer †† , not only enjoys the intrinsically strong representation ability of Transformer but also maintains the ideal translation invariance and intact local relationships.
Revised: The resulting transformer, named Stoformer †† , not only enjoys the intrinsically strong representation ability of Transformer but also maintains the ideal translation invariance and intact local relationships.

Original: Figure 3 illustrates the overall architecture.
Revised: Figure 3 illustrates the overall architecture.

Original: [Figure]
Revised: [Figure]

Original: 3.3 Implementation Details
Revised: 3.3 Implementation Details

Original: Loss Function.
Revised: Loss Function.

Original: The loss function adopted for training is the Charbonnier loss [1], whose mathematical expression is:
Revised: The loss function adopted for training is the Charbonnier loss [1], whose mathematical expression is:

Original: [Equation]
Revised: [Equation]

Original: where I ′ and I are the restored and ground-truth image respectively.
Revised: where I ′ and I are the restored and ground-truth image respectively.

Original: The constant ϵ is empirically set to 10 − 3 .
Revised: The constant ϵ is empirically set to 10 − 3 .

Original: Training Detail.
Revised: Training Detail.

Original: Stoformer †† employs a four-level encoder-decoder structure.
Revised: Stoformer †† employs a four-level encoder-decoder structure.

Original: The numbers of
Revised: The numbers of

Original: StoBlock are { 1 , 2 , 8 , 8 } for level- 1 to level- 4 of Encoder and the blocks for Decoder are mirrored.
Revised: StoBlock are { 1 , 2 , 8 , 8 } for level- 1 to level- 4 of Encoder and the blocks for Decoder are mirrored.

Original: The number of channel is set to 32 and the window size is 8 × 8 .
Revised: The number of channel is set to 32 and the window size is 8 × 8 .

Original: We train the network with Adam optimizer ( β 1 = 0 .
Revised: We train the network with Adam optimizer ( β 1 = 0 .

Original: 9 , β 2 = 0 .
Revised: 9 , β 2 = 0 .

Original: with the initial learning rate 3 × 10 − 4 gradually reduced to 1 × 10 − 6 with the cosine annealing.
Revised: with the initial learning rate 3 × 10 − 4 gradually reduced to 1 × 10 − 6 with the cosine annealing.

Original: The training samples are augmented by the horizontal flipping and rotation of 90 ◦ , 180 ◦ , or 270 ◦ .
Revised: The training samples are augmented by the horizontal flipping and rotation of 90 ◦ , 180 ◦ , or 270 ◦ .

Original: Please refer to the supplemental material for task-specific settings.
Revised: Please refer to the supplemental material for task-specific settings.

Original:  Experiments 4.1 Experimental Setup
Revised:  Experiments 4.1 Experimental Setup

Original: In this section, we validate the effectiveness of the stochastic window strategy.
Revised: In this section, we validate the effectiveness of the stochastic window strategy.

Original: Except the previous state-of-the-art methods are included for comparisons, we also elaborate four Stoformer variants:
Revised: Except the previous state-of-the-art methods are included for comparisons, we also elaborate four Stoformer variants:

Original: Stoformer ♢♢ , Stoformer † ♢ , Stoformer ♢ † , Stoformer †† .
Revised: Stoformer ♢♢ , Stoformer † ♢ , Stoformer ♢ † , Stoformer †† .

Original: The first symbol aims to indicate whether the stochastic window is adopted for training( † : the stochastic window; ♢ : the fixed window) and the second symbol represents whether the layer expectation propagation is adopted for testing( † : layer expectation propagation; ♢ : the fixed window).
Revised: The first symbol aims to indicate whether the stochastic window is adopted for training( † : the stochastic window; ♢ : the fixed window) and the second symbol represents whether the layer expectation propagation is adopted for testing( † : layer expectation propagation; ♢ : the fixed window).

Original: In all experiments, we use the shifted window to specify the fixed window (Equation (1)).
Revised: In all experiments, we use the shifted window to specify the fixed window (Equation (1)).

Original: These variants are identical except for the aforementioned window strategy.
Revised: These variants are identical except for the aforementioned window strategy.

Original: In particular, Stoformer ♢♢ is reduced to the traditional fixed window transformer and Stoformer †† is the proposed stochastic window transformer.
Revised: In particular, Stoformer ♢♢ is reduced to the traditional fixed window transformer and Stoformer †† is the proposed stochastic window transformer.

Original: 4.2 Toy Experiment
Revised: 4.2 Toy Experiment

Original: To illustrate the broken translation invariance of the fixed window strategy and the remedy of our proposed stochastic window strategy to this effect, a toy experiment is performed: we add a Gaussian noise patch with σ = 50 whose spatial size is 64 × 64 to an 8 -bit pure color image with spatial size 256 × 256 , the pixel value of which is constantly set to the medium value 127 .
Revised: To illustrate the broken translation invariance of the fixed window strategy and the remedy of our proposed stochastic window strategy to this effect, a toy experiment is performed: we add a Gaussian noise patch with σ = 50 whose spatial size is 64 × 64 to an 8 -bit pure color image with spatial size 256 × 256 , the pixel value of which is constantly set to the medium value 127 .

Original: We move the
Revised: We move the

Original: Gaussian noise patch with stride 1 and obtain total 192 × 192 = 36864 noisy images.
Revised: Gaussian noise patch with stride 1 and obtain total 192 × 192 = 36864 noisy images.

Original: Then we denoise these noisy images using the fixed and stochastic window transformer, respectively and the PSNR distributions with respect to ( ξ h , ξ w ) are normalized to the fixed range [0 , 1] and plotted in Figure 4.
Revised: Then we denoise these noisy images using the fixed and stochastic window transformer, respectively and the PSNR distributions with respect to ( ξ h , ξ w ) are normalized to the fixed range [0 , 1] and plotted in Figure 4.

Original: To further highlight the translation invariance of stochastic window transformer, we also include a typical CNNs-based denoiser VDN [60], which can ensure the translation invariance well due to inherent structure of CNNs, as the inference.
Revised: To further highlight the translation invariance of stochastic window transformer, we also include a typical CNNs-based denoiser VDN [60], which can ensure the translation invariance well due to inherent structure of CNNs, as the inference.

Original: Figure 4 indicates clearly that there are great variations, especially for the interior, in the PSNR distribution of fixed window transformer while stochastic window transformer flattens the distribution significantly.
Revised: Figure 4 indicates clearly that there are great variations, especially for the interior, in the PSNR distribution of fixed window transformer while stochastic window transformer flattens the distribution significantly.

Original: Both stochastic window transformer and VDN contain the relatively uniform distribution for the interior, which indicates that the stochastic window strategy helps to maintain the desired translation invariance quite well.
Revised: Both stochastic window transformer and VDN contain the relatively uniform distribution for the interior, which indicates that the stochastic window strategy helps to maintain the desired translation invariance quite well.

Original: It is worth noting that all these distributions have variations on the image boundary, which can be attributed to the difference of the contexts of the image boundary.
Revised: It is worth noting that all these distributions have variations on the image boundary, which can be attributed to the difference of the contexts of the image boundary.

Original: 4.3 Experiments on Image Deraining
Revised: 4.3 Experiments on Image Deraining

Original: [Table] [Table]
Revised: [Table] [Table]

Original: We validate the stochastic window strategy on image deraining task.
Revised: We validate the stochastic window strategy on image deraining task.

Original: Specifically, except the elaborate four Stoformer variants, existing seven deraining methods are included: GMM [31], DDN [13],
Revised: Specifically, except the elaborate four Stoformer variants, existing seven deraining methods are included: GMM [31], DDN [13],

Original: SPANet [51], JORDER-E [58], RCDNet [50], SPAIR [40], Uformer [54].
Revised: SPANet [51], JORDER-E [58], RCDNet [50], SPAIR [40], Uformer [54].

Original: All these methods are evaluated on SPA-Data [51] and the performance comparison is shown in Table 1.
Revised: All these methods are evaluated on SPA-Data [51] and the performance comparison is shown in Table 1.

Original: Due to the limited space, visualization of image deraining is included in the supplement material.
Revised: Due to the limited space, visualization of image deraining is included in the supplement material.

Original: 4.4 Experiments on Image Denoising
Revised: 4.4 Experiments on Image Denoising

Original: We conduct denoising experiments on the additive white Gaussian Noise benchmark datasets, which include Set 12 [66], BSD 68 [37], Urban 100 [19], Kodak 24 [12] and McMaster [71].
Revised: We conduct denoising experiments on the additive white Gaussian Noise benchmark datasets, which include Set 12 [66], BSD 68 [37], Urban 100 [19], Kodak 24 [12] and McMaster [71].

Original: Following previous works [66, 60, 67], a single model is trained to tackle with various noise levels.
Revised: Following previous works [66, 60, 67], a single model is trained to tackle with various noise levels.

Original: Tablesand 4 report PSNR scores of existing methods as well as Stoformer variants for color and grayscale image denoising, respectively.
Revised: Tablesand 4 report PSNR scores of existing methods as well as Stoformer variants for color and grayscale image denoising, respectively.

Original: The supplemental material contains visualization of image denoising.
Revised: The supplemental material contains visualization of image denoising.

Original: 4.5 Experiments on Image Deblurring
Revised: 4.5 Experiments on Image Deblurring

Original: We also perform deblurring experiments on the benchmark datasets (GoPro [39] and HIDE [45]).
Revised: We also perform deblurring experiments on the benchmark datasets (GoPro [39] and HIDE [45]).

Original: The model is trained only on the GoPro dataset and directly evaluated on the HIDE.
Revised: The model is trained only on the GoPro dataset and directly evaluated on the HIDE.

Original: Table 2 presents
Revised: Table 2 presents

Original: PSNR and SSIM scores of different deblurring methods and Stoformers.
Revised: PSNR and SSIM scores of different deblurring methods and Stoformers.

Original: Figure 5 presents an image deblurring example from GoPro [39].
Revised: Figure 5 presents an image deblurring example from GoPro [39].

Original: More visual results are provided in the supplemental material.
Revised: More visual results are provided in the supplemental material.

Original: Remark.
Revised: Remark.

Original: From above experiments on various image restoration tasks, we can make the following observations and analyses:
Revised: From above experiments on various image restoration tasks, we can make the following observations and analyses:

Original: [Table] [Table]
Revised: [Table] [Table]

Original: • Stoformer ♢♢ vs. Stoformer ♢ † : Without the stochastic window strategy for training, directly applying the layer expectation propagation will degrade performance dramatically, which reveals that the performance gain cannot simply attribute to feature ensemble at test time and the stochastic window for training matters.
Revised: • Stoformer ♢♢ vs. Stoformer ♢ † : Without the stochastic window strategy for training, directly applying the layer expectation propagation will degrade performance dramatically, which reveals that the performance gain cannot simply attribute to feature ensemble at test time and the stochastic window for training matters.

Original: • Stoformer ♢♢ vs.
Revised: • Stoformer ♢♢ vs.

Original: Stoformer † ♢ :
Revised: Stoformer † ♢ :

Original: Even without the layer expectation propagation algorithm for testing, the stochastic window partition for training alone is also conductive to boost performance.
Revised: Even without the layer expectation propagation algorithm for testing, the stochastic window partition for training alone is also conductive to boost performance.

Original: This is reasonable since a large amount of local information lost by Stoformer ♢♢ is re-exploited by Stoformer † ♢ .
Revised: This is reasonable since a large amount of local information lost by Stoformer ♢♢ is re-exploited by Stoformer † ♢ .

Original: • Stoformer †† vs. Others: Equipped with the stochastic window for training and layer expectation propagation for testing (Stoformer †† ), the model achieves the highest performance.
Revised: • Stoformer †† vs. Others: Equipped with the stochastic window for training and layer expectation propagation for testing (Stoformer †† ), the model achieves the highest performance.

Original: 4.6 Analytic Experiments
Revised: 4.6 Analytic Experiments

Original: Complexity Analysis.
Revised: Complexity Analysis.

Original: We provide a detailed analysis about the time and space complexity of the stochastic window strategy based on a single attention layer.
Revised: We provide a detailed analysis about the time and space complexity of the stochastic window strategy based on a single attention layer.

Original: Specifically, we compare our stochastic window strategy with fixed and sliding window strategy for training and testing, respectively.
Revised: Specifically, we compare our stochastic window strategy with fixed and sliding window strategy for training and testing, respectively.

Original: For training, our strategy is as efficient as the fixed window for both speed and memory cost, as shown in Table 5.
Revised: For training, our strategy is as efficient as the fixed window for both speed and memory cost, as shown in Table 5.

Original: But the fixed window loses translation invariance and locality.
Revised: But the fixed window loses translation invariance and locality.

Original: The sliding window can fulfill the translation invariance and locality while suffer from huge memory burden, which often incurs OOM.
Revised: The sliding window can fulfill the translation invariance and locality while suffer from huge memory burden, which often incurs OOM.

Original: In contrast, the stochastic window enjoys efficient training while maintains translation invariance and locality.
Revised: In contrast, the stochastic window enjoys efficient training while maintains translation invariance and locality.

Original: For testing, as shown in Table 6, the fixed window can inference efficiently but suffer from the broken translation invariance and locality loss.
Revised: For testing, as shown in Table 6, the fixed window can inference efficiently but suffer from the broken translation invariance and locality loss.

Original: Compared with the sliding window, our stochastic window strategy requires more computations but less memory cost
Revised: Compared with the sliding window, our stochastic window strategy requires more computations but less memory cost

Original: [Table] [Table]
Revised: [Table] [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Revised: [Equation]

Original:  [Equation] [Table]
Revised: [Table] [Equation] [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table][Equation] [Table]
Revised: [Table]

Original: [Table]
Revised: [Table]

Original: Eliminating Blocking Artifacts.
Revised: Eliminating Blocking Artifacts.

Original: Since the fixed window strategy implies infinite preference towards specific window partition, the learned feature map will consequently contain annoying blocking artifacts.
Revised: Since the fixed window strategy implies infinite preference towards specific window partition, the learned feature map will consequently contain annoying blocking artifacts.

Original: In contrast, the proposed stochastic window strategy treats all window partitions equally, facilitating to keep translation invariance and model locality.
Revised: In contrast, the proposed stochastic window strategy treats all window partitions equally, facilitating to keep translation invariance and model locality.

Original: Therefore, the blocking artifacts should be significantly eliminated by the stochastic window strategy.
Revised: Therefore, the blocking artifacts should be significantly eliminated by the stochastic window strategy.

Original: Figure 6 provides experimental evidence to support our inference.
Revised: Figure 6 provides experimental evidence to support our inference.

Original: Specifically, the learned feature map by the fixed window strategy (Figure 6(b)) contains obvious blocking artifacts while the stochastic window strategy (Figure 6(c)) can remove these side effects as expected.
Revised: Specifically, the learned feature map by the fixed window strategy (Figure 6(b)) contains obvious blocking artifacts while the stochastic window strategy (Figure 6(c)) can remove these side effects as expected.

Original: It is surprising that with the stochastic window strategy, the feature is divided into non-overlapping local windows for further processing but in turn nearly no blocking artifacts are preserved.
Revised: It is surprising that with the stochastic window strategy, the feature is divided into non-overlapping local windows for further processing but in turn nearly no blocking artifacts are preserved.

Original: Please refer to the supplemental material for more visualization.
Revised: Please refer to the supplemental material for more visualization.

Original: [Figure]
Revised: [Figure]

Original: Trade-off in Layer Expectation Propagation.
Revised: Trade-off in Layer Expectation Propagation.

Original: Although the layer expectation propagation accelerates inference significantly, the strict conduction still requires to take average of s × s features yielded by self-attention.
Revised: Although the layer expectation propagation accelerates inference significantly, the strict conduction still requires to take average of s × s features yielded by self-attention.

Original: Here, we investigate the influence of the number of averaged features to model performance.
Revised: Here, we investigate the influence of the number of averaged features to model performance.

Original: Figure 7 shows the trend of model performance and FLOPs of testing 3 on a 256 × 256 image with respect to the number of averaged features each layer.
Revised: Figure 7 shows the trend of model performance and FLOPs of testing 3 on a 256 × 256 image with respect to the number of averaged features each layer.

Original: Figure 7 reveals that
Revised: Figure 7 reveals that

Original: PSNR first rises rapidly and then tends to saturate with the increasing of the averaged feature number while FLOPs increases linearly with that number, which provides a promising trade-off between performance and FLOPs.
Revised: PSNR first rises rapidly and then tends to saturate with the increasing of the averaged feature number while FLOPs increases linearly with that number, which provides a promising trade-off between performance and FLOPs.

Original: For instance, we can decrease the averaged feature number from 64 to 8 to keep relatively high performance and efficient inference as well.
Revised: For instance, we can decrease the averaged feature number from 64 to 8 to keep relatively high performance and efficient inference as well.

Original: Boosting Model Generalization.
Revised: Boosting Model Generalization.

Original: Figure 8 presents the loss curve on SPA-Data during training and the PSNR curve on the test set is plotted in Figure 9.
Revised: Figure 8 presents the loss curve on SPA-Data during training and the PSNR curve on the test set is plotted in Figure 9.

Original: Compared with the fixed window strategy, our proposed stochastic window strategy attains higher PSNR on the test set with similar training error, which demonstrates that the stochastic window strategy boosts model generalization.
Revised: Compared with the fixed window strategy, our proposed stochastic window strategy attains higher PSNR on the test set with similar training error, which demonstrates that the stochastic window strategy boosts model generalization.

Original: This can be anticipated since the stochastic window strategy can exploit complete local relationships, which are discarded by the fixed window strategy.
Revised: This can be anticipated since the stochastic window strategy can exploit complete local relationships, which are discarded by the fixed window strategy.

Original:  Limitation and Discussion
Revised:  Limitation and Discussion

Original: In this work, we propose the stochastic window strategy and extensive experiments validate the effectiveness on several image restoration tasks, including image deraining, denoising and deblurring.
Revised: In this work, we propose the stochastic window strategy and extensive experiments validate the effectiveness on several image restoration tasks, including image deraining, denoising and deblurring.

Original: Our experiments are mainly based on the widely-used U-shaped architecture.
Revised: Our experiments are mainly based on the widely-used U-shaped architecture.

Original: We will further validate the stochastic strategy on more architectures (e.g., isotropic and multi-stage architecture).
Revised: We will further validate the stochastic strategy on more architectures (e.g., isotropic and multi-stage architecture).

Original: Indeed, the performing improvements can be positively anticipated due to the desired translation invariance and locality.
Revised: Indeed, the performing improvements can be positively anticipated due to the desired translation invariance and locality.

Original: Since the importance of translation invariance and locality is quite general, we also plan to extend the stochastic window strategy based transformer to more low-level vision tasks, such as image super-resolution, JPEG deblocking, and video restoration.
Revised: Since the importance of translation invariance and locality is quite general, we also plan to extend the stochastic window strategy based transformer to more low-level vision tasks, such as image super-resolution, JPEG deblocking, and video restoration.

Original: Local attention based transformers, e.g., Swin Transformer [33], have achieved remarkable success on several CV tasks, including image classification, dense prediction, and semantic segmentation.
Revised: Local attention based transformers, e.g., Swin Transformer [33], have achieved remarkable success on several CV tasks, including image classification, dense prediction, and semantic segmentation.

Original: However, they also compute attention based on specific local windows and cannot treat all the local windows fairly.
Revised: However, they also compute attention based on specific local windows and cannot treat all the local windows fairly.

Original: The consequences are that unexpectable information loss occurs during feature processing and the translation invariance is also absent.
Revised: The consequences are that unexpectable information loss occurs during feature processing and the translation invariance is also absent.

Original: Our proposed stochastic window strategy can provide an effective mechanism towards this problem.
Revised: Our proposed stochastic window strategy can provide an effective mechanism towards this problem.

Original: Therefore, it is also promising to extend the strategy to other tasks where local attention based transformers have exhibited competitive performance.
Revised: Therefore, it is also promising to extend the strategy to other tasks where local attention based transformers have exhibited competitive performance.

Original:  Conclusion
Revised:  Conclusion

Original: In this paper, we analyze the deficiencies of existing transformers for image restoration.
Revised: In this paper, we analyze the deficiencies of existing transformers for image restoration.

Original: We figure out that the fixed window strategy will inevitably lead to the broken translation invariance and loss of locality.
Revised: We figure out that the fixed window strategy will inevitably lead to the broken translation invariance and loss of locality.

Original: To tackle with these issues, the novel stochastic window strategy is proposed.
Revised: To tackle with these issues, the novel stochastic window strategy is proposed.

Original: Specifically, we utilize the window partition with stochastic shift to replace the fixed window partition for training, which can ensure the translation invariance and intact locality.
Revised: Specifically, we utilize the window partition with stochastic shift to replace the fixed window partition for training, which can ensure the translation invariance and intact locality.

Original: For testing, we propose the efficient layer expectation propagation to approximately take expectation of the introduced stochastic shift.
Revised: For testing, we propose the efficient layer expectation propagation to approximately take expectation of the introduced stochastic shift.

Original: Based on this strategy, we conduct extensive comparison experiments on various image restoration tasks to validate the effectiveness of the proposed stochastic window strategy.
Revised: Based on this strategy, we conduct extensive comparison experiments on various image restoration tasks to validate the effectiveness of the proposed stochastic window strategy.

Original: Broad Impact
Revised: Broad Impact

Original: Nowadays, image acquisition system inevitably suffers from various degradation, ranging from inherent noise of capturing instruments, the shaking during shooting, to unpredictable weather condition.
Revised: Nowadays, image acquisition system inevitably suffers from various degradation, ranging from inherent noise of capturing instruments, the shaking during shooting, to unpredictable weather condition.

Original: Hence, image restoration itself has important value of research and application.
Revised: Hence, image restoration itself has important value of research and application.

Original: Our proposed stochastic window strategy can make transformer more sophisticated in restoring degraded images.
Revised: Our proposed stochastic window strategy can make transformer more sophisticated in restoring degraded images.

Original: However, from a societal point of view, negative consequences may also come along.
Revised: However, from a societal point of view, negative consequences may also come along.

Original: For example, the deviation from the actual image textures caused by image restoration technology may affect fair judgment in medical and criminal situations.
Revised: For example, the deviation from the actual image textures caused by image restoration technology may affect fair judgment in medical and criminal situations.

Original: In these cases, it is necessary to combine expert knowledge to make rational decisions.
Revised: In these cases, it is necessary to combine expert knowledge to make rational decisions.

Original: Acknowledgement
Revised: Acknowledgement

Original: This work was supported by the National Key R&D Program of China under Grant 2020AAA0105701, the National Natural Science Foundation of China (NSFC) under Grants 62225207 and U19B2038, the University Synergy Innovation Program of Anhui Province under Grant GXXT-2019- and the USTC Research Funds of the Double First-Class Initiative under Grant YD
Revised: This work was supported by the National Key R&D Program of China under Grant 2020AAA0105701, the National Natural Science Foundation of China (NSFC) under Grants 62225207 and U19B2038, the University Synergy Innovation Program of Anhui Province under Grant GXXT-2019- and the USTC Research Funds of the Double First-Class Initiative under Grant YD

Original: Checklist
Revised: Checklist

Original: For all authors... (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope?
Revised: For all authors... (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope?

Original: [Yes] (b) Did you describe the limitations of your work?
Revised: [Yes] (b) Did you describe the limitations of your work?

Original: [Yes] See Sec.
Revised: [Yes] See Sec.

Original: (c) Did you discuss any potential negative societal impacts of your work?
Revised: (c) Did you discuss any potential negative societal impacts of your work?

Original: [Yes] See Sec. 6.
Revised: [Yes] See Sec. 6.

Original: (d) Have you read the ethics review guidelines and ensured that your paper conforms to them?
Revised: (d) Have you read the ethics review guidelines and ensured that your paper conforms to them?

Original: [Yes] 2.
Revised: [Yes] 2.

Original: If you are including theoretical results... (a) Did you state the full set of assumptions of all theoretical results?
Revised: If you are including theoretical results... (a) Did you state the full set of assumptions of all theoretical results?

Original: [Yes] (b) Did you include complete proofs of all theoretical results?
Revised: [Yes] (b) Did you include complete proofs of all theoretical results?

Original: [Yes] 3.
Revised: [Yes] 3.

Original: If you ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
Revised: If you ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?

Original: [Yes] (b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?
Revised: [Yes] (b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?

Original: [Yes] See supplemental material.
Revised: [Yes] See supplemental material.

Original: (c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
Revised: (c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?

Original: [No] The random seed is fixed for reproducibility.
Revised: [No] The random seed is fixed for reproducibility.

Original: (d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)?
Revised: (d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)?

Original: [Yes] See supplemental material.
Revised: [Yes] See supplemental material.

Original: If you are using existing assets (e.g., code, data, models) or curating/releasing new assets... (a) If your work uses existing assets, did you cite the creators?
Revised: If you are using existing assets (e.g., code, data, models) or curating/releasing new assets... (a) If your work uses existing assets, did you cite the creators?

Original: [Yes] (b) Did you mention the license of the assets?
Revised: [Yes] (b) Did you mention the license of the assets?

Original: [Yes] (c) Did you include any new assets either in the supplemental material or as a URL?
Revised: [Yes] (c) Did you include any new assets either in the supplemental material or as a URL?

Original: [No] (d) Did you discuss whether
Revised: [No] (d) Did you discuss whether

Original: and how consent was obtained from people whose data you’re using/curating?
Revised: and how consent was obtained from people whose data you’re using/curating?

Original: [N/A]
Revised: [N/A]

Original: The data used for experiments are all publicly available.
Revised: The data used for experiments are all publicly available.

Original: (e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content?
Revised: (e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content?

Original: [N/A]
Revised: [N/A]

Original: The data we used contain no personally identifiable information or offensive content.
Revised: The data we used contain no personally identifiable information or offensive content.

Original: If you used crowdsourcing or conducted research with human subjects... (a) Did you include the full text of instructions given to participants and screenshots, if applicable?
Revised: If you used crowdsourcing or conducted research with human subjects... (a) Did you include the full text of instructions given to participants and screenshots, if applicable?

Original: [N/A] (b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?
Revised: [N/A] (b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?

Original: [N/A] (c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?
Revised: [N/A] (c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?

Original: [N/A]
Revised: [N/A]