Predicting masked tokens in stochastic locations improves masked image modeling

Thank you for the thoughtful consideration of the paper and constructive feedback. We’ve incorporated your feedback and uploaded a new revision of the paper.

Q: “As in Figure 2, the context and masked representations are computed by adding their tokens and positional embeddings together”

Thank you for pointing this out, there is a mistake in Figure 2 and we apologize for the confusion. To clarify, the masked and context tokens are computed as follows (as in Algorithm 11):

11: $n \sim \mathcal{N}(0, \sigma I)$

12: $m = An + \psi_{B_y} + \tilde{m}$

13: $c = As_x + b + \psi_{B_x}$

We uploaded a new paper revision and fixed Figure 2.

Q: “For step 11, I suppose $\psi_{B_x}$ should refer to the positional embedding, and $As_x + b$ should refer to context token?

You are right.

Q: Why do we need an additional linear transformation on $s_x$? Some explanations may be needed for this part.

The mapping projects $s_x$ from the output dimension of the encoder to the input dimension of the predictor (this is standard for other approaches like MAE and I-JEPA).

Q: The authors seem to let $s_{x_i}$ (resp. $n_j$) as context (resp. masked) tokens, and $b$ (resp. $\tilde{m}$) corresponds to the bias for context (resp. masked) tokens. However, I suppose $n_j$ should simply be used to compute stochastic positional embedding as in (2), and $s_{x_i}$ is computed from encoder $f_{\theta}$ to encode context information. How can they have the same role?”

We assume that you ask why both the context $s_{x_i}$ and noise $n_j$ linearly projected by matrix $A$. In the original MIM formulation there is no stochastic positions and the masked tokens and context tokens are computed as follows:

12: $m = \psi_{B_y} + \tilde{m}$

13: $c = Bs_x + b + \psi_{B_x}$

When applying StoP with the reparameterization trick (Eq.3, revised manuscript), the noise is now linearly projected with a matrix $A$ (defined in Eq. 1) and summed with the positional embeddings. Therefore, both the sampled noise and the context tokens are now linearly projected:

11: $n \sim \mathcal{N}(0, \sigma I)$

12: $m = An + \psi_{B_y} + \tilde{m}$

13: $c = Bs_x + b + \psi_{B_x}$

Note that here the context tokens and noise use different projections $A$ and $B$. However, we find that the weights of A are quickly scaled down during training, setting $An=0$ which overcomes the noise during training, resorting to the original MIM introduced before, and the same empirical downstream accuracy. We discussed this in Section 3.1 and 3.2 (see “Avoiding a degenerate determinism solution” and “Masked tokens in stochastic locations”.)

To avoid this, we use the same matrix $A$ to also project the context tokens $s_x$ (line 13), instead of using a different projection matrix $B$:

11: $n \sim \mathcal{N}(0, \sigma I)$

12: $m = An + \psi_{B_y} + \tilde{m}$

13: $c = As_x + b + \psi_{B_x}$

The motivation for using $A$ to project both the context features and noise can be understood by considering two extreme cases. When $A=0$, there is complete certainty about the positional embeddings but all context is lost ($As_{x}=0$). On the other hand, when $A$ is large the context information is preserved, but due to the large magnitude of $A$ the noise is amplified and camouflages the positional embedding features of the masked tokens: $An + \psi_{B_y}$. This dual role of matrix A forces the model to balance between location certainty and the influence of context features in predictions. It optimizes the trade-off for each feature, balancing their presence in predictions against the need for precise spatial locations.

We discussed this in Section 3.2 in the original submission but following the comments we revised the manuscript to make this more clear.

Q: I would also like to see more discussions on the connection between StoP and vanilla MIM.

The above answer should clarify this comment as well. We follow your advice and clarify this in Algorithm 1 caption and highlight the differences in the Algorithm (see revised manuscript).

Dear reviewer, thank you for the reply and we are happy that some of your previous concerns are resolved.

Q: I suppose there might be some other implementations, a straight-forward idea is to use an additional matrix B

Thank you for this suggestion. The idea to use the matrix B would cancel the noise and lead to a deterministic solution (i.e., removing our novel noise component), and thus this is undesirable. For example: If $A = \epsilon I$, then the noise is scaled down via $An$ and the positional embedding $\psi_{B_y}$ is unaffected. B can then be set to be $B = A^{-1} = \frac{1}{\epsilon} I$, and this will preserve the context tokens information.

There might be other ways to regularize A that can be explored, for example, by incorporating additional (multiple) loss terms that ensure A has a large enough norm, and that it is full rank. However, our solution is simpler as it doesn't require additional losses and hyperparam tuning.

Q: I am now a bit confused on the experiments on regularization. I suppose you are trying to prove that the improvements of StoP do not solely come from regularizing $A$? (which is used as the projection matrix for context token)? However, I am not sure if the matrix $A$ in StoP is really regularized towards a sparse matrix. Given that you observed that the norm of $A$ decreases with increasing $\sigma$, I suppose you should try ℓ2 regularization (which regularizes the norm of matrix $A$) instead of ℓ1, and see if that can lead to much improvement.

Indeed, we wanted to show that the improvements of StoP are not just due to reducing the norm of A. Clearly, there are several notions of norm, and these can be explored. We focused on $L_1$ because this is a standard approach to regularizing the rank of $A$ for the diagonal case. Furthermore, it is well known that optimization with SGD implicitly regularizes $L_2$ norm (e.g., see https://arxiv.org/abs/1906.05890), so we wanted to test a norm that is not implicitly regularized.

We note that our $L_1$ regularization experiments also resulted in low $L_2$ (the higher the $L_1$ regularization loss coefficient , the lower the $L_2$ norm, see table below).

The authors may consider using some other regularization (and also remove A in computing context tokens)

I suppose you should try ℓ2 regularization (which regularizes the norm of matrix $A$) instead of ℓ1, and see if that can lead to much improvement.

Dear reviewer, we follow up on your suggestion and include additional experiments applying $L_2$ regularization on $A$.

Specifically, we trained ViT-B/16 baseline models using deterministic sine-cosine positional embeddings for 150 epochs while adding $L_2$ regularization loss weighted by $\alpha \in${$1.0, 0.1, 0.01, 0.001$}. We then applied the ImageNet linear probing protocol, then report the results below.

These results indicate that StoP cannot be merely replaced by $L_2$ regularization over $A$. Please let us know if there are any other concerns, and we are open to hear more feedback or provide further clarification if needed.

Thank you for the thoughtful consideration of the paper and very positive feedback.

Q: To strengthen the evaluation, it would be nice to see linear probes/finetuning results on the larger set of downstream datasets.

Thank you for the comment. We’ve evaluated StoP on 5 different datasets (ImageNet, iNat, Places, DAVIS 2017, CLEVR). Following your comment, we will run additional evaluations on CUB-200, Flowers-102 and IN-100 and will include it in the final manuscript.

Q: it could be nice to have a model pretrained on a larger dataset rather than Imagenet-1000 as it could lead to stronger model and will enable better transfer to downstream problems which is important to have such representations for the community.

Thank you for the suggestion. We think that running large scale experiments (e.g, on LAION 5B) with StoP is exciting. Since this might require non trivial engineering efforts and amounts of resources, we leave this for future work.

Thank you for the thoughtful consideration of the paper and constructive feedback. We’ve incorporated your feedback and uploaded a new revision of the paper.

W1: It is unclear to me why it is necessary to learn optimal Σ via additional parameterization. What is the benefits of introducing additional degree of freedom here to learn Sigma? What if we fix Sigma without learning? Isn't it a simpler way to avoid degeneracy of matrix A? Please explain the motivation.

We compare fixed $\Sigma$ to learned $\Sigma$ in Figure 3. Like you mentioned, using a fixed $\Sigma$ indeed prevents degeneracy of $A$. However, learned $\Sigma$ works better empirically ($+3.5$ compared to $+1.9$, see Figure 3). Implementation wise, both approaches are very simple.

The motivation to use a learned $\Sigma$ is to avoid having to perform an extensive grid search to find the optimal $\Sigma$ values. It is easier to let the model find the values itself.

W2: I understand that adding stochastic perturbation to position of the images makes sense in regularizing the model. However, why the same spectral decomposition is applied to features $s_x$ (by multiplying with $A$)? This step also lacks motivation and seems to be heuristic, please clarify on this point.

Without posing any constraint over $A$, we find that the weights of $A$ are quickly scaled down during training, setting $An=0$ to overcome the noise:

11: $n \sim \mathcal{N}(0, \sigma I)$

12: $m = An + \psi_{B_y} + \tilde{m}$

Therefore, we resort to the basic MIM without stochasticity. To avoid this, we use $A$ to project both the noise tokens $n$ and the context tokens $s_x$:

11: $n \sim \mathcal{N}(0, \sigma I)$

12: $m = An + \psi_{B_y} + \tilde{m}$

13: $c = As_x + b + \psi_{B_x}$

The motivation for using $A$ to project both the context features and noise can be understood by considering two extreme cases, when $A=0$, there is complete certainty about the positional embeddings of the masked tokens $\psi_{B_y}$ but all context is lost ($As_{x}=0$), thus making the MIM prediction task impossible. On the other hand, when $A$ is large the context information $As_{x}$ is preserved, but due to the large magnitude of $A$ the noise is amplified and camouflages the positional embedding features of the masked tokens: $An + \psi_{B_y}$, which makes the prediction task hard as well. This dual role of matrix $A$ forces the model to balance between location certainty and the influence of context features in predictions. It optimizes the trade-off for each feature, balancing their presence in predictions against the need for precise spatial locations.

We discuss this in Section 3.1 (see “Avoiding a degenerate determinism solution”) and in Section 3.2 in the original submission but following the comment we revised the manuscript to make this more clear.

W3: What exactly architecture did the paper use to parameterize the matrix Σ? An architecture flow illustration will help better illustrate this mechanism. (figure 1 does not have this part ) Currently, I am not sure how the back-propagation of Σ flows back to the network and how it affects the SSL learning with a positive gain.

We defined $\Sigma= \sigma AA^t$ (See revised manuscript Eq. 2) where $\sigma$ is a scalar hyperparameter and $A$ is a learned matrix. However, instead of sampling from Eq.1 (where we cannot backprop through $A$), we use the reparametrization trick to sample noise $n \sim \mathcal{N}(0, \sigma I)$, and multiplying by $A$ to get the stochastic positional embeddings: $An + \psi_{B_y}$ (see revised manuscript Eq 3). This is differentiable w.r.t $A$ because the sampling distribution does not depend on $A$. Note that $Cov(An + \psi_{B_y}) = \Sigma$.

We followed your suggestion and revised the architecture figure (Figure 2) to include the reparameterization trick to make it more clear (see new paper revision).

W4: I am not sure of the significance of proposition 1. I do not see why using this optimal predictor can help achieve better generalization ability of the SSL pretraining on downstream tasks.

The main goal of Proposition 1 is to provide insight to what is learned with StoP in a simple setting (one input and one output). In this case, we show that the optimal predictor explicitly models location uncertainty by performing spatial smoothing. We do not claim this property leads to better generalization (although we do see empirical downstream gains).

To summarize, we think Proposition 1 provides a nice further analysis, but we are open to moving this into the appendix.

Dear reviewer, thank you for the reply and we are happy that some of your previous concerns are resolved.

However, in terms of reusing the matrix $A$ to $s_x$, I am still not convinced (W2). The current version of doing this projection lacks clear motivation and thus leaving it hard to judge the correctiveness.

The reason for applying matrix $A$ to $s_x$ is to prevent the stochastic positional embeddings from collapsing into deterministic positional embeddings. Let's begin by describing why this collapse phenomenon happens, and subsequently, we will outline how the use of $A$ with $s_x$ provides an effective solution to address it.

Stochastic positional embeddings collapse to deterministic

By using the reparametrization trick, we generate stochastic positions as follows: $\hat{\psi}_i = An_i + \psi_i$ (Eq 3).

It's important to note that $A$ regulates the noise level, and this noise disrupts the positional embeddings of the masked tokens. Therefore, better MIM predictions may be achieved without the presence of noise. Consequently, during training there is a risk of collapse into deterministic positional embeddings by setting $A=0$.

Our experimental results confirm this. Without introducing a mechanism to prevent collapse, the empirical results resemble those obtained using deterministic Sine-Cosine features (see, for example, Table 6, under "Sine Cosine").

Preventing Collapse

Hence, in order to effectively capture location uncertainty through stochastic positional embeddings, it is crucial to prevent the occurrence of this collapse. While there might exist other ideas to address this issue, we employ a simple yet effective approach. This approach stands out as it doesn't necessitate additional losses, hyperparameters, or even learned weights.

The idea is to use the matrix $A$ to project both $n$ and $s_x$. It's worth noting that in MIM models, there is a linear projection of $s_x$ from the encoder's dimension to the predictor's dimension and we replace it with $A$. For a detailed view of the differences between StoP and MIM, please refer to the revised paper, Algorithm 1.

How does reusing $A$ to both $s_x$ and $n$ prevents collapse while promoting the modeling of location uncertainty?

Using $A$ to project $s_x$ serves as a preventive measure against setting $A=0$, as doing so would eliminate crucial context information, making the MIM prediction task impossible. Or as pointed out by reviewer gTwQ: "StoP effectively prevents $A=0$ as it will lead to $As_x=0$".

However, the model has to learn a matrix A that doesn't excessively amplify $s_x$ as this would result in amplifying the noise $n$ as well. Excessive amplification of the noise would camouflage the positional embeddings of the masked tokens, making their location very uncertain.

Summary

To summarize, without introducing a mechanism to prevent collapse, the positional embeddings become deterministic. We proposed to mitigate that by applying the same matrix $A$ both to the noise $n$ and to $s_x$. To learn a good $A$, the model has to trade off the importance of input context and the certainty in the masked tokens location.

Please let us know if there are any other concerns, and we are open to hear more feedback and provide further clarification if needed. We discuss this topic at length in the recent revision (Section 3.2: “Avoiding a degenerate deterministic solution”).

Thanks for the response! After reading the rebuttal, I think some of my concerns are addressed (empirical evidence showing the benefits of using learned Sigma vs the fixed Sigma). However, in terms of reusing the matrix A to $s_x$, I am still not convinced (W2). The current version of doing this projection lacks clear motivation and thus leaving it hard to judge the correctiveness. In this regard, I am afraid I agree with reviewer gTwQ, and I look forward to a better justification of the formulation.