Sketched Lanczos uncertainty score: a low-memory summary of the Fisher information

• "For the claim that "the disadvantage of introducing noise through sketching is outweighed by a higher-rank approximation" (lines 55~57), the authors should add a simple experiment with synthetic data to verify it."

We thank the reviewer for the valuable comment. We run an experiment on synthetic data and the results (Figure 1 in the main rebuttal pdf) confirms our claim. We now proceed to explain the setting.

For a given number of parameters $p$ and a given ground-truth rank $R$ we generate an artificial Fisher matrix $M$. To do so, we sample $R$ uniformly random orthonormal vectors $v_1 \dots v_R\in\mathbb{R}^p$ and define $M = \sum_i \lambda_i v_i v_i^\top$ for some $\lambda_i>0$. Doing so allows us to (1) implement $x \mapsto M x$ without instantiating $M$ and (2) have explicit access to the exact projection vector product so we can measure both the sketch-induded error and the lowrank-approximation-induced error.

For various values of $r$ and $s$, we run Skeched Lanczos on $x \mapsto M x$ for $r$ iteration with a sketch size $s$, and we obtain the sketched low rank approximation $U_S$.

We generate a set of test Jacobians as random unit vectors conditioned on their projection onto $Span(v_1\dots v_R)$ having norm $\frac{1}{\sqrt2}$. We compute their score both exacltly (as in Equation (9)) and sketched (as in Equation (10)). In the Figure we show the difference between these two values.

As expected, higher rank $r$ leads to lower error, and higher sketch size $s$ leads to lower error. Note that the memory requirement is proportional to the product $rs$, and the figure is in log-log scale.

To further clarify this results, on the x-axis we added the column "inf" which refers to the same experiments done without any sketching (thus essentially measuring Lanczos lowrank-approximation-error only) which coincides with the limit of $s\rightarrow\infty$. The memory requirement of this "inf" setting is $Pr$, that is the same as the second to last column where $s=P$.

• "Experiment settings in this paper are too simple; expect to see results with larger-scale models such as Vision Transformers."

We have extended the experiments with a 4M parameter Vision Transformer on CelebA dataset. Results are great (much better than the ResNet model included in the submission) and we provide figures and tables in the general rebuttal pdf. In the $3p$ budget setting our method clearly outperforms all baselines. In the $10p$ budget Deep Ensemble slightly outperforms in detecting FOOOD101 OoD, but our method still outperform all baselines for the more challenging hold-out-classes OoD datasets.

• "Are there other ways of reducing the memory bottleneck in computing uncertainty scores? Since Eq.2 which authors leverage to compute the score is an approximation of the true variance, there may be better alternatives to it. For instance, is there a more computation/memory-efficient form of the covariance matrix $M$ that leads to better scores with less cost?"

This is a great question and the whole line of research that we reviewed in our related work section aims at answering exactly this question.

There are indeed several works that study memory efficient approximation of the matrix $M$, as referred to in the paper, the most notable are diagonal, block diagonal, and Kronecker factorization. Out of these three, we compare with the first one and outperform it, we didn't include the second one because it requires too much memory and, lastly, we didn't include the third one because it is an architecture-dependent approximation and current implementations are very restricted only to some specific types of network (differently from our model, which is architecture agnostic).

• "For Tab.1, why do the authors fix the memory budget as $3p$? Looking forward to justifications for this choice."

This is an interesting question. Our goal was to show the superiority of SLU in a low-memory regime. Given the number of baselines and ID-OoD dataset pairs, we chose to fix the memory budget in order to display our results more easily.

Then why $3p$ and not $2p$ or $5p$? This was indeed an arbitrary choice. We used both $3p$ and $10p$ (in the Appendix) as we thought this was enough to show the behavior of the method. If the reviewer thinks some different values would be interesting we are happy to run those experiments as well and include them in the paper.

Moreover, we are open to suggestions on experiments without a fixed memory budget and we would be happy to include them.

• "I do not see significant weaknesses with this paper. I think it would be nice if the choice of embedding procedure was better described in the main text. I believe that this might be a problem dependent choice, and whilst the authors have provided a good default choice, it would be good to know when their choice might give poor results, albeit if this seldom occurs."

The embedding is done through multiplication with a sketch matrix. It is true that several choices can be made, an alternative to the srft (used in all the reported experiments) is the dense sketch (Theorem 4, [3] ). You can try this alternative with the code provided in the submission by simply changing the command --sketch srft" to --sketch dense".

The choice of the sketch comes from the theoretical guarantees they offer. Specifically, we require it to be $(1\pm \epsilon)$ oblivious subspace embedding (Definition 2.1) and we require it to be memory efficient. We considered three candidate sketch matrices: SRFT (also named Fast JL transform) [1], Sparse JL transform [2] and Dense JL transform (Theorem 4, [3]). These three sketch matrices give different trade-offs in terms of matrix-vector product evaluation time and memory footprint, which we summarize in the following table. Here, $p$ is the original dimension (i.e., the number of parameters), $\varepsilon$ is the approximation parameter for subspace embedding (as in Definition 2.1) and $s$ is the reduced dimension. Let $\omega$ be such that the current best matrix-multiplication algorithm runs in time $n^\omega$.

From the table above, it is clear that SRFT is best if our goal is to minimize memory footprint. At the same time, evaluation time is still quasilinear.

[1] Ailon, Nir, and Bernard Chazelle. "The fast Johnson–Lindenstrauss transform and approximate nearest neighbors." SIAM Journal on computing 39.1 (2009): 302-322.

[2] Kane, Daniel M., and Jelani Nelson. "Sparser johnson-lindenstrauss transforms." Journal of the ACM (JACM) 61.1 (2014): 1-23.

[3] Woodruff, David P. "Sketching as a tool for numerical linear algebra." Foundations and Trends® in Theoretical Computer Science 10.1–2 (2014): 1-157.

• "Can the authors provide examples of when the embedding method works/fails. I believe it would be greatly beneficial to the reader if some intuition is shed on when this compression step is a useful computational resource and when it may greatly hinder performance."

The embedding is theoretically guaranteed to work (i.e. to have a small error $\epsilon$ with high probability) as long as the sketch size $s$ is big enough. This means that, as long as this condition is satisfied, there exist no examples where the embedding fails.

To be more explicit, for any choosen sketch size $s$, there exists some $\epsilon$ such that the sketching is guaranteed to not exceed $\epsilon$-error (with high probability). This of course means that for this technique to be of any use, such $\epsilon$ needs to be small compared to the ground truth, and consequently $s$ needs to be ``big enough". If $\epsilon$ is on the same scale of the ground truth, then you essentially measure noise, and the method fails.

For example, in the synthetic data experiment (Figure 1 in the general rebuttal pdf) the ground truth is $~0.7$ (i.e. $1/\sqrt2$). You can see that $s=200$ leads to an error of $0.3$ (so likely practically useless), while $s=5000$ leads to an error of $0.07$ (so likely practically useful).

Despite significant investigations, we have been unable to find failure modes beyond those predicted by the theory. Nonetheless, for a fixed sketch matrix $S$, one can construct an adversarial example on which the error is huge. One such example is to build a GGN matrix whose eigenvectors are exactly orthogonal to the span of the sketch in parameter space (which is an easier task with smaller $s$). This failure, however, depends on the specific $S$, so it will not fail when $S$ is randomized.

Finally it should be mentioned that the sketch size $s$ has to be larger than the rank $k$. If $s$ is smaller then we will attempt to orthogonalize more vectors than the dimensionality of the space they live in, which is guaranteed to fail.

• "As I alluded to earlier, the models used in experiments are primarily vision models ranging from 40K params to 300K params, it's not clear whether this approach would work for NLP pre-trained models typically ranging in millions of parameters."

We made additional experiments, specifically, we used a 4M parameter Visual Transformer trained on the CelebA dataset. We also included an extra (easier) OoD dataset: FOOD101. The results are very good and included in the general response pdf. In the $3p$ budget setting our method clearly outperforms all baselines. In the $10p$ budget Deep Ensemble slightly outperforms in detecting FOOOD101 OoD, but our method still outperform all baselines for the more challenging hold-out-classes OoD datasets.

• "As I alluded to The datasets used are also primarily small scale vision datasets like CIFAR and MNIST."

The submission also worked on the CelebA dataset which is not a small-scale vision problem. With the updated experiments, we show again favorable performance on CelebA OoD detection with another and more modern architecture (Visual Transformer), in addition to the previously shown favorable performance on ResNet architecture. This experiments suggests that our method works irrespective of the choosen architecture.

• "Given that memory consumption of the algorithm is $O(k^2 \epsilon^{-2})$. (Line 43) Since this does not depend on the number of parameters. Does this mean this can theoretically be applied to larger models? Also, how does this grow logarithmically in the size of model parameters?"

The actual space consumption is $O(k\varepsilon^{-2} \cdot \log p \cdot \log(k / \delta))$ as reported in Lemma 3.1 and 3.2. In line 43, we say that the memory consumption of our algorithm scales essentially as $O(k^2\varepsilon^{-2})$. In the introduction, we chose to ignore factors $\log p$ and $\log k/\delta$ to make our result more understandable. The word essentially suggests that $O(k^2\epsilon^{-2})$ is not exact but conveys the most relevant dependence. Indeed yes, this means this can theoretically be applied to larger models since the $\log p$ term is neglectable even for a billion parameters model.

• "What is the reason for dip in the AUROC beyond a certain memory cost in Figure 3 and Figure 4?"

In a nutshell, a sketch of size $O(k^2 \epsilon^{-2})$ can only handle up to $k$ vectors and guarantee an absolute error $\epsilon$ on our uncertainty score. If more than $k$ vectors are considered, the sketch becomes more noisy, which makes our uncertainty score more noisy and thus performs worse on OoD task.

• "What is meant by "number of NNs" ? Do you mean number of neural network parameters?"

In our figures we measure the memory footprint of our algorithm as the number of trained models (NNs) one could fit in such space, exactly as the size of a Deep Ensemble would. We believe that such a unit of measurement is natural. Indeed, for large models, we expect to dispose of storage comparable with that used for the model itself but not orders of magnitude more than that.