# Sub-token ViT Embedding via Stochastic Resonance Transformers

Dong Lao1 Yangchao Wu1 Tian Yu Liu1 Alex Wong2 Stefano Soatto1
UCLA Vision Lab1 Yale Vision Lab2
<EMAIL_ADDRESS>
<EMAIL_ADDRESS>
###### Abstract

We discover the presence of quantization artifacts in Vision Transformers
(ViTs), which arise due to the image tokenization inherent in these
architectures. These artifacts result in coarsely quantized features, which
negatively impact performance, especially on dense prediction tasks. We
present a zero-shot method to improve how pre-trained ViTs handle spatial
quantization. We propose to ensemble the features obtained from perturbing the
input via sub-token spatial translations, inspired by Stochastic Resonance, a
method traditionally applied to signal processing. We term our method
“Stochastic Resonance Transformer” (SRT), which we show can effectively super-
resolve features of pre-trained ViTs, capturing more of the local fine-grained
structures that might otherwise be neglected as a result of tokenization. SRT
can be applied at any layer, on any task, and does not require any fine-
tuning. The advantage of the former is evident when applied to monocular depth
prediction, where we show that ensembling model outputs are detrimental while
applying SRT on intermediate ViT features outperforms the baseline models by
an average of $4.7\%$ and $14.9\%$ on the RMSE and RMSE_log metrics across
three different architectures. On semi-supervised video object segmentation,
SRT also improves over the baseline models uniformly across all metrics, and
by an average of $2.4\%$ in F&J score. We further show that these quantization
artifacts can be attenuated to some extent via self-distillation. On the
unsupervised salient region segmentation, SRT improves upon the base model by
an average of 2.1% on the maxF metric. Finally, despite operating purely on
pixel-level features, SRT generalizes to non-dense prediction tasks such as
image retrieval and object discovery, yielding consistent improvements of up
to $2.6\%$ and $1.0\%$ respectively.

##  1 Introduction

The Transformer architecture, that takes quantized or “tokenized” inputs,
seems ill fit for vision tasks, since images do not have a natural
discretization scale: The same object can disappear within a pixel or fill the
entire image plane depending on its distance from the camera. Yet Vision
Transformers (ViTs) have been shown effective especially in semantic inference
tasks, so we focus on simple methods to use pre-trained ViT models while
addressing some of the shortcomings of a fixed spatial quantization of the
input tokens.

The standard remedy for quantization artifacts is anti-aliasing. In one-
dimensional signals such as audio, anti-aliasing refers to averaging nearby
samples, or equivalently translated versions of the signal. For images, in
addition to quantization of the translation group reflected in the size of the
pixels, there is also the scale (semi-)group, reflected in the size of
projection of objects onto the image plane. Various network architectures
comprise spatial average pooling, which is translational anti-aliasing, and
the notion of domain-size pooling has been introduced in local features by
(Dong & Soatto, 2015). While traditionally antialiasing is performed via
convolution with a fixed kernel, Stochastic Resonance simply perturbs the data
with respect to an artificial distribution and then averages the results.
Stochastic resonance can be thought of as a way of performing data
augmentation, or adaptive quantization. This simplistic approach is well
suited to pre-trained transformers since it only requires acting on inputs and
outputs without modifying (or even knowing) the weights of the model.
Stochastic Resonance is used to resolve coarsely quantized signals beyond the
native resolution of the sensor and has found wide applicability in cochlear
implants. We apply the same mechanism to the spatial dimension, by perturbing
the input signal to a ViT, which results in highly variable outcomes for the
embeddings. Such outcomes are aggregated statistically to first-order (mean or
median) and second-order (dispersion) to yield a sub-token embedding along
with a measure of confidence or consistency, which have broad potential
applications. For instance, the median outcome can be used as a seed for
unsupervised object segmentation (along with motion masks), and the dispersion
can be used as a weight for an adaptive regularizer.

Input |  |  Visual-language CLIP |  Classification Supervised |  Contrastive DINO |  Segmentation SAM |  Reconstruction MAE
---|---|---|---|---|---|---
|  ViT Feature |  |  |  |  |
|  Ours |  |  |  |  |
|  ViT Feature |  |  |  |  |
|  Ours |  |  |  |  |
Figure 1: High-resolution ViT features Computed by Stochastic Resonance.
Stochastic Resonance enables super-resolving tokenized ViT features during
inference without the need for additional training or modifying ViT forward
pass. Features visualized via Principal Component Analysis are helpful in
analyzing different pre-trained ViT models: CLIP (Radford et al., 2021)
captures major image components. Interestingly, although Supervised
(Dosovitskiy et al., 2020) and DINO (Caron et al., 2021) are trained by
different pipelines and training loss, they prioritize similar regions. This
may be due to they are trained on the same dataset and thus capture similar
inductive bias. In contrast, SAM (Kirillov et al., 2023) and MAE (He et al.,
2022) capture local features over high-level semantics. Stochastic Resonance
not only serves as a powerful visualization tool but also enhances model
performance across multiple downstream tasks, as demonstrated in Sect. 3.

We call the resulting methods “Stochastic Resonance Transformer” although we
do not modify the transformer itself. Instead, we can leverage ViTs, pre-
trained on large datasets, such as CLIP and DINO, to improve their handling of
spatial quantization. This may help attenuate some of the biases of these
datasets, for instance the object-centric nature of DINO, which biases the
representation towards centered objects that occupy a large portion of the
visual field. Stochastic Resonance Transformer can be used to super-resolve
and ensemble the feature maps in ViT, outputting features that reveal some of
the local fine-grained image structure. This can be done at any ViT layer, on
any task, without altering network architecture or pre-trained network
weights. We can optionally distill the fine-grained features back to the
original ViT scale, where we notice performance increase at equal inference
time and cost. Our contributions are listed as follows:

  * •

We introduce a novel technique, namely the Stochastic Resonance Transformer
(SRT), that super-resolves ViT embeddings without additional training or
modifications to the ViT’s forward pass.

  * •

SRT yields a versatile visualization tool that can be applied to any layer of
any pre-trained ViT model, offering valuable insights into ViT model
characteristics.

  * •

The enhanced embeddings from SRT can be seamlessly integrated into any task
that utilizes ViT as a feature extractor, thereby serving as a test-time
augmentation and ensemble method.

  * •

We showcase the effectiveness of SRT by consistent improvement on a range of
diverse vision tasks. Notably, it demonstrates significant enhancements on
dense prediction tasks, of up to 14.9% on depth prediction, 2.4% on video
object segmentation, and 2.1% on salient region segmentation.

  * •

We provide an efficient implementation SRT, including parallelization and
recursive aggregation, which reduces computational and memory requirements.

##  2 Method

In this section, we first describe the conceptual framework and pipeline of
SRT in Sect. 2.1, then we formalize SRT in Sect. 2.2, and present details for
efficient implementation in Secton 2.3

###  2.1 Super-resolving ViT embeddings by Stochastic Resonance.

Figure 2: Schematic for Stochastic Resonance Transformer.  Stochastic
Resonance Transformer (SRT) applies controlled perturbations via translations
to input images, extracting features through Vision Transformers (ViTs). These
features are then upsampled to higher resolution and aligned using the inverse
of the applied perturbations. Statistical aggregation, including mean and
median, along the perturbation dimension, produces fine-grained feature
representations. These features find utility in visualization and can also be
seamlessly integrated back into the network for enhanced performance in
downstream tasks.

Given an image $x$ with $N\times M$ resolution, a Vision Transformer (ViT)
divides it into tokens, where each token represents a $n\times m$ rectangular
patch. While tokens can technically overlap, practical ViT models often use
non-overlapping tokens for efficiency due to the quadratic complexity of
transformers with respect to the number of tokens. Consequently, in a certain
layer of ViT, this approach yields a feature map with dimensions
$\frac{N}{n}\times\frac{M}{m}\times C$, where $C$ is the size of the feature
vector determined by architecture, downsampled from the original image and
subsequently losing sub-token spatial information.

Given a trained ViT model, we aim to obtain features in a higher resolution
that preserves the spatial information on a pixel level, ideally matching with
the original image input. Fig. 2 illustrates our proposed pipeline of SRT. To
achieve super-resolution of the features, we introduce random sub-token
perturbation to the input, i.e. transforming the coordinates of the input and
resampling onto a new image plane, and extract embeddings from the resulting
perturbed image. We then upsample the resulting low-resolution embeddings back
to the original image resolution $N\times M$ and apply an inverse of the
perturbation to the spatial coordinates of the embeddings, and through an
inverse warp, align it with the original input image.

By repeating this process on different sub-token perturbations for $t$ times,
we generate a collection of embeddings, denoted by $N\times M\times C\times
t$, that are spatially aligned to the input frame of reference. We can then
compute statistics, e.g. mean or median, along the $t$ dimension.
Consequently, we obtain a feature field $N\times M\times C$, with the same
spatial resolution as the original image. As showcased in Fig. 1, the
embeddings are “super-resolved” to sub-token resolution. This process is
similar to Stochastic Resonance, where introducing white noise to the input
signal super-resolves a signal beyond the native resolution. These embeddings
offer promising downstream applications, as discussed in more detail in Sect.
3.

For any task that utilizes ViT as a feature extractor, we can take an
additional step by applying average pooling to again tokenize this high-
resolution feature, to map it to $\frac{N}{n}\times\frac{M}{m}\times C$. It’s
important to note that this feature differs from the one obtained from one
single forward pass of ViT, as it is an aggregate of multiple perturbed
inputs. This process can be viewed as test-time augmentation and ensemble.
Since this feature is compatible with the original ViT architecture, it can be
seamlessly integrated into any model at any layer, regardless of pre-training,
without requiring additional learned modules or altering the forward pass.
Such a pipeline improves performance on diverse computer vision tasks, as
validated by Sect. 3. Next, we formalize the aforementioned pipeline.

###  2.2 Formalization

$x\in{\mathbb{R}}^{N\times M\times K}$ is a $K$-channel signal (e.g., $K=3$
for a color image.) Let $\pi:{\mathbb{R}}^{N\times
M}\rightarrow{\mathbb{R}}^{n\times m};x\mapsto x$ a projection (subsampling,
$n\ll N,m\ll M$), with the corresponding inverse (interpolation) map
$\pi^{-1}:{\mathbb{R}}^{n\times m}\rightarrow{\mathbb{R}}^{N\times M};x\mapsto
x$ be piecewise constant. This is a trivial form of subsampling and
interpolation with a constant kernel.

Now, let $\phi:{\mathbb{R}}^{NMK}\rightarrow{\mathbb{R}}^{nmC}$ a trained
model with $C$ channels of feature maps, typically $C\gg K$. Finally, let
$T:{\mathbb{R}}^{N\times M}\rightarrow{\mathbb{R}}^{N\times M};x\mapsto Tx$ a
compact and invertible transformation, for instance, e padded shift by a
number of pixels smaller than $(N-n)/n\times(M-m)/m$. We consider uniform
random padded shifts (translation) and consider the following measurement
process:

$y_{t}=\phi(T_{t}x)$ (1)

for all random transformations $T_{t}$. We wish to super-resolve111We call
this process immersion since each point $x$ maps to $z=\phi(x)$ but $z\neq
T^{-1}\phi(Tx)$. In other words, $x$ is mapped injectively but not
bijectively, since there are as many (vector)-values as sampled value of $T$.
the output of $\phi$ from $n\times m$ to $N\times M$. We do so iteratively by
averaging (or by a trainable linear transformation $K_{t}$) with respect to
the innovation process:

$\epsilon_{t}=\underbrace{\pi\left(T_{t}^{-1}\pi^{-1}y_{t}\right)}_{\hat{y}_{t}}-K_{t}\phi(x)$
(2)

now the super-resolved features which we call $\hat{x}_{t}$ are obtained by an
observer architecture, which implements a closed-loop dynamical system of the
form:

$\begin{cases}\hat{x}_{t+1}=\hat{x}_{t}+T_{t}^{-1}\pi^{-1}y_{t}\quad\hat{x}_{0}=0;\\\
y_{t}=\phi(T_{t}x)\end{cases}$ (3)

This is just a super-resolved moving average, whereby the variance of
$\hat{x}$ will decrease to a steady state (by Central Limit Theorem),
following the practice of stochastic resonance. It is a mixture of
upsampling/interpolation and inference-time data augmentation, or ensembling.

###  2.3 Efficient Implementation

In theory, there is no limitation on the types of sub-token transformations
that can be employed. We opt for a straightforward approach by applying
translations (with padding) and this practice demonstrates effective results.
We sample translations at the pixel level, avoiding the need for sub-pixel
interpolation, which could introduce unwanted artifacts.

For a ViT utilizing token sizes of $m\times n$, we impose a constraint on the
maximum magnitude of translation, limiting it to
$\frac{m}{2}\times\frac{n}{2}$. This constraint allows the model to explore
all possible token selections within the image. It is worth noting that
excessive translation can be counterproductive when applied to downstream
vision tasks, as it can result in information loss at the image boundaries. A
detailed discussion can be found in Sect3.2, where we study the relation
between perturbation level and model performance.

Furthermore, our framework facilitates computational acceleration through
batching. Each ViT forward pass on different augmented images operates
independently, enabling parallelization. Nevertheless, in practical vision
tasks, GPU memory constraints may pose challenges. This is particularly true
after upsampling the embeddings, as each $M\times N\times C$ tensor consumes a
significant amount of memory, and there could be numerous such tensors. To
address this issue, we employ recursive mean computation to iteratively
aggregate information. In the cases where the sole objective is to obtain the
embeddings after average pooling, we can simplify the pipeline by even
bypassing the upsampling step. Instead, we explicitly hard-code the average
pooling of SRT in a recursive manner, by exploiting the property of
translation as one averaged token can be explicitly computed by 4 neighboring
tokens222Our implementation will be publicly available upon paper
publication.. This approach enhances inference speed by a factor of 25
compared to a naive implementation of SRT and substantially reduces GPU memory
usage. With ViT-16/S architecture, on DAVIS-2017 (Pont-Tuset et al., 2017) our
implementation of SRT runs at 2.2 seconds per image on a Nvidia 1080Ti GPU
using a perturbation level of 3 pixels, which is 14 times slower than running
a single forward pass, despite SRT requires 49 independent ViT forward passes.
To further speed up, one may optionally fine-tune the ViT model by distilling
utilizing SRT, so that the inference time and cost remain, as demonstrated in
Sect. 3.2.

##  3 Experiments

In this section, we first showcase the Stochastic Resonance Transformer (SRT)
as a powerful visualization tool. We use it to visualize fine-grained ViT
embeddings from well-established ViT models, illustrating SRT’s capacity to
provide insights into the properties of ViTs. We then extend the application
of SRT to various downstream tasks that utilize ViT as a backbone, employing
SRT in a zero-shot manner without fine-tuning the network during inference.
This evaluation validates SRT’s ability to enhance ViT embeddings for diverse
vision tasks, and these enhanced embeddings seamlessly integrate with the
original tasks, eliminating the need for any additional training. We test on
three dense prediction tasks: semi-supervised video object segmentation,
monocular depth estimation, and unsupervised saliency detection. To address
concerns about potential bias towards dense prediction, we conduct a sanity
check on image retrieval and unsupervised object detection. Remarkably, SRT
consistently improves performance across all five tasks.

|
---|---
Figure 3: Relative improvement on DAVIS-2017 dataset vs different noise
levels. There exists an inherent trade-off between perturbation level and
performance gain. Smaller perturbation ranges result in weaker improvements
from the baseline model due to lower input diversity, while larger
perturbations are susceptible to greater information loss. 3 pixels is found
to be the optimal perturbation level on both ViT-S/16 and Vit-B/16.

###  3.1 Dense ViT Feature Visualization

SRT demonstrates significant promise in visualizing features of ViT models. It
achieves this by enhancing feature resolution through the aggregation of
perturbed inputs, thus without necessitating modifications to the ViT’s
forward pass, in contrast to using overlapping tokens (Amir et al., 2021) that
impose substantial GPU memory demands. Notably, all visualizations in this
paper are computed by a standard consumer laptop.

In Fig. 1, we present visualizations of the final layer features from five
popular ViT models, all employing the ViT-B/16 architecture. We employ SRT
with a turbulence level of 7 pixels to traverse non-overlapping augmented
tokens extensively. The resultant high-dimensional features then go through
Principal Component Analysis (PCA), with the top three components mapped to
RGB channels to facilitate effective visualization. Despite sharing the same
architecture, the five models exhibit distinct characteristics owing to
variations in their pre-training supervision. For instance, CLIP (Radford et
al., 2021) is trained through contrastive visual-language pre-training and
captures major image components in the displayed examples. The Supervised
model (Dosovitskiy et al., 2020) is trained for ImageNet classification, while
DINO (Caron et al., 2021) undergoes contrastive learning. Interestingly,
despite their diverse training regimes, both models prioritize similar image
regions, potentially due to their shared dataset and resulting common
inductive bias. In contrast, SAM (Kirillov et al., 2023) is trained on massive
segmentation masks without semantic labels or object-centric priors, and MAE
(He et al., 2022) is trained through inpainting of randomly masked image
regions. Both methods emphasize local image features over high-level
semantics. Our versatile visualization tool provides valuable insights into
the characteristics of ViT models, offering substantial potential for
practical applications.

###  3.2 Semi-supervised Video Object Segmentation

We apply the Stochastic Resonance Transformer (SRT) to evaluate its
performance using the DAVIS-2017 video instance segmentation benchmark (Pont-
Tuset et al., 2017). We adhere to the experimental methodology established in
(Jabri et al., 2020), which employs a ”semi-supervised” video object
segmentation approach. Provided with the initial annotation of the objects of
interest in the first frame, this method subsequently propagates the
segmentation between consecutive frames. Notably, the method utilizes the last
layer feature of the Vision Transformer (ViT) to guide this segmentation
propagation process. Consequently, the quality of the ViT features directly
impacts the final segmentation results. For optimal outcomes, these features
must possess discriminative and semantically meaningful characteristics to
effectively support this segmentation task.

In our study, we evaluate various Vision Transformer (ViT) models pre-trained
using the DINO (Caron et al., 2021) contrastive scheme. We adopt three
different architectures, specifically ViT-S/16, ViT-B/16, and ViT-S/8, each
varying in their spatial patch size (16x16 pixels and 8x8 pixels). Our
results, presented in Tab. 1, indicate that, on average, the Stochastic
Resonance Transformer (SRT) enhances the original baseline models by a
relative 2.4% in terms of the F&J score. The most significant improvement is
observed with ViT-S/16, where the improvement reaches 4.1%. Importantly, these
enhancements are achieved without any modifications to the model or pre-
trained weights. However, we address a potential criticism of our approach,
which could be seen as trivial test-time augmentation combined with feature-
level ensemble. To counter this concern, we perform a heuristic by naively
augmenting images by color jitter and performing feature-level ensemble, and
we find that this method is, in fact, detrimental to performance. We also
investigate whether inference costs induced by SRT can potentially be
mitigated via distillation. Toward this goal, we attempt to learn the
ensembled SRT representations using the following self-distillation objective:

$\min_{w}\sum_{x\in\mathcal{D}}||\phi_{w}(x)-SRT(x,w_{0})||,$ (4)

where $\phi$ and ($w_{0}$) $w$ are the ViT and its (original) parameters, and
$x$ the image in the target dataset. Our preliminary results on DINO-ViT/16
improve from the baseline by $1.3\%$ after the self-distillation step. Note
that Eq. equation 4 is agnostic to the task and requires no label, rendering
distillation by SRT a fine-tuning scheme that adapts pre-trained ViT features
to new target datasets. We leave the investigation of this to future work.

Fig. 3 illustrates the relative improvement across different perturbation
levels of SRT applied to ViT-S/16 and ViT-B/16. While higher perturbation
levels offer greater input diversity, they are also susceptible to information
loss. We anticipate a trade-off between perturbation level and performance
gain and empirically identify a perturbation level of 3 pixels as the optimal
point for both.

Method | F&J-Mean | J-Mean | J-Recall | F-Mean | F-Recall
---|---|---|---|---|---
DINO-ViT-S/16 | 0.617 | 0.602 | 0.740 | 0.634 | 0.764
\+ SRT | 0.642 | 0.632 | 0.783 | 0.653 | 0.819
Distill by SRT | 0.625 | 0.609 | 0.745 | 0.642 | 0.780
DINO-ViT-B/16 | 0.622 | 0.608 | 0.748 | 0.637 | 0.760
\+ SRT | 0.630 | 0.623 | 0.766 | 0.637 | 0.795
DINO-ViT-S/8 | 0.706 | 0.675 | 0.815 | 0.737 | 0.846
\+ SRT | 0.720 | 0.688 | 0.827 | 0.752 | 0.868
Naive ensemble* | 0.477 | 0.455 | 0.468 | 0.500 | 0.542
Table 1: Results on DAVIS-2017 video object segmentation. Applying SRT
improves over the baseline models uniformly over all metrics, as measured
across 3 variants of ViTs trained using the DINO (Caron et al., 2021)
contrastive learning objective. SRT yields significant improvements even for
ViT-S/8 trained with finer patch sizes (8x8). SRT also uniformly outperforms
naively augmenting images and averaging features at inference time, as denoted
by $\ast$. One may optionally fine-tune the original ViT model by  distilling
by SRT, which increases performance while inference time and cost remains.

###  3.3 Monocular Depth Prediction

We extend the application of SRT to monocular depth estimation, a task that
leverages ViT features from multiple ViT layers, in contrast to video object
segmentation which primarily utilizes the last layer features. This choice of
task highlights the versatility of SRT, showcasing its seamless compatibility
with various ViT layers and architectures. Specifically, we evaluate three ViT
architectures: ViT-S/14, ViT-B/14, and ViT-L/14, each equipped with two
prediction heads (linear and DPT (Ranftl et al., 2021)). We adopt the
experimental settings provided by DINOV2, which offers pre-trained backbones
and corresponding prediction heads. Our assessment utilizes the NYU-V2 dataset
(Nathan Silberman & Fergus, 2012).

Tab. 2 presents the results, demonstrating consistent improvements over
baseline methods. The most significant enhancements are observed in the RMSE
and RMSE_log metrics, where we achieve relative improvements of 4.7% and 14.9%
with linear heads, and 3.6% and 11.0% with DPT heads, respectively. Notably,
these metrics are sensitive to outliers, highlighting the effectiveness of our
approach in mitigating instability in ViT features and enhancing robustness.
For completeness, we compare our method with output-space ensemble (marked as
”OE”), which employs the perturbations as SRT, and aggregates the model output
instead of intermediate features. We find no significant improvements, and in
some cases, this method is even detrimental. This underscores the robustness
of SRT’s feature ensemble scheme.

Backbone | Head | Method | RMSE | RMSE_log | AbsRel | SqRel | a1 | a2 | a3
---|---|---|---|---|---|---|---|---|---
DINOV2-ViT-B/14 | Linear | Baseline | 0.396 | 0.135 | 0.100 | 0.061 | 0.903 | 0.983 | 0.996
+OE | 0.376 | 0.121 | 0.093 | 0.059 | 0.918 | 0.984 | 0.997
+SRT | 0.349 | 0.108 | 0.087 | 0.052 | 0.930 | 0.990 | 0.998
DPT | Baseline | 0.323 | 0.109 | 0.074 | 0.044 | 0.941 | 0.987 | 0.996
+OE | 0.314 | 0.101 | 0.073 | 0.043 | 0.944 | 0.988 | 0.997
+SRT | 0.305 | 0.096 | 0.073 | 0.043 | 0.945 | 0.989 | 0.997
DINOV2-ViT-S/14 | Linear | Baseline | 0.471 | 0.162 | 0.125 | 0.084 | 0.853 | 0.972 | 0.994
+OE | 0.486 | 0.153 | 0.126 | 0.095 | 0.858 | 0.974 | 0.994
+SRT | 0.457 | 0.140 | 0.118 | 0.085 | 0.876 | 0.980 | 0.996
DPT | Baseline | 0.336 | 0.114 | 0.080 | 0.048 | 0.933 | 0.986 | 0.996
+OE | 0.347 | 0.114 | 0.080 | 0.053 | 0.932 | 0.985 | 0.996
+SRT | 0.334 | 0.104 | 0.080 | 0.051 | 0.935 | 0.988 | 0.996
DINOV2-ViT-L/14 | Linear | Baseline | 0.373 | 0.127 | 0.093 | 0.054 | 0.916 | 0.985 | 0.996
+OE | 0.401 | 0.131 | 0.097 | 0.062 | 0.908 | 0.982 | 0.996
+SRT | 0.365 | 0.113 | 0.090 | 0.053 | 0.924 | 0.989 | 0.998
DPT | Baseline | 0.311 | 0.105 | 0.070 | 0.042 | 0.946 | 0.988 | 0.997
+OE | 0.317 | 0.103 | 0.072 | 0.044 | 0.942 | 0.987 | 0.996
+SRT | 0.297 | 0.092 | 0.070 | 0.041 | 0.947 | 0.991 | 0.997
Table 2: Results on NYU-V2 depth prediction. Our method can be extended
without modification to improve intermediate features to yield improved
performance on the downstream depth prediction tasks. While ensembling of
outputs (OE) can often be detrimental to performance, applying SRT on the
features from pre-trained backbones (inputs to prediction heads) can improve
performance over baselines by $4.7\%$ and $14.9\%$ on RMSE and RMSE_log, using
the linear prediction head and by $3.6\%$ and $11.0\%$ using the DPT head.

###  3.4 Unsupervised Salient Region Segmentation

We employ SRT in conjunction with TokenCut (Wang et al., 2022) for
unsupervised salient region segmentation tasks. TokenCut is a graph-based
approach that applies the Normalized Cut algorithm to partition ViT features
into two distinct clusters, representing the salient foreground and the
background. The key challenge here is to ensure that the features are not only
discriminative across clusters but also consistent within clusters. We execute
TokenCut without any post-processing, such as Conditional Random Fields (CRF),
to assess the raw quality of ViT embeddings.

Our method is evaluated on three datasets: ECSSD (Shi et al., 2015), DUTS
(Wang et al., 2017), and DUT-OMRON (Yang et al., 2013). Due to computational
limitations, we exclusively use the ViT-S/16 architecture, as TokenCut demands
substantial memory resources when applied to a larger number of ViT tokens.
The results are in Tab. 3. Across all three datasets, we consistently observe
improvements, with an average increase in the maxF metric of 2.1%. Notably,
this improvement is constrained by the model architecture, as TokenCut
operates at the coarse segmentation level of ViT tokens. Given that SRT has
the capability to provide finer-grained features (directly applying TokenCut
at this level is computationally impractical due to its $O(n^{2})$ complexity,
where $n$ is the number of tokens), we anticipate that future research will
develop feasible methods to leverage SRT’s high-resolution embeddings
effectively.

Datasets | ECSSD | DUTS | DUTS-OMRON
---|---|---|---
Feature Extractor | maxF | IoU | Acc. | maxF | IoU | Acc. | maxF | IoU | Acc.
DINO ViT-S/16 | 80.3 | 71.2 | 91.8 | 67.2 | 57.6 | 90.3 | 60.0 | 53.3 | 88.0
+SRT | 82.4 | 71.7 | 92.1 | 68.8 | 58.3 | 90.6 | 60.8 | 53.7 | 88.0
Table 3: Results on unsupervised salient region segmentation. Despite
architectural constraints, our method yields consistent improvement on all
three datasets, with an average increase of 2.1% in the maxF metric.

###  3.5 Sanity Check: Image Retrieval and Unsupervised Object Detection

Incorporating SRT into vision tasks involves updating ViT features based on
fine-tuned super-resolved features. However, questions remain regarding
whether the observed enhancements in dense prediction tasks are solely due to
increased awareness of semantic boundaries in images and whether this method
extends to non-dense prediction tasks. To address these concerns, we conducted
a sanity check using image retrieval and unsupervised object detection tasks.

For image retrieval, we applied a nearest-neighbor protocol following DINO,
using the Oxford image retrieval datasets (Radenović et al., 2018). Notably,
our base model’s pre-training poses a substantial domain gap to the target
datasets. Although image retrieval primarily requires distinctive image-level
features (rather than pixel-level), aiming to match images to queries at a
higher level, SRT exhibited effective adaptation, resulting in a notable 2.6%
relative improvement in accuracy.

Regarding unsupervised object detection, we utilized TokenCut and the VOC07
dataset (Everingham et al., 2010). Unsupervised object detection focuses on
region-level discriminative features, utilizing bounding boxes instead of
segmentation masks for object delineation. Despite this, we observed a 1.0%
relative improvement in the detection rate, reaffirming that SRT does not
compromise the information within the original ViT embeddings. These results
serve as a critical validation of SRT’s capacity to enhance ViT features
without distorting their original information.

Task | Metric | Baseline | 1 | 2 | 3 | 4 | 5 | 6
---|---|---|---|---|---|---|---|---
Image Retrieval | Medium | 34.6 | 34.8 | 35.1 | 35.2 | 35.3 | 35.3 | 35.5
Hard | 13.0 | 13.1 | 13.2 | 13.1 | 13.2 | 13.2 | 13.1
Object Discovery | Detection Rate | 68.7 | 68.9 | 68.9 | 69.2 | 69.4 | 69.3 | 69.2
Table 4: Results on Image Retrieval and Object Discovery. Despite operating
purely on pixel-level features, SRT generalizes to non-dense prediction tasks
operating on higher-level region/image features to yield equal or better
performance compared to the standard inference baseline. On the Oxford image
retrieval task, SRT on the DINO-ViT-16 model yields up to $2.6\%$ relative
improvement from the baseline model. On the unsupervised object detection
task, SRT improves the detection rate by up to $1.0\%$.

##  4 Discussion

###  4.1 Related Work

Stochastic Resonance  was proposed by Benzi et al. (1981) and first applied in
climate dynamics (Benzi et al., 1982) and later in signal processing (Wellens
et al., 2003; Kosko & Mitaim, 2001; Chen et al., 2007) and acoustics (Shu-Yao
et al., 2016; Wang et al., 2014). It is used to super-resolve a signal beyond
the native resolution of the sensor by adding white noise. We use the same
principle to adapt generic ViT image features for dense prediction downstream
tasks. By randomly translating the images, (i.e. introducing noise in the
spatial dimension), we are able to super-resolve ViT image features to be
smoother and better suited for dense prediction tasks. We leave extensions to
other groups or semi-groups of transformations (e.g., scale or domain size) to
future work.

Test-time data augmentation involves aggregating model predictions from
augmented test input to a final prediction. Applying such a technique
increases the robustness of predictions (Prakash et al., 2018; Song et al.,
2017; Cohen et al., 2019) and prediction accuracy (Krizhevsky et al., 2012;
Szegedy et al., 2015; Simonyan & Zisserman, 2014; Jin et al., 2018; Matsunaga
et al., 2017) in a variety of tasks. It can also used to estimate the
uncertainty of the model (Matsunaga et al., 2017; Smith & Gal, 2018; Ayhan &
Berens, 2022; Wang et al., 2019). Different transformations are used to target
different potential tasks: Pang et al. (2019) linearly combines the testing
input and a randomly sampled clean image to generate classification
prediction. Isensee et al. (2018) performs flipping and rotation to the test
input image to generate 64 different inputs and finally aggregates the outputs
to perform medical image segmentation. Krizhevsky et al. (2012) crops the
images into smaller patches and ensemble the results for classification. Self-
ensembling (Bousselham et al., 2021) is also closely related to our work.
Bousselham et al. (2021) leverages multi-scale features fed into multiple
independent decoders to create an ensemble within a single model. Liu et al.
(2018) ensembles outputs from networks augmented with random noise layers to
improve model robustness. SRT aggregates information via adding spatial
translations as noise and can be considered a general case of test-time
augmentation, where ensembling is performed at the feature level at
intermediate layers of a ViT, instead of the output level, which is novel.

Knowledge distillation aims to transfer the knowledge from stronger teacher
models to weaker student models to improve their performance. Hinton et al.
(2015) trains a student model to mimic the soft output distribution of the
teacher model. Romero et al. (2014) extends this idea to distill the
intermediate features learned by the teacher models. We consider a form of
self-distillation (Zhang et al., 2019), in which the student itself is used as
the teacher to improve learned representations.

Dense ViT feature extractor.  Our work is closely related to (Amir et al.,
2021),which employs ViT for generating dense visual descriptors. To extract
these fine-grained features, (Amir et al., 2021) reduce the stride allowing
for overlapping tokens and performing a single forward pass with ViT. In SRT,
instead of a single pass, we conduct multiple passes using perturbed inputs.
This modification reduces the computational complexity from quadratic to
linear.

###  4.2 Conclusion, Limitation, and Future Work

Averaging sub-token perturbations functions as a versatile visualization tool
for ViT embeddings and offers test-time augmentation and ensemble capabilities
usable with any ViT architecture. We found the technique especially useful for
dense prediction tasks, where even coarse-scale obejcts have fine-scale
occlusion boundaries that must be resolved.

In contrast to most test-time augmentation and ensemble methods that operate
at the output level, and require task-specific designs, our method can be
applied to any layer within any architecture and any task that utilizes ViT as
a feature extractor, eliminating the need for modifications to the forward
pass, such as increasing token numbers, which would significantly increase GPU
usage. Instead, SRT accomplishes this by conducting multiple inferences,
avoiding the quadratic complexity related to the number of ViT tokens, and
enabling super-resolution of ViT embeddings to match the input resolution.
This approach is amenable to parallelization through batching, ensuring
computational efficiency. Furthermore, the method allows ensembling without
memory-intensive resizing of all embeddings to full resolution, which can be
executed recursively, as described in Sect. 2.3. Practical implementations
demonstrate efficient execution on even laptop GPUs.

With all that said, Stochastic Resonant Transformers have several limitations.
The basic embodiment increases inference cost and latency, as each perturbed
image necessitates a ViT forward pass. To address this, one viable approach is
knowledge distillation, which involves fine-tuning the network to mimic the
feature-level output of SRT. We illustrate this process using the DAVIS-2017
training dataset with DINO-ViT-S/16, achieving improved results (F&J-score
$0.617\Rightarrow 0.625$) without the use of labels or operations on the
validation set. This establishes a label-free, task-free transductive fine-
tuning scheme that adapts pre-trained ViT features to new target datasets.
Future directions may involve refining the distillation process on different
layers and exploring the integration of Stochastic Resonance directly into ViT
architectures.

Additionally, our findings underscores the segmentation capabilities of ViT
embeddings, aligning with recent claims in the field (Caron et al., 2021; Yu
et al., 2023). Super-resolved features exhibit sharp, fine-grained
semantically relevant boundaries. Furthermore, our method leverages the
convexity properties (Park & Kim, 2022) of ViT embeddings, enabling convex
combinations (average pooling as a special case) during inference, resulting
in improvements across various tasks. It is worth noting that Stochastic
Resonance is not limited to ViT architectures nor to spatial quantization. It
can be applied to architectures like CNN as well as to other forms of
quantization, such as sale or domain size. However, our emphasis in this paper
is on ViTs that mostly use non-overlapping tokens, making them particularly
suited to our approach.

## References

  * Amir et al. (2021) Shir Amir, Yossi Gandelsman, Shai Bagon, and Tali Dekel.  Deep vit features as dense visual descriptors.  _arXiv preprint arXiv:2112.05814_ , 2(3):4, 2021\.
  * Ayhan & Berens (2022) Murat Seckin Ayhan and Philipp Berens.  Test-time data augmentation for estimation of heteroscedastic aleatoric uncertainty in deep neural networks.  In _Medical Imaging with Deep Learning_ , 2022.
  * Benzi et al. (1981) Roberto Benzi, Alfonso Sutera, and Angelo Vulpiani.  The mechanism of stochastic resonance.  _Journal of Physics A: mathematical and general_ , 14(11):L453, 1981.
  * Benzi et al. (1982) Roberto Benzi, Giorgio Parisi, Alfonso Sutera, and Angelo Vulpiani.  Stochastic resonance in climatic change.  _Tellus_ , 34(1):10–16, 1982.
  * Bousselham et al. (2021) Walid Bousselham, Guillaume Thibault, Lucas Pagano, Archana Machireddy, Joe Gray, Young Hwan Chang, and Xubo Song.  Efficient self-ensemble for semantic segmentation.  _arXiv preprint arXiv:2111.13280_ , 2021.
  * Caron et al. (2021) Mathilde Caron, Hugo Touvron, Ishan Misra, Hervé Jégou, Julien Mairal, Piotr Bojanowski, and Armand Joulin.  Emerging properties in self-supervised vision transformers.  In _Proceedings of the IEEE/CVF international conference on computer vision_ , pp. 9650–9660, 2021.
  * Chen et al. (2007) Hao Chen, Pramod K Varshney, Steven M Kay, and James H Michels.  Theory of the stochastic resonance effect in signal detection: Part i—fixed detectors.  _IEEE transactions on Signal Processing_ , 55(7):3172–3184, 2007.
  * Cohen et al. (2019) Jeremy Cohen, Elan Rosenfeld, and Zico Kolter.  Certified adversarial robustness via randomized smoothing.  In _international conference on machine learning_ , pp. 1310–1320. PMLR, 2019.
  * Dong & Soatto (2015) Jingming Dong and Stefano Soatto.  Domain-size pooling in local descriptors: Dsp-sift.  In _Proceedings of the IEEE conference on computer vision and pattern recognition_ , pp. 5097–5106, 2015.
  * Dosovitskiy et al. (2020) Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al.  An image is worth 16x16 words: Transformers for image recognition at scale.  _arXiv preprint arXiv:2010.11929_ , 2020.
  * Everingham et al. (2010) Mark Everingham, Luc Van Gool, Christopher KI Williams, John Winn, and Andrew Zisserman.  The pascal visual object classes (voc) challenge.  _International journal of computer vision_ , 88:303–338, 2010.
  * He et al. (2022) Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár, and Ross Girshick.  Masked autoencoders are scalable vision learners.  In _Proceedings of the IEEE/CVF conference on computer vision and pattern recognition_ , pp. 16000–16009, 2022.
  * Hinton et al. (2015) Geoffrey Hinton, Oriol Vinyals, and Jeff Dean.  Distilling the knowledge in a neural network.  _arXiv preprint arXiv:1503.02531_ , 2015.
  * Isensee et al. (2018) Fabian Isensee, Jens Petersen, Andre Klein, David Zimmerer, Paul F Jaeger, Simon Kohl, Jakob Wasserthal, Gregor Koehler, Tobias Norajitra, Sebastian Wirkert, et al.  nnu-net: Self-adapting framework for u-net-based medical image segmentation.  _arXiv preprint arXiv:1809.10486_ , 2018.
  * Jabri et al. (2020) Allan Jabri, Andrew Owens, and Alexei Efros.  Space-time correspondence as a contrastive random walk.  _Advances in neural information processing systems_ , 33:19545–19560, 2020.
  * Jin et al. (2018) Hongsheng Jin, Zongyao Li, Ruofeng Tong, and Lanfen Lin.  A deep 3d residual cnn for false-positive reduction in pulmonary nodule detection.  _Medical physics_ , 45(5):2097–2107, 2018.
  * Kirillov et al. (2023) Alexander Kirillov, Eric Mintun, Nikhila Ravi, Hanzi Mao, Chloe Rolland, Laura Gustafson, Tete Xiao, Spencer Whitehead, Alexander C Berg, Wan-Yen Lo, et al.  Segment anything.  _arXiv preprint arXiv:2304.02643_ , 2023.
  * Kosko & Mitaim (2001) Bart Kosko and Sanya Mitaim.  Robust stochastic resonance: Signal detection and adaptation in impulsive noise.  _Physical review E_ , 64(5):051110, 2001.
  * Krizhevsky et al. (2012) Alex Krizhevsky, Ilya Sutskever, and Geoffrey E Hinton.  Imagenet classification with deep convolutional neural networks.  _Advances in neural information processing systems_ , 25, 2012.
  * Liu et al. (2018) Xuanqing Liu, Minhao Cheng, Huan Zhang, and Cho-Jui Hsieh.  Towards robust neural networks via random self-ensemble.  In _Proceedings of the European Conference on Computer Vision (ECCV)_ , pp. 369–385, 2018.
  * Matsunaga et al. (2017) Kazuhisa Matsunaga, Akira Hamada, Akane Minagawa, and Hiroshi Koga.  Image classification of melanoma, nevus and seborrheic keratosis by deep neural network ensemble.  _arXiv preprint arXiv:1703.03108_ , 2017.
  * Nathan Silberman & Fergus (2012) Pushmeet Kohli Nathan Silberman, Derek Hoiem and Rob Fergus.  Indoor segmentation and support inference from rgbd images.  In _ECCV_ , 2012.
  * Pang et al. (2019) Tianyu Pang, Kun Xu, and Jun Zhu.  Mixup inference: Better exploiting mixup to defend adversarial attacks.  _arXiv preprint arXiv:1909.11515_ , 2019.
  * Park & Kim (2022) Namuk Park and Songkuk Kim.  How do vision transformers work?  _arXiv preprint arXiv:2202.06709_ , 2022.
  * Pont-Tuset et al. (2017) Jordi Pont-Tuset, Federico Perazzi, Sergi Caelles, Pablo Arbeláez, Alexander Sorkine-Hornung, and Luc Van Gool.  The 2017 davis challenge on video object segmentation.  _arXiv:1704.00675_ , 2017.
  * Prakash et al. (2018) Aaditya Prakash, Nick Moran, Solomon Garber, Antonella DiLillo, and James Storer.  Deflecting adversarial attacks with pixel deflection.  In _Proceedings of the IEEE conference on computer vision and pattern recognition_ , pp. 8571–8580, 2018.
  * Radenović et al. (2018) Filip Radenović, Ahmet Iscen, Giorgos Tolias, Yannis Avrithis, and Ondřej Chum.  Revisiting oxford and paris: Large-scale image retrieval benchmarking.  In _Proceedings of the IEEE conference on computer vision and pattern recognition_ , pp. 5706–5715, 2018.
  * Radford et al. (2021) Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al.  Learning transferable visual models from natural language supervision.  In _International conference on machine learning_ , pp. 8748–8763. PMLR, 2021.
  * Ranftl et al. (2021) René Ranftl, Alexey Bochkovskiy, and Vladlen Koltun.  Vision transformers for dense prediction.  In _Proceedings of the IEEE/CVF international conference on computer vision_ , pp. 12179–12188, 2021.
  * Romero et al. (2014) Adriana Romero, Nicolas Ballas, Samira Ebrahimi Kahou, Antoine Chassang, Carlo Gatta, and Yoshua Bengio.  Fitnets: Hints for thin deep nets.  _arXiv preprint arXiv:1412.6550_ , 2014.
  * Shi et al. (2015) Jianping Shi, Qiong Yan, Li Xu, and Jiaya Jia.  Hierarchical image saliency detection on extended cssd.  _IEEE transactions on pattern analysis and machine intelligence_ , 38(4):717–729, 2015.
  * Shu-Yao et al. (2016) Ji Shu-Yao, Yuan Fei, Chen Ke-Yu, and Cheng En.  Application of stochastic resonance technology in underwater acoustic weak signal detection.  In _OCEANS 2016-Shanghai_ , pp. 1–5. IEEE, 2016.
  * Simonyan & Zisserman (2014) Karen Simonyan and Andrew Zisserman.  Very deep convolutional networks for large-scale image recognition.  _arXiv preprint arXiv:1409.1556_ , 2014.
  * Smith & Gal (2018) Lewis Smith and Yarin Gal.  Understanding measures of uncertainty for adversarial example detection.  _arXiv preprint arXiv:1803.08533_ , 2018.
  * Song et al. (2017) Yang Song, Taesup Kim, Sebastian Nowozin, Stefano Ermon, and Nate Kushman.  Pixeldefend: Leveraging generative models to understand and defend against adversarial examples.  _arXiv preprint arXiv:1710.10766_ , 2017.
  * Szegedy et al. (2015) Christian Szegedy, Wei Liu, Yangqing Jia, Pierre Sermanet, Scott Reed, Dragomir Anguelov, Dumitru Erhan, Vincent Vanhoucke, and Andrew Rabinovich.  Going deeper with convolutions.  In _Proceedings of the IEEE conference on computer vision and pattern recognition_ , pp. 1–9, 2015.
  * Wang et al. (2019) Guotai Wang, Wenqi Li, Michael Aertsen, Jan Deprest, Sébastien Ourselin, and Tom Vercauteren.  Aleatoric uncertainty estimation with test-time augmentation for medical image segmentation with convolutional neural networks.  _Neurocomputing_ , 338:34–45, 2019.
  * Wang et al. (2014) Jun Wang, Qingbo He, and Fanrang Kong.  Adaptive multiscale noise tuning stochastic resonance for health diagnosis of rolling element bearings.  _IEEE Transactions on instrumentation and measurement_ , 64(2):564–577, 2014.
  * Wang et al. (2017) Lijun Wang, Huchuan Lu, Yifan Wang, Mengyang Feng, Dong Wang, Baocai Yin, and Xiang Ruan.  Learning to detect salient objects with image-level supervision.  In _Proceedings of the IEEE conference on computer vision and pattern recognition_ , pp. 136–145, 2017.
  * Wang et al. (2022) Yangtao Wang, Xi Shen, Yuan Yuan, Yuming Du, Maomao Li, Shell Xu Hu, James L Crowley, and Dominique Vaufreydaz.  Tokencut: Segmenting objects in images and videos with self-supervised transformer and normalized cut.  _arXiv preprint arXiv:2209.00383_ , 2022.
  * Wellens et al. (2003) Thomas Wellens, Vyacheslav Shatokhin, and Andreas Buchleitner.  Stochastic resonance.  _Reports on progress in physics_ , 67(1):45, 2003\.
  * Yang et al. (2013) Chuan Yang, Lihe Zhang, Huchuan Lu, Xiang Ruan, and Ming-Hsuan Yang.  Saliency detection via graph-based manifold ranking.  In _Proceedings of the IEEE conference on computer vision and pattern recognition_ , pp. 3166–3173, 2013.
  * Yu et al. (2023) Yaodong Yu, Tianzhe Chu, Shengbang Tong, Ziyang Wu, Druv Pai, Sam Buchanan, and Yi Ma.  Emergence of segmentation with minimalistic white-box transformers.  _arXiv preprint arXiv:2308.16271_ , 2023.
  * Zhang et al. (2019) Linfeng Zhang, Jiebo Song, Anni Gao, Jingwei Chen, Chenglong Bao, and Kaisheng Ma.  Be your own teacher: Improve the performance of convolutional neural networks via self distillation.  In _Proceedings of the IEEE/CVF international conference on computer vision_ , pp. 3713–3722, 2019.