Paper: Predicting masked tokens in stochastic locations improves masked image modeling (ICLR-2024)

Reviewer comment (Writing):
I am a bit confused on step 11 in Algorithm 1. As in Figure 2, the context and masked representations are computed by adding their tokens and positional embeddings together. Then for step 11, I suppose $\psi_{B_x}$ should refer to the positional embedding, and $A s_x+b$ should refer to context token? Why do we need an additional linear transformation on $s_x$? Some explanations may be needed for this part.

Author response:
**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).

Reviewer comment (Theory):
Based on the above concern, I am also confused by later explanations in section 3.2 and 4.3. 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?

Author response:
**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.

Reviewer comment (Theory):
Moreover, with the above correspondence, we should have $A s_x+b$ (resp. $An+\tilde{m}$) as context (resp.) tokens, then the positional embedding is simply $\psi_{B_x}$ (resp. $\psi_{B_y}$), and where is the stochasticity? I suppose there might be some misunderstanding.

Author response:
**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.

Reviewer comment (Writing):
I would also like to see more discussions on the connection between StoP and vanilla MIM. I suppose we can replace step 10 with $\tilde{m} + \psi_{B_y}$, and step 11 with $s_x+\psi_{B_x}$ to reduce to vanilla MIM, is it correct? Such discussions may make it easier to understand the proposed method.

Author response:
**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).

Reviewer comment (Experiments):
While the authors have mentioned the necessity of regularization on A, the regularization with context token is a bit confusing. I note that the authors have conducted additional experiments in section 4.3 that uses L1 regularization on A. Nevertheless, L1 regularization should aim to obtain a sparse matrix A, which seems to contradict with the original aim to avoid zero A. The authors may consider using some other regularization (and also remove A in computing context tokens) and see how such modification works compared to Algorithm 1.

Author response:
**Q: While the authors have mentioned the necessity of regularization on A, the regularization with context token is a bit confusing. I note that the authors have conducted additional experiments in section 4.3 that uses L1 regularization on A. Nevertheless, L1 regularization should aim to obtain a sparse matrix A, which seems to contradict with the original aim to avoid zero A. The authors may consider using some other regularization (and also remove A in computing context tokens) and see how such modification works compared to Algorithm 1.**

To clarify, in the regularization experiments we followed a similar setting to what you suggested (no stochasticity). We used the basic MIM recipe plus regualrization over the projection matrix that projects the context tokens from the encoder output dimension to the predictor input dimension:

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

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

Since there is no stochasticity, we do not need to worry about avoiding A going to zero to scale the noise down ($An=0$). There is still the tradeoff between the MIM reconstruction loss and regularization loss (but this is always the case with regularization).

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

Author response:
**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.

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

Author response:
**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.

Reviewer comment (Theory):
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.

Author response:
**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.

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

Author response:
**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).

Reviewer comment (Theory):
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.

Author response:
**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.