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

Original: Anonymous Author(s)

Original: AffiliationAddress email

Original: Abstract
Revised: Abstract

Original: Thanks to the strong representation ability, transformers have attained impressiveresults for image restoration.
Revised: Thanks to the strong representation ability, transformers have attained impressive results for image restoration.

Original: However, existing transformers do not carefully takeinto account the particularities of image restoration.
Revised: However, existing transformers do not carefully take into account the particularities of image restoration.

Original: Basically, image restorationrequires that the ideal approach should be invariant to translation of degradation,i.e., undesirable degradation should be removed irrespective of its position withinthe image.
Revised: Basically, image restoration requires that the ideal approach should be invariant to translation of degradation, i.e., undesirable degradation should be removed irrespective of its position within the image.

Original: Moreover, local relationships play a vital role and should be faithfullyexploited for recovering clean images.
Revised: Moreover, local relationships play a vital role and should be faithfully exploited for recovering clean images.

Original: Nevertheless, most of transformers haveresorted to either fixed local window basedattention or global attention, whichunfortunately breaks the translation invariance and further causes huge loss of localrelationships.
Revised: Nevertheless, most of transformers have resorted to either fixed local window based or global attention, which unfortunately breaks the translation invariance and further causes huge loss of local relationships.

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

Original: Specifically, we introduce the window partition withstochastic shift to replace the original fixed window partition for training andelaborate the layer expectation propagation algorithm to efficiently approximate theexpectation of the induced stochastic transformer for testing.
Revised: Specifically, we introduce the window partition with stochastic shift to replace the original fixed window partition for training and elaborate the layer expectation propagation algorithm to efficiently approximate the expectation of the induced stochastic transformer for testing.

Original: The stochastic windowtransformer can not only enjoy powerful representation but also maintain the desiredproperty of translation invariance and locality.
Revised: The stochastic window transformer can not only enjoy powerful representation but also maintain the desired property of translation invariance and locality.

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

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

Original:  Introduction
Revised:  Introduction

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

Original: In general, image restoration methodsshould satisfy a fundamental constraint: the ideal approach should remove undesirable degradationirrespective of its position within the image.
Revised: In general, image restoration methods should satisfy a fundamental constraint: the ideal approach should remove undesirable degradation irrespective of its position within the image.

Original: In other words, maintaining the translation invariance isthe basic requirement for image restoration methods.
Revised: In other words, maintaining the translation invariance is the basic requirement for image restoration methods.

Original: In this sense, convolutional neural networks(CNNs) are well-suited for various image restoration tasks [52, 21, 9,  69,66, 57, 62, 54, 68, 6, 14].
Revised: In this sense, convolutional neural networks (CNNs) are well-suited for various image restoration tasks [53, 22, 9, 70, 67, 58, 63, 55, 69, 6, 14].

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

Original: These two priors play a vitalrole in image restoration, i.e., pixels within local regions tend to exhibit strong correlations whiletranslation invariance is  the desiderata for any 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 methods [2, 51, 30, 58] for image restoration are not compatiblewith the translation invariance and locality.
Revised: On the other hand, transformer-based methods [2, 52, 31, 59] for image restoration are not compatible with the translation invariance and locality.

Original: For now, the patterns that transformers adopted forrecovering high-resolution images can be concluded as small patch with global attention (e.g.,48 × 48 for IPT [2]) and large patch with local attention (e.g., Uformer [51] and SwinIR [30]).
Revised: For now, 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 [52] and SwinIR [31]).

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

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

Original: [Figure]
Revised: [Figure]

Original: The latter pattern partitions the feature map into fixed non-overlapped windows and self-attention isrestricted inside these local windows.
Revised: The latter pattern 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 breakthe 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 Fig. 1(a),for the fixed window partition strategy, translation invariance can only be ensured among specificpartitioned windows (blue windows) 1 rather than general windows with arbitrary shifts (red windows),which disagrees with the translation invariance.
Revised: As shown in Fig. 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 windowpartition imposes the artificial prior to windows with specific shifts, which breaks the translationinvariance.
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: Fig. 1(b) intuitively illustrates that the fixed window partition cannot capture intact localrelationships.
Revised: Fig. 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 theneighbor 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 acertain token (central red dot), many tokens (green dots) within its neighbor space (encompassedby the red circle with radius r ) cannot participate in the calculation of self-attention, which leads totremendous 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 , theloss will be more severe.
Revised: As the neighbor space expands, i.e., increasing radius r , the loss will be more severe.

Original: Fig.
Revised: Fig.

Original: 1(c) quantitatively depicts the proportion of local relationships thatare captured by the fixed window partition over the total relationships existing in the neighbor space.
Revised: 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 expectedthat 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, wecan observe that there exists significant loss of local relationships even within the neighbor spacewhose 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 toadopt 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 will incur huge overhead in terms of both memory and computing speed.
Revised: However, this strategy will incur 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 thetransformer 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 thefeature map into fixed non-overlapped windows, we choose to cover the whole feature map withstochastically shifted window partition (see Fig. 2).
Revised: Specifically, instead of partitioning the feature map into fixed non-overlapped windows, we choose to cover the whole feature map with stochastically shifted window partition (see Fig. 2).

Original: By introducing stochastic shift, all windows aretreated 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 Fig.
Revised: As shown in Fig.

Original: 1(c),our stochastic window strategy can capture complete local relationships until the radius of neighborspace reaches the window size.
Revised: 1(c), our stochastic window strategy can 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 timeand memory overhead consumed for training stochastic window transformer is comparable to thefixed 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 algorithmtoefficiently approximate the effect of taking expectation with respect to  the stochastic shift, duringwhich the translation invariance and locality can also be maintained.
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: The contributions are threefold.1 .
Revised: In conclusion, the contributions are threefold:

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

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

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

Original:  We propose anew stochastic window strategy, comprising the stochastic window for training and layer expectationpropagation 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, comprising the stochastic window 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:  We conduct extensive experiments to validatethat by equipping our stochastic window strategy, the performance on various image restoration tasks(e.g. , image deraining, denoising and deblurrring) can be constantly improved.
Revised: • We conduct extensive experiments to validate that by equipping our stochastic window strategy, the performance on various image restoration tasks (e.g., image deraining, denoising and deblurrring) can be constantly improved.

Original:  Related works
Revised:  Related works

Original: Image Restoration.
Revised: Image Restoration.

Original: Image Restoration [63, 11, 5, 33, 31, 35, 59, 37, 27, 26] aims to restore theclean image from its degraded version.
Revised: Image Restoration [64, 11, 5, 34, 32, 36, 60, 38, 28, 27] aims to restore the clean image from its degraded version.

Original: In recent years, remarkable progress against traditional modeldriven methods [4, 50, 15] has been achieved due to the development of deep learning technologies[25], especially CNNs [24, 20, 16].
Revised: In recent years, remarkable progress against traditional modeldriven methods [4, 51, 15] has been achieved due to the development of deep learning technologies [26], especially CNNs [25, 21, 16].

Original: Instead of relying on preset image priors, learning based methodsdirectly 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 [23, 56, 64, 62, 7, 70], image super-resolution[68, 6, 14, 18, 71], imagedeblurring [9, 21, 36, 21, 44, 61, 40, 39], image deraining [54, 47, 13, 60, 48, 38, 28], etc.
Revised: Numerous representative CNNs have sprung up across various classical image restoration tasks, including image denoising [24, 57, 65, 63, 7, 71], image super-resolution[69, 6, 14, 19, 72], image deblurring [9, 22, 37, 22, 45, 62, 41, 40], image deraining [55, 48, 13, 61, 49, 39, 29], etc.

Original: In general,the translation invariance (weight sharing) and locality have been hard-coded into the inherentstructure 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 [46, 8] in the NLP field, Vision
Revised: Recently, with great success of transformer [47, 8] in the NLP field, Vision

Original: Transformers [10, 45, 49, 32, 53, 55, 3] have also been prevalent in vision community.
Revised: Transformers [10, 46, 50, 33, 54, 56, 3] have also been prevalent in vision community.

Original: ViT [10] treatedimage 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 [32]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 [33] introduced local window based attention and established a hierarchical architecture.

Original: Inspired bythe key insights from high-level vision, a few transformers [2, 30, 51, 58] for low-level vision havearose.
Revised: Inspired by the key insights from high-level vision, a few transformers [2, 31, 52, 59] for low-level vision have arose.

Original: But most of them directly transfer high-level designs (e.g., local attention) into low-level visionwithout 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, transformersdo 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 usuallyrequire to accomplish pixel-level regression, where more strict locality and translation invarianceare 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 ofmemory 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 translationinvariance 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 forimage 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 shouldhave 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 oflocal 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 Fig.
Revised: As shown in Fig.

Original: 2, for a feature map, the windowpartition with arbitrary shift ( ξ h , ξ w ) contains a comparable quantity oflocal relationships and implies the translation invariance among dividedwindows.
Revised: 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 treatedequally.
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 stochasticwindow 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: Thenetwork can be trained based on local attention but with stochastic ratherthan fixed shift.
Revised: The network can be trained based on local attention but with stochastic rather than fixed shift.

Original: Therefore, transformer with the stochastic window canbe 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 thelayer expectation propagation algorithm to approximate the expectationof the introduced stochastic shift, which also helps to maintain the desiredproperty 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  translation invariance to transformer and make it faithfullyexploit 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 self-attention (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 attentionand 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 featuremap 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 andshift of the whole window partition is ( ξ h , ξ w ) (see Fig. 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 Fig. 2).

Original: With these notations, the shifted windowtransformer can be formulated 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 notationuncluttered, we overlook layer normalization in Eq.
Revised: To keep notation uncluttered, we overlook layer normalization in Eq.

Original: We figure out that all the window partitionsare equally informative and preserving only one configuration each layer will lead to the brokentranslation 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 allthe 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.
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  random variablesand sampled from ℜ (  s )independently and uniformly .
Revised: During training, ( ξ lh , ξ lw ) are treated as i.i.d. 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 independentso 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 propagationrequires only  a set of sampled shifts { ( ξ lh , ξ lw ) } N − 1 l [-=
-] so that the stochastic window transformer canbe 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 averagedaccording 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 P is the uniform distribution.
Revised: where F ( · ) denotes the function of the overall transformer and U is the uniform distribution.

Original: Thederivation of Eq.
Revised: The derivation of Eq.

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

Original: Obviously, according to Eq. 4, thesingle exact inference requires forward propagation exponential times so that the computationaloverhead of exact inference grows exponentially with the depth N , which is prohibitively expensive.
Revised: Obviously, according to Eq. 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 translationinvariance 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 approximationin Dropout [43], we propose the general layer expectation propagation algorithm to approach theoriginal exact inference process.
Revised: Inspired by the approximation in Dropout [44], 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 Eq. 4, the single approximate inference merely requires one forward propagation of theexpected signal, which accelerates inference considerably.
Revised: According to Eq. 5, the single approximate inference merely requires one forward propagation of the expected signal, which accelerates inference considerably.

Original: We then examine whether the stochasticwindow 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 shiftedwindow equally and eliminates the particularity of certain windows introduced by the fixed windowpartition (blue windows in Fig. 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 Fig. 1(b)).

Original: Therefore, the stochastic window strategy can maintainthe 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, eachindividual 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 provideexperimental evidence to support the translation invariance in Sec. 4.2.
Revised: We also provide experimental evidence to support the translation invariance in Sec. 4.2.

Original: Locality.
Revised: Locality.

Original: We have argued that the fixed window partition will lead to huge loss of local relationshipssince 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 Fig.1(b), there exists some token pairs whose distance is small enough but they are not present in thesame local window, which results in tremendous loss of locality.
Revised: As shown in Fig. 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: By contrast, with the stochasticwindow strategy, the window partition is randomly shifted, which will guarantee that any token pairwill 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 shownin Fig.
Revised: As shown in Fig.

Original: 1(c), the stochastic window strategy can capture intact local relationships when the radius ofneighbor space is not larger than window size while the fixed window partition suffers from severeloss of local relationships.
Revised: 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 allthe 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.

Revised: Implicit Model Ensemble.

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

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 [17].

Original:  3.2 Network Architecture
Revised: Each self-attention layer contains w 2 window partition, which results in w 2 N possible network combinations. For each training mini-batch, one of the w 2 N networks is sampled and then updated. For testing phase, all the networks are implicitly integrated using the proposed layer expectation propagation. 3.2 Network Architecture

Original: The proposed stochastic window transformer provides an effective and efficient way to integratethe 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 imagerestoration tasks which require pixel-level regression, faithfully exploiting these priors is especiallyimportant.
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 itwith the widely-used U-shaped architecture [41, 19].
Revised: To validate the effectiveness of the proposed stochastic window strategy, we integrate it with the widely-used U-shaped architecture [42, 20].

Original: The resulting transformer, named Stoformer †† ,not only enjoys the intrinsically strong representation ability of Transformer but also maintains theideal 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: Fig.
Revised: Fig.

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

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 mathematicalexpression 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 setto 10 − 3 .
Revised: The constant ϵ is empirically set to 10 − 3 .

Original: Training details.
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: [Figure]

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 Adamoptimizer ( β 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 to1 × 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 flippingand 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 
Revised:  Experiments 4.1 Experimental Setup

Original: 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 previousstate-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 whetherthe stochastic window is adopted for training( † : the stochastic window; ♢ : the fixed window) andthe 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 windowto specify the fixed window (Eq. 1).
Revised: In all experiments, we use the shifted window to specify the fixed window (Eq. 1).

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

Original: In particular, Stoformer ♢♢ is reduced to the traditional fixed window transformerand 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 ourproposed stochastic window strategy to this effect, a toy experiment is performed: we add a Gaussiannoise patch with σ = 50 whose spatial size is 64 × 64 to an 8 -bit pure color image with spatialsize 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 wedenoise these noisy images using the fixed and stochastic window transformer, respectively and the
Revised: Then we denoise these noisy images using the fixed and stochastic window transformer, respectively and the

Original: PSNR distributions with respect to ( ξ h , ξ w ) are normalized to the fixed range [0 , 1] and plotted in Fig.4.
Revised: PSNR distributions with respect to ( ξ h , ξ w ) are normalized to the fixed range [0 , 1] and plotted in Fig.

Original: To further highlight the translation invariance of stochastic window transformer, we also includea typical CNNs-based denoiser VDN [56], which can ensure the translation invariance well due toinherent 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 [57], which can ensure the translation invariance well due to inherent structure of CNNs, as the inference.

Original: Fig.
Revised: Fig.

Original: 4 indicates clearly that there are great variations,especially for the interior, in the PSNR distribution of fixed window transformer while stochasticwindow transformer flattens the distribution significantly.
Revised: 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
Revised: Both stochastic window transformer and

Original: VDN contain the relatively uniform distribution for the interior, which indicates that the stochasticwindow strategy helps to maintain the desired translation invariance quite well.
Revised: 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 thatall these distributions have variations on the image boundary, which can be attributed to the differenceof 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

Revised: [Table]

Revised: [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 elaboratefour Stoformer variants, existing seven deraining methods are included: GMM [29], DDN [13],[Figure] [Table] [Table]
Revised: Specifically, except the elaborate four Stoformer variants, existing seven deraining methods are included: GMM [30], DDN [13],

Original: SPANet [48], JORDER-E [54], RCDNet [47], SPAIR [37], Uformer [51].
Revised: SPANet [49], JORDER-E [55], RCDNet [48], SPAIR [38], Uformer [52].

Original: All these methods areevaluated on SPA-Data [48] and the performance comparison is shown in Tab. 1.
Revised: All these methods are evaluated on SPA-Data [49] and the performance comparison is shown in Tab. 1.

Original: Due to the limitedspace, 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, whichinclude Set 12 [62], BSD 68 [34], Urban 100 [17], Kodak 24 [12] and McMaster [67].
Revised: We conduct denoising experiments on the additive white Gaussian Noise benchmark datasets, which include Set 12 [63], BSD 68 [35], Urban 100 [18], Kodak 24 [12] and McMaster [68].

Original: Followingprevious works [62, 56, 63], a single model is trained to tackle with various noise levels.
Revised: Following previous works [63, 57, 64], a single model is trained to tackle with various noise levels.

Original: Tabs.
Revised: Tabs.

Original: 3 and4 report PSNR scores of existing methods as well as Stoformer variants for color and grayscale imagedenoising, respectively.
Revised: 3 and 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: [Table]

Original: 59 VDN [56]

Original: [-- - - - - - - - FuncNet [33] 34 .- - 35 .- - - - - - - - DRUNet [66] 
-]
Revised: [-- - - - - - - - FuncNet [33] 34 .- - 35 .- - - - - - - - DRUNet [66] 
-]

Original: [-61 Restormer [58] 
-]
Revised: [-61 Restormer [58] 
-]

Original: [-02 Stoformer ♢♢ 
-]
Revised: [-02 Stoformer ♢♢ 
-]

Original: Stoformer

Original: [Table]

Original: [Table]

Original: [Table]

Original: [Table]

Original: [Table]

Original: 61 Stoformer

Original: [Table]

Original: ♢

Original: [Table]

Original: [Table]

Original: [Table]

Original: [Table]

Original: [Table]

Original: [Table]

Original: [Table]

Original: [Table]

Original: [Table]

Original: [-Stoformer †† 35 .
-]
Revised: [-Stoformer †† 35 .
-]

Original: [Table]

Original: [Table]

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

Original: We also perform deblurring experiments on the benchmark datasets (GoPro [36] and HIDE [42]).
Revised: We also perform deblurring experiments on the benchmark datasets (GoPro [37] and HIDE [43]).

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: Tab.
Revised: Tab.

Original: 2 presents
Revised: 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: Fig.
Revised: Fig.

Original: 5 presents an imagedeblurring example from GoPro [36].
Revised: 5 presents an image deblurring example from GoPro [37].

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 followingobservations and analyses:
Revised: From above experiments on various image restoration tasks, we can make the following observations and analyses:

Original:  • Stoformer ♢♢ vs. Stoformer ♢ † : Without the stochastic window strategy for training, directlyapplying the layer expectation propagation will degrade performance dramatically, whichreveals that the performance gain cannot simply attribute to feature ensemble at test timeand the stochastic window for training matters.• Stoformer ♢♢ vs.
Revised: [Table] [Table] • 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. • Stoformer ♢♢ vs.

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

Original: Even without the layer expectation propagation algorithmfor testing, the stochastic window partition for training alone is also conductive to boost per- formance.
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 expec- tation 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

Revised: Complexity Analysis.

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

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

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

Original: Eliminating Blocking Artifacts.
Revised: But the fixed window loses translation invariance and locality. The sliding window can fulfill the translation invariance and locality while suffer from huge memory burden, which often incurs OOM. In contrast, the stochastic window enjoys efficient training while maintains translation invariance and locality. For testing, as shown in Tab. 6, the fixed window can inference efficiently but suffer from the broken translation invariance and locality loss.

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

Revised: [Table]

Revised: [Table]

Revised: Θ( BHWCs 2 + BHWC 2 ) Θ( BHWC + BHWhs 2 ) (cid:37)

Revised: [Table]

Revised: Θ( BHWCs 2 + BHWC 2 ) Θ( BHWCs 2 + BHWhs 2 ) (cid:33)

Revised: [Table]

Revised: Θ( BHWCs 2 + BHWC 2 ) Θ( BHWC + BHWhs 2 ) (cid:33)

Revised: Strategy

Revised: [Table]

Revised: Space Complexity

Revised: Trans.

Revised: Inva. & Locality

Revised: Fixed window Θ( BHWCs 2 + BHWC 2 ) Θ( BHWC + BHWhs 2 ) (cid:37)

Revised: Sliding

Revised: [Table]

Revised: Θ( BHWCs 2 + BHWC 2 ) Θ( BHWCs 2 + BHWhs 2 ) (cid:33)

Original:  Since the fixed window strategy implies infinite preference towardsspecific window partition, the learned feature map will consequently contain annoying blockingartifacts.
Revised: [Table] window Θ( BHWCs 4 + BHWC 2 ) Θ( BHWC + BHWhs 2 ) (cid:33) Eliminating Blocking Artifacts. Since the fixed window strategy implies infinite preference towards specific window partition, the learned feature map will consequently contain annoying blocking artifacts.

Original: By 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 shouldbe significantly eliminated by the stochastic window strategy.
Revised: Therefore, the blocking artifacts should be significantly eliminated by the stochastic window strategy.

Original: Fig.
Revised: Fig.

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

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

Original: It is surprising that with the stochastic window strategy, the featureis divided into non-overlapping local windows for further processing but in turn nearly no blockingartifacts 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. [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 yieldedby 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 modelperformance.
Revised: Here, we investigate the influence of the number of averaged features to model performance.

Original: Fig.
Revised: Fig.

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

Original: Fig.
Revised: Fig.

Original: 7 reveals that PSNR first risesrapidly and then tends to saturate with the increasing of the averaged feature number while FLOPsincreases linearly with that number, which provides a promising trade-off between performance and[Figure]
Revised: 7 reveals that 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

Original: FLOPs.
Revised: FLOPs.

Original: For instance, we can decrease the averaged feature number from 64 to 8 to keep relativelyhigh 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: More Powerful Representation.
Revised: Boosting Model Generalization.

Original: Fig.
Revised: Fig.

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

Original: Compared with the fixed window strategy, ourproposed stochastic window strategy attainsnot only higher PSNR on the test set but also lowertraining error, which demonstrates that the stochastic window strategy endows model more powerfulrepresentation.
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 completelocal 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.

Revised:  Limitation and Discussion

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.

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

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

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

Original: Conclusions
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. Local attention based transformers, e.g., Swin Transformer [31], have achieved remarkable success on several CV tasks, including image classification, dense prediction, and semantic segmentation. However, they also compute attention based on specific local windows and cannot treat all the local windows fairly. The consequences are that unexpectable information loss occurs during feature processing and the translation invariance is also absent.

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

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

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 figureout that the fixed window strategy will inevitably lead to the broken translation invariance and loss oflocality.
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 efficientlayer 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 restorationtasks 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 frominherent noise of capturing instruments, the shaking during shooting, to unpredictable weathercondition.
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: Ourproposed stochastic window strategy can make transformer more sophisticated in restoring degradedimages.
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: Forexample, the deviation from the actual image textures caused by image restoration technology mayaffect 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 combineexpert knowledge to make rational decisions.
Revised: In these cases, it is necessary to combine expert knowledge to make rational decisions.

Original: Checklist

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

Original: [Yes] See Sec.5.(d)

Original: Have you read the ethics review guidelines and ensured that your paper conforms tothem?

Original: [Yes]2.

Original: If you are including theoretical results...

Original: (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?

Original: [Yes]

Original: If you ran experiments...

Original: (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]

Original: (b) Did you specify all the training details (e.g., data splits, hyperparameters,

Original: how they

Original: were chosen)?

Original: [Yes] See Sec.

Original: 3.3 and supplemental material.

Original: (c) Did you report error bars (e.g., with respect to the random seed after running experi-

Original: ments multiple times)? [No] The random seed is fixed for reproducibility. (d) Did you include the total amount of compute and the type of resources used (e.g., typeof GPUs, internal cluster, or cloud provider)? [Yes] See supplemental material.
Revised: Acknowledgement

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? [Yes] (b) Did you mention the license of  the assets?
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, 61872327 and

Original: [Yes]

Original: (c) Did you include any new assets either in the supplemental material or as a URL? [No] (d) Did you discuss whether and how consent was obtained from people whose data you’re using/curating?
Revised: 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: [N/A]

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

Original: (e) Did you discuss whether the data you are using/curating contains personally identifiable

Original: information or offensive content?

Original: [N/A]

Original: The data we used contain no personally

Original: identifiable information or offensive content.

Original: If you used crowdsourcing or conducted research with human subjects...

Original: (a) Did you include the full text of instructions given to participants and screenshots, if

Original: [Table]

Original: applicable?

Original: [N/A]

Original: (b) Did you describe any potential participant risks, with links to Institutional Review

Original: Board (IRB) approvals, if applicable?

Original: [N/A]

Original: (c) Did you include the estimated hourly wage paid to participants and the total amountspent on participant compensation?

Original: [N/A]