An In-depth Investigation of Sparse Rate Reduction in Transformer-like Models

It will benefit the paper if the authors can provide more insights and investigation from section 5. For example, examining using SSR to measure other pretrained transformer-like architectures.

Evaluating $R^c(Z;U)=\sum_{k=1}^K \frac{1}{2}\log \operatorname{det}(I+\gamma (U_k^TZ)^T (U_k^TZ))$ in SRR need to specify $U$. However, it is not well defined on models other than CRATE and its variants like general transformers. Moreover, conducting correlation analysis needs to train a large number of models, but currently we lack the resources and time to accomplish it. Despite that, we will provide some preliminary results on evaluating SRR on a standard transformer in Figure 5 in the pdf. Specifically, we set $U=[U_1,\dots,U_K]$ equal to the query, key, and value projections respectively. This shows we may generalize the usage to a standard transformer, or even more general models. We will surely add this experiment and discussion in the revision.

The analysis in Section 4 and experiments in Section 6 uses different alpha values. It could further improve the work if consistent alpha is used.

We will change the Figure 1(a) using consistent $\alpha$ with that in Section 6 in the revision.

W1: Lines 125 and 135-136 make conclusions based on how (6) is related to (7), but these conclusions are not justified.

The derivations in (6-7) are rephrased from the original CRATE paper [1]. The softmax introduced after omitting the first-order term is quite intuitive, which converts auto-correlation to the distribution of membership of different subspaces. We refer the reviewer to Section 2.3 of [1] for discussion.

W1: Note however a softmax is also added from (6) to (7). Maybe omitting the first-order term is not the (only) culprit, adding softmax could also play a role.

Your concern that introducing softmax may also affect the optimization is understandable. But note that, as we mention in line 128, the update we use in the toy experiment is (7), which contains the softmax. Therefore, Figure 1(a) should be able to corroborate our analysis, and adding softmax does not change the conclusion.

W1: Therefore, one step of update (7) secretly maximizes $R^c$. I’m not sure about this either. Note that the green and orange curves are both quadratic functions, and are shifted/scaled versions of each other. Why would doing gradient descent with the orange be better than with the green?

The two answers you mention are reasonable. But we are not arguing which one is better. After all, the goal is to minimize the blue curve$-$the log function. If using the orange curve, we need to constrain the eigenvalues on the left of the peak, which is hard. The update (7) discards the first-order term, corresponding to the green curve. Now, minimizing the green curve makes the eigenvalues $(\ge 1)$ increase, hence maximizing the blue curve which corresponds to $R^c$. The quadratic function is the most straightforward approximation for designing the theory-grounded structure$-$MSSA operator. By optimizing other approximations of log functions, more powerful and more mathematically meaningful structures may emerge.

W1: part of the issues above can be addressed by making a more complete experiment

We follow these settings and provide the visualization in Figure 4 in the pdf. The results shows that the original gradient update of $R^c$ (a) with no approximation will monotonically decrease $R^c$, which is expected. Using the Taylor expansion approximation (b) will increase $R^c$, meaning that the eigenvalues $\gamma$ are mostly greater than one and the optimization happens to the right of the peak of the orange curve in Figure 1(b). Removing the first-order term of the Taylor expansion (d) monotonically increases $R^c$ which is the issue we specify in CRATE-C. Finally, further adding the softmax (e), leading to the update in (7), will not change the monotonic increase of $R^c$. We believe this figure will address most of the concerns in this part.

W2: I do not understand the motivation of the variants proposed in the 4.2. ​​Granted, (7) is not the best approximation to the SRR. How do (8) and (9) make better approximations?

These two variants have different motivations. CRATE-N is motivated to counteract the issue in CRATE-C where the update (7) increases the coding rate $R^c$, conflicting with the SRR principle. By moving in the opposite direction of CRATE-C, CRATE-N with (8) aims to implement the decrease in $R^c$ more faithfully, aligning better with the SRR principle. On the other hand, CRATE-T is not motivated for better alignment with the SRR principle. Instead, it is motivated by designing competitive alternatives of CRATE-C or even CRATE without introducing new parameters, since CRATE replaces the output matrix $U=[U_1,\dots,U_K]$ in CRATE-C with learnable $W$ (line 107) which sacrifices interpretability.

W2: If the argument, as the author states in lines 135-136, is that omitting the first-order terms in (6) is bad, then isn’t it natural to see what if one keeps the first-order terms?

Omitting the first-order term is bad when doing gradient descent because this will maximize $R^c$, but it is good when doing gradient ascent since this minimizes $R^c$ which is (8). It’s indeed natural to keep it, but we have to constrain the eigenvalues on the left to the peak of the orange curve, which is difficult.

W3: Where is the generalization gap $g(\theta)$ defined?

We measure the generalization gap as the difference between validation loss and training loss at convergence (training loss reaches 0.01), i.e., $L_{val}(w)- L_{train}(w)$. We can give a formal definition in the appendix in the revision.

W3: In table 2, the change of the accuracies by adding SRR regularization seems relatively small.

Our goal is not to demonstrate the superiority of SRR over other kinds of regularization or for substantial performance gains. Instead, we would like to validate and complement the conclusions in Section 5. SRR measure is shown to be better than sharpness as a predictor of generalization, then it should be reasonable to incorporate it into the training for improved generalization, similar to sharpness-aware minimization [2]. One direct approach is through regularization. We only provide preliminary results as a proof-of-concept. We believe there is room for engineering to make the results more significant.

[1] Yu et al. White-Box Transformers via Sparse Rate Reduction. NeurIPS, 2023.

[2] Foret, et al. Sharpness-aware Minimization for Efficiently Improving Generalization. ICLR, 2021.

Thanks to the authors for their rebuttal!

Section 4.1/1

I agree that using softmax does not change the conclusion that update (7) is increasing the $R^c$ term (which is counterproductive) as shown in Figure (1).

What I am not sure about is the apriori attribution of this counterproductive phenomenon to solely omitting first-order terms, as the paper writes on lines 135-136. Couldn't it be due to the softmax approximation? Do I miss something here?

Section 4.1/2

I see! So the thing that I was missing was the fact that $\lambda_i^k \geq 1$ (which you mentioned on line 122), so one is always at the right half of the green curve. This could be highlighted better: e.g., you could make the left half of the green curve dashed, and/or reiterate the fact that $\lambda_i^k \geq 1$ near line 140.

Section 4.1/3

Thanks for providing the experiments. They make the picture clear, and therefore I would suggest adding/referring to them in section 4.

Section 4.2

I think reviewers RBPW (W1), uGk9 (W1) and I share the same concern. It seems that CRATE-N and CRATE-T do not connect tightly with the findings in section 4.1.

Perhaps the interesting finding is that of Table 2: Implementation-wise, one has to deviate from CRATE-C and use CRATE for maximized performance, and CRATE-N/-T serve as some sort of performance interpolations between CRATE-C and CRATE. It remains an open problem what the changes (from CRATE-C) do to the representation/optimization dynamics in terms of first principles.

W1:Unclear motivations/soundness of variants

CRATE-N aims to counteract issues in CRATE-C where the update can increase $R^c$, opposing SRR principle. By moving in the opposite way of CRATE-C, CRATE-N implements decrease in $R^c$ more faithfully, aligning better with SRR principle. CRATE-T addresses the discrepancy between CRATE and CRATE-C by questioning the replacement of the output matrix in MSSA operator with learnable parameters in CRATE. By exploring alternatives without new parameters, CRATE-T aims to enhance comprehension of SRR principle and improve performance of CRATE-C.

W1:Authors point to (7) and differentiate it from CRATE...how CRATE-C is different?

Sorry for making you confused. (5-7) describe the approximation and derivation of part of CRATE-C, and (7) is the update in CRATE-C. On the other hand, CRATE involves further modification-replacing matrix $U=[U_1,\dots,U_K]$ with learnable parameters $W$(line 106). The title of Section 4.1 should refers to "CRATE-C", not "CRATE".

W1:"...update step takes the gradient of a lower bound..." weird as the objective is minimizing $R^c$

You are right, minimizing the lower bound of the objective cannot reliably lead to the minimization of the original objective $R^c$. This is exactly the issue we want to highlight for CRATE-C.

W1:"taking gradient ascent...consequently $R^c$", not true for general objective functions

For the general function, you are correct. Yet, we're not using a general function. As in line 122, eigenvalues of a PSD matrix plus an identity matrix are definitely no less than one. When $\lambda \ge 1$, gradient ascent on the second-order term of $R^c$ (green curve in Figure 1(b)) will decrease $\lambda$, hence decreasing $R^c$ (blue curve). Therefore, our statement still holds.

W1:Authors motivate CRATE-T in "replacing output matrix...breaks the structures...", but then still use learnable output matrix, unclear arguments to show

The reason for sacrificing the interpretability is not about "learnable", it is about learnable parameters are "different" for output matrix and Q/K/V matrices. This can greatly enhance the performance. We then want to explore other designs not entirely derived from theory. For example, we have experimented with fixed weights and identity matrix in Table 1 in the pdf. We want to highlight designs different from $[U_1,\dots,U_K]$ in CRATE-C without new parameters.

W2:Increasing SRR after certain depth (Figure2/3), not agree with Figure4 of [1], need discussion.

There are several reasons: 1) Most importantly, different y axis. They plot the figures of compression and sparsity term separately, which fail to reflect the true tendency of compression term because they use $Y^\ell$ in (2). Instead, we visualize SRR measure as a whole and use $Z^\ell$ in (3); 2) They apply $\ell_2$ normalization before measurement of $R^c$ while we follow the definition. 3) CRATE-Base v.s. CRATE-Tiny configurations; 4) ImageNet-1k v.s. CIFAR10/100.

W2:Motivated from pitfalls of the original CRATE, CRATE-N still shows worsening SRR after certain depth.

SRR has two terms: $R^c(Z)$ and $\lambda \|Z\|_0-R(Z)$. Based on analysis in Section 4, the update in (8) can decrease $R^c$. Yet, the operation in (3) that aims to optimize the latter term also matters for the whole objective and may affect $R^c$ as well. This structure could have influences on the rise of the curves in Figure 2 and 3. How the ISTA block in (3) interacts with the MSSA block is an open question.

W2:Limited experiments compared to [1] that used ImageNet-1k.

We provide SRR measure similar to Figure 2 and 3 on ImageNet-1k in Figure 2 in the pdf. The models are trained with the official recipe [1]. We also provide the sparsity term in SRR measure $\|\boldsymbol{Z}\|_0/(d*N)$ in Figure 3 in the pdf. The sparsity term has a similar tendency with the SRR measure, and it goes up in the last few layers, which means that ISTA block also matters and need further research.

W3:Unclear what "lacks fine-grained controls over parameters" means.

What here emphasize the reg term in (10) is an average SRR measure over all layers, and is highly inefficient to compute. The SRR measure at some layer might be more important and should be optimized in isolation. Therefore we propose to regularize one specific layer at each iteration. We'll make this point clear in the revision.

W3:Reg on the last layer, marginal performance gain.

Our goal is NOT to demonstrate its superiority in performance gains. Instead, we want to complement the conclusions in Section 5. SRR measure is shown to be better than sharpness, then it should be reasonable to incorporate it into the training for generalization, similar to sharpness-aware minimization [2]. One direct approach is through regularization. We only provide preliminary results and leaves room for engineering to improve the results.

W3:Why layer 12 in Table2.

Specifying which layer to regularize could be computationally prohibitive, especially when the model size grows. We intuitively select the last layer, which should be reasonable if the depth of the models scales. Results in Table 2 indicate that this intuitive choice can already give consistent performance gains in different settings.

W3:Random layer reg in Table5 performs worse than w/o it, reasons unclear.

We also provide the accuracy on the training set of CIFAR10 in Table 2 in the pdf. The performance drops more on the training set than validation set. So underfitting could be the problem that hurts the generalization.

Q1:How is the generalization gap measured?

It's the difference between validation and training loss at convergence (training loss reaches 0.01).

Limitations. We will include the limitation section in the revision.

[1] Yu et al. White-Box Transformers via Sparse Rate Reduction. [2] Foret, et al. Sharpness-aware Minimization for Efficiently Improving Generalization.

W1: What is the point of proposing CRATE-N and CRATE-T? What problems do they solve?

The goal of developing CRATE-N is to address the potential issue of CRATE-C. As we pointed out that the update (7) in CRATE-C, which performs a gradient descent only for the second-order term, could maximizing $R^c$, rather than minimizing it. This contradicts the principle of SRR. Therefore, CREATE-N is designed by seeking the opposite direction of CRATE-C, which could perform decreases of $R^c$ more faithfully, thus aligning better with the SRR principle.

The goal of developing CRATE-T is to address the misalignment between CRATE and CRATE-C (i.e., the misalignment between CRATE and the SRR principle). As mentioned in line 107, CRATE replaces the output matrix $U=[U_1,\dots,U_K]$ in the MSSA operator with learnable $W$ (which is different from $U$). We then raise the following question on the manipulation of the output matrix: if we are free to adjust the output matrix while sacrificing interpretability, can we find more alternatives that can outperform CRATE-C or even CRATE? Therefore, CRATE-T is a feasible choice without introducing new parameters, which can be utilized to better understand the SRR principle and its connection to the performance. We will revise our Section 4 accordingly to improve clarity.

W1: Conclusions drawn from the comparison with CRATE

We want to clarify that our analysis intends to compare the variants with CRATE-C, not CRATE, because CRATE introduces learnable parameters $W$ that are less interpretable. We believe there are at least some interesting conclusions from the comparison: 1) CRATE-N achieves better performance by following the SRR principle more faithfully, shedding light on the connection of SRR to generalization; 2) We need to explore more design choices (e.g., CRATE-T, which may deviate from directly optimizing the SRR but still exhibit a similar architecture) to gain a complete understanding of the SRR principle for model performance (this motivates our Section 5).

W2: SRR at the last layer close to that at the first layer, especially for the two variants. Authors not suggesting pruning the last layers of CRATE and using the intermediate representations.

The SRR objective contains minimizing two terms: $R^c(Z)$ and $\lambda\|Z\|_0-R(Z)$. Based to analysis in Section 4, the update in (8) can decrease $R^c$. Yet, the operation in (3), i.e., ISTA operation, designed to optimize the latter term also matters for the whole objective. This structure may have effects on the rise of curves in Figure 2 and 3. In fact, the sparsity term $\|Z\|_0$ has been discovered to shoot up in the last layer in Figure 3 of the original CRATE paper [1]. We also make similar discoveries in other variants (see Figure 3 in the pdf).

To see if this rise affects the intermediate presentations, we present the result of linear probing on them in Figure 1 in the pdf. It shows that the representations becomes more linearly separable as the layer goes deeper, although SRR is not well-optimized in the last few layers. This suggests there might be a trade-off between SRR, presumably sparsity, and representation learning in the last few layers. This also means SRR does not necessarily indicate linear separability, especially when the architecture does not faithfully implement the SRR objective.

W3 & Q3: Why is SRR a "complexity measure". Don't see any function class here.

Here we use "complexity measure" as we think higher SRR measure implies more complicated models, then it can be viewed as a "complexity measure" for a model. We agree that the rigorous definition of it should be for a certain function class. We will replace it with "proxy/measure of generalization" in the revision.

W3 & Q4: Worried that the high correlation between SRR and generalization is a special property of CRATE models. Could you also use models that are not CRATE in the experiment of Section 5, and see if SRR can still predict their generalization?

This is a good question. This is also one of the motivations to consider CRATE-T, which is not intentionally designed to minimize or maximize SRR. This variant can expand the candidate classes of transformer-like models in the correlation analysis. We will use more models such as those replacing the output matrix with randomly initialized and fixed weights or just identity matrix. We will also include more general models, such as standard transformers and CNNs for the experiment in Section 5. We have performed the preliminary analysis for transformers (see Figure 5 in the pdf). However, since the correlation analysis requires training a large number of models, we are short of time and computing resources to complete during the rebuttal phase. We will definitely include the experiments and discussions in the revision.

Q1: Cannot see why CRATE-N makes sense. If it is gradient ascent, how does it minimize $R^c$?

CRATE-N actually performs gradient ascent on the second-order term in (5). This corresponds to the green curve in Figure 1(b), which tends to make the eigenvalue $\lambda$ smaller (as the optimization is performed for $\lambda\ge1$ and small $\lambda$ increases the green curve). On the other hand, the right-hand side of (5) itself, which can be understood as the orange curve in Figure 1(b), has an opposite pattern as the green curve, which will decrease when $\lambda$ becomes smaller enough. This leads to the minimization of $R^c$ (which can be expressed with a logarithmic function of eigenvalue $\lambda$).

Q2: In line 170 the authors claimed that SRR "...rises slightly...". But it seems to me that it rises significantly.

We will make our description more accurate in the revision.

Add a limitation section in the revision. We will add a limitation section in the appendix in the revision.

[1] Yu et al. White-Box Transformers via Sparse Rate Reduction. NeurIPS, 2023.