Efficient Sketches for Training Data Attribution and Studying the Loss Landscape
Authors: Andrea Schioppa
Venue: NeurIPS 2024 poster (2024)
Keywords: Deep Learning
Language Models
Interpretability
Training Data Attribution
Loss Landscape
Primary area: interpretability_and_explainability
TL;DR: We introduce new gradient sketching algorithms as a toolkit for model analysis.

Abstract: The study of modern machine learning models often necessitates storing vast quantities of gradients or Hessian vector products (HVPs). Traditional sketching methods struggle to scale under these memory constraints.  We present a novel  framework for scalable gradient and HVP sketching, tailored for modern hardware.  We provide theoretical guarantees and demonstrate the power of our methods in applications like training data attribution, Hessian spectrum analysis, and intrinsic dimension computation for pre-trained language models. Our work sheds new light on the behavior of pre-trained language models, challenging assumptions about their intrinsic dimensionality and Hessian properties.

Efficient Sketches for Training Data Attribution and Studying the Loss Landscape
Andrea Schioppa
Google DeepMind
Amsterdam, the Netherlands arischioppa@googee.cco
Abstract
The study of modern machine learning models often necessitates storing vast quantities of gradients or Hessian vector products (HVPs). Traditional sketching methods struggle to scale under these memory constraints. We present a novel framework for scalable gradient and HVP sketching, tailored for modern hardware. We provide theoretical guarantees and demonstrate the power of our methods in applications like training data attribution, Hessian spectrum analysis, and intrinsic dimension computation for pre-trained language models. Our work sheds new light on the behavior of pre-trained language models, challenging assumptions about their intrinsic dimensionality and Hessian properties.
Introduction
Overview In this work, we investigate gradient and Hessian vector product (HVP) sketching to address the memory constraints inherent in applications that require the storage of numerous such vectors. Examples of such applications include training data attribution (TDA), cigenvalue estimation, and the computation of the intrinsic dimension. While previous studies the intrinsic dimension have employed the Fastfood Transform to mitigate memory demands, we demonstrate its theoretical limitations for sketching and its persistent memory bottlenecks modern accelerators. To resolve these issues, we propose novel sketching algorithms designed for modern hardware and underpinned by robust theoretical guarantees. Our experiments pre-trained language models demonstrate the scalability of our methods while offering new perspectives on the intrinsic dimension and the Hessian of generative language models.
Motivation
Training data attribution (TDA) [20 [1] and Hessian eigenvalue estimation [2] offer powerful insights into neural network behavior. TDA requires storing vectors of the same dimen- sionality N as the network's parameters for each training point, scaling linearly with dataset size (ONN)) for T training points). Similarly, numerically stable Hessian eigenvalue estimation al- gorithms demand repeated Hessian-vector product (HVP) computations and storage, scaling with network size and iterations (O(NT) for T iterations). These memory bottlenecks hinder the study of large-scale models; to address this, sketching (27) provides a compelling solution. By projecting gradients HVPs into liwer-dimensional random subspaces, sketching preserves their essential geometric properties while drastically reducing memory requirements. However, in TDA sketching has been employed with limited effectiveness. Random projections have been carried out with dense matrices, that introduce significant scaling constraints (O(N D) memory for target dimension D) which necessitate the restriction of gradients subset of layers. Current TDA scaling methods require layer selection, as demonstrated with BERT in ( where only 15M parameters were used out of 10M. Beyond the computational cost of dense matrices, our experiments show that layer selection introduces substantial distortion in the estimation of TDA scores (Sec. Recent work on neural network geometry [4] suggests using the Fastfood Transform for gradient sketching
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

with O(N) memory; however, as the Fastfood Transform is random feature generation algorithm, it necessitates indirect application in order to sketch gradients. This has both theoretical and practical consequences. On the one hand, our theoretical analysis reveals that FFD fails to fully satisfy the algorithmic requirements of sketching (Thm. 3.T On the other hand, our experiments (Tab. 3) demonstrate that FFD exhibits unacceptable run-time performance on TPUs. These limitations, both theoretical and practical, underscore the need for novel sketching algorithms. To this end, we investigate new sketching paradigms optimized for modern accelerators like GPUs and TPUs. Our work makes the following contributions:
(abbr. FFD [3]
Scalable Gradient Sketching for Modern Accelerators and Networks: We introduce al- gorithms (AFFD, AFJL, QK) (Sec. 3 designed to overcome performance limitations of existing sketching techniques with modern neural networks on architectures like GPUs and TPUs. Our design analysis provides further insights into the effectiveness of our approach.
2. Robust Theoretical Foundations: We establish theoretical guarantees for AFFD and QK (Thms. 3.2 3.3) for sketching and demonstrate limitations in the theoretical basis of the Fastfood Transform (Thm. B.T.. Our analysis further indicates a dimensionality reduc- tion advantage for AFFD over QK, finding supported by our TDA experimental results
(Sec.5..3)
3. Algorithmic Improvements: We propose more efficient algorithms for the intrinsic dimension estimation and Hessian eigenvalue computation (Sec.
We demonstrate how our methods enable large-scale applications in training data attribution, intrinsic dimension computation, and Hessian spectra analysis with pre-trained language models. This leads to the following insights that advance the understanding of pre-trained language models:
1. Limitations of Layer Selection: We demonstrate that layer selection methods yield inaccurate influence score estimations in training data attribution, their usage should be avoided
(Sec.5.. 5.2).
2. Hlgh Intrinsic Dimension: In contrast to assumptions drawn from classification-based stud- ies, we demonstrate that the intrinsic dimension of LLMs can approach their full parameter count (Sec. (5.4). This challenges prevailing beliefs about the intrinsic dimensionality of these models [I7))
3. LLM Hessian Spectra: Our analysis shows distinct characteristics of LLM Hessian spectra (Sec. 5,5). contrasting with conjectures based on findings related to smaller networks 9 21
Paper organization Sec. 2 provides survey of relevant research in the field, contextualizing our contributions. Sec. intrrouuue our novel sketching algorithms. We begin with necessary background material, analyze design choices, and provide a step-by by-sep implementation tutorial in Appendix B Sec. Joutiness proposed techniques for efficient intrinsic dimension search and Hessian eigenvalue computation. Sec. dddescrbbs our experimental setup: subsections are aligned with Sections and to enhance the connection between theory and empirical results
2 Related Work
Our algorithms, AFFD and AFJL, draw inspiration from the seminal FJL algorithm 2 and the FFD approach (13. However, these techniques were designed before the of modern accelerators like GPUs and TPUs. Therefore, our work revisits their design, optimizing them for modern neural networks and hardware. Theoretically, while (2) established FJL as a sketching algorithm, their proof relies on independence assumptions that do not hold in our setting. To address this, we employ more sophisticated concentration tools tailored bi-linear forms (Thm. (32)) and the special orthogonal group (Thm. 3.3) Recent work on PAC bounds [17 leveraged Kronecker-produce decompositions to accelerate gradient sketching when computing intrinsic dimensionality. Our QK algorithm extends the concepts introduced in [17]. Importantly, we provide a proof that QK is sketching algorithm, absent 17. Additionally, we demonstrate that the Kronecker structure used
Sketching Sketching algorithms have been extensively studied (see surveys [27 18 15)
[7 is not essential for performance gains, highlighting that the true bottleneck in FFD and FJL is memory access. For comprehensive comparison with [17, please refer Appendix C Finally, to support our eigenvalue estimation, we leverage the theoretical guarantees outlined in 225

---

Intrinsic dimension The concept of intrinsic dimension (ID) offers a valuable metric for under- standing the complexity of learning tasks. Originally employed to analyze loss landscapes [[4]. intrinsic dimension has expanded into the study of language models. []] demonstrated its role in explaining the generalization power of fine-tuned pre-trained language models. However, their focus was limited to classification tasks. In contrast, our work extends the analysis of ID to generative tasks. This distinction is crucial as we identify scenarios where the task's intrinsic dimension approaches the full model size, a phenomenon not typically observed in classification settings. Additionally, while FFD has been used for efficient language model fine-tuning [66. memory constraints limited their investigation of the intrinsic dimension (ID) in generative tasks to 500k; in other words, because of scalability contraints, [16 could only work with a target sketching dimension I 500k, which prevented searching for the true value of ID as we demonstrate on summarization task where ID approaches the model dimension. Therefore, our work overcomes this limitation, allowing us to compute the intrinsic dimension of such tasks.
Scaling up influence functions Scaling influence functions for training-date attribution remains crucial research direction. Previous works like 8 7 have improved index creation, retrieval speed, and Hessian estimation. However, these approaches can still be computationally demanding. Our work takes a distinct path, aiming to make influence function calculations fundamentally more efficient. Our HVP sketching methods seamlessly replace existing HVP and gradient computations in frameworks proposed within [8 7 120). Furthermore, we offer eigenvector sketching to enhance methods like the Arnoldi iteration 224] [9 has proposed use dense random projections by materializing them in chunks and fly:he-fly: drawbacks of this approach are: (1) the lack of the scalability in the target sketching dimension and (2) the need of hardware-
custom implementations; on the other hand, our approach removes the requirement for specialized implementations (e.g., custom CUDA kernels). This flexibility enables easy integration into standard ML workflows using higher-level languages such as Jax. An orthogonal direction to scaling up via sketching is that of using surrogate models, compare
Hessian evolution during training Investigating how the Hessian evolves during training has shed light on the dynamics of deep learning, with seminal works [21 9 offering intriguing findings. These studies suggest the progressive disappearance of negative eigenvalues and the confinement of gradient descent within small subspace. While [12] developed a numerically stable Hessian analysis algorithm, its computational demands hinder its application to large-scale models over numerous iterations. Our work addresses this limitation by introducing sketching techniques to enable the efficient construction of large Krylov subspaces (e.g.. 10 -dmmensional) for models like GPT-2L (770M parameters). This advancement significantly surpasses the memory constraints of the method utilized by [2]: a 3TB fp32 storage requirement would have been necessary for a comparable analysis using their approach. Consequently, we are uniquely positioned to rigorously examine the conjectures proposed 21 within the context of pre-trained language models.
3 Design Principles for Efficient Sketching Algorithms
In this section, we explore a design space for more efficient sketching algorithms. To establish foundation, we first analyze the performance bottlenecks of existing algorithms,, specifically FJL and FFD, within the context of modern accelerators. This examination highlights two critical design choices: whether the gradient is sketched implicitly or explicitly, and the kind of pre-conditioner that is used. Informed by this analysis, we propose three novel algorithms: AFFD, AFJL, and QK. We then delve into the theoretical underpinnings of AFFD and QK, providing rigorous proofs for their guarantees. Additionally, demonstrate that FFD lacks the theoretical foundation required for sketching. Relevant experimental findings are presented in Sections 5.2 and 5.3
Dense sketches and FJL. A D-dimensional sketch of the gradient of a real-valued function (0) (0 R is random projection of the gradient VL. N to RDD To ensure this projection preserves essential geometric properties, the random projection operator t : N - TRD must, with high probability, concentrate the norm of (z)) around the norm of æ. Mathematically, for each s and 6 there exists a large enough target dimension ((s, 6) so that for D 2 D)s,
Prob (1))
|| > =||x|12) - 6.
(I)

---

This concept generalizes to sketching higher-order derivatives (see Appendix C J. For our purposes, consider the Hessian vector product operator HVP: RW RN defined as (v) = (()) A sketch of the HVP can be obtained
>> (((HP(("')) v)), e E RD. This sketch defines linear mapping TR: - RDD 25 While simple Dense Sketch (using a random x N Gaussian matrix) ensures the norm property, it has O(DN) memory and O(DN2) compute requirements. The FJL
algorithm [2] addresses the compute cost with a sparser projection matrix:
VD
N G. HN -(x),
4 x
(2)
where: B E IRNN is a diagonal matrix with B = +1 (random signs); HN is the N- dimensional Walsh-Hadamard transform (Appendix C G. is a sparse DxN Gaussian matrix with (in-expectation) @ M) non-zero entries (M is a parameter). The Hy B component preconditions sparse inputs for which ① might otherwise fail. Implementing G. efficiently presents challenges on modern hardware and frameworks like Jax that lack native sparsity support.
FFD: Implicit Gradient sketching The FFD transform, introduced by [3] in the context of kernel machines, provides a computationally efficient way to approximate high-dimensional feature maps. As a random feature generator [1),, FFD constructs high-dimensional random features (E RM) from lower- dimer-dimessionl input vector (u E P). Given u E RP, FFD concatenates N D vectors of the form:
b (u) = op HD G. II. HD B(u) (1sis N )
where B and HD are as in ( (and are both D x D-matrices), II is a permutation matrix, Gy is a diagonal D D-matrix with i.i.d. standard Gaussian entries, and ar is normalization constant (see [33: for a practical implementation see the code released by [16].. FFD has the key advantage of constant memory cost, O(N), regardless of the input dimension D. Since FFD defines map RDD + N direct sketching of a gradient is not possible. To address this, [14 perform what we call an Implicit Gradient Sketch:
S(Ve, L.)
I + FFD(w));
this formulation effectively applies the transpose of FFD to the gradient. While [ establish certain properties of FFD, complete proof of its suitability random feature generator is missing. Additionally, whether the FFD satisfies sketching guarantees ( remains to be fully investigated (we will address it in Thm. b.1b
Explicit sketches In light of equation 4 it's natural to consider whether direct sketch of the gradient could be achieved using a map N + RDD. We define this as an Explicit Gradient Sketch:
(5) This approach offers flexibility. For random feature generation methods like FFD, + needs to be implemented the transpose; for sketching algorithms like FJL, the explicit sketch can be applied directly, while the transpose would be needed for the implicit form (4) In Appendix B we provide a Jax-basee tutorial on transposing sketching algorithms. Which one the right approach? Intuitively, implicit sketches may seem more efficient since they avoid direct gradient materialization. However, as we'll Il demonstrate advantages. Table
S(Von, L) & (Vojo,
Section 5.3 explicit sketches surprisingly offer significant performance 4 quantifies the substantial wall-time reductions (approximately 70% on average) across various algorithms when using explicit sketches.
Removing the lookup bottleneck. In both FJL and FFD algorithms, multiplications by G, and II introduce lookup-barc based memory bottleneck on modern accelerators. This hinders performance, as seen in unacceptable early TPU results (compare Tab. B3 To address this, we propose randomizing the pre-conditioner HN. Efficient implementations of H leverage the fact that it can be decomposed using Kronecker products; specifically, HAR HA @ HB, which allows a recursive multiplication by H in O(N log -immime and O(N) storage. We exploit this by permuting rows/columns within Kronecker factors, reducing memory access costs from O(AB) to O(A + B) in the previous example. For optimal accelerator usage, we limit factors to roughly 1024 columns. Building on this, we
4

---

In the resulting AFFD algorithm 6 we remove II, use row- row-permuted factors in H and column--ermuute factors in H
modify 3
+ RD(y N
GG. B(x)),
(6)
where RD denotes restriction to the first coordinates. While all matrices are now N x N, efficient Hadamard implementations avoid large matrix materializations (see Appendix B for our Jax code).
We further introduce AFJL
7 where Hy is removed:
2(x) = -RolV N G. .H HN..B.B B(x))
This can be seen as replacing G, in FJL with a diagonal Gaussian matrix. Generalizations with multiple (G., H pairs are possible but weren't explored due to the vanilla version's strong results. Empirical results on TPUs (T) show the crucial impact of our changes. AFJL achieves 100x wall-time reduction over FJL, AFFD a 64x reduction over FFD. While GPU wall-time remains similar, peak memory usage significantly improves. Appendix A highlights FJL's scaling issues beyond D 2 on GPUs.
Alternative pre-conditioners. To address the smoothing of sparse inputs in sketching algorithms, (2) introduced H in (2). The theoretical basis lies in the Heisenberg uncertainty principle, leveraging the Fourier Transform on discrete group (see (2). However, the computationally efficient Fast Fourier Transform (FFT) shares this desirable property. This raises the question of whether the FFT might yield performance gains for sketching. We find that replacing HN with the FFT offers significant advantages. Experimentally. it reduces wall time by 62% on GPUs (Tab {.). Inspired by the Kronecker product structure enabling efficient H implementation, we propose another generalized pre-conditioner, Q. This random N x N orthogonal matrix has Kronecker decomposition of K independent orthogonal matrices of sizes {B: x B,) sampled according to the Haar measure on (B,).. Our approach allows direct application of Q without the additional diagonal matrix B, and offers 40% wall time reduction GPUs (Table 4). This Q pre-conditioner has the potential to unlock broader optimizations within sketching algorithm design as discussed in the next subsection.
1=1:
Direct usage of the pre-conditioner Q. Inspired by the design of the pre-conditioner Q, we introduce novel sketching algorithm, QK. This ablation explores the direct use of Q transform the input, potentially leading to improved efficiency and memory usage. QK is defined as:
.(x) N
Q(x).
(8)
Here, Q is random D x N-orthogonal matrix with a Kronecker product decomposition of K independent orthogonal matrices of sizes (D; x B,) Each factor is generated by sampling from S0(B) according to the Haar measure and restricting to the first D; rows. Importantly, QK generalizes the approach of [17 by employing more Kronecker factors. This offers the potential for significant memory reductions (Appendix(C)
Diagrams. Figure
| illustrates AFFD, AFJL and QK with diagrams.
Theoretical results
We now delve into the thcoretical underpinnings of our proposed algorithms, emphasizing the interplay between theory and experimentation. A key finding is a limitation of the FFD algorithm:
Theorem 3.1. There are some inputs x for which FFD does not satisfy the sketching property (
Note that this limitation is specific FFD and not inherent implicit-mode algorithms. This distinction is important: the explicit-mode formulation allowed to simplify the theoretical analysis to prove Thm. 3. Next, establish theoretical guarantee for our AFFD algorithm. Compared to (2) there is added complexity as independence arguments cannot be used for G.; we thus apply the Hanson- Hanson-Wright Wright inequality for quadratic forms prove:
Theorem 3.2. AFFD satisfies (T) with
D
8 = 82 + 2exp
-CE²
(9)
4 1og
2N K,
5

---

Haddsmarrl Srawehhnn
Nadwward traetferr
AFFD
dovaatio vecttr
ITLL drop the tecond Hadaward
Reshase eabrie +
cose, (f acts ooloers, mus)
Figure 1:
Diagram to illustrate our proposed sketching algorithms.
for a universal constant C and for any 81 > O.
Finally, analyze QK. Its structure allows for concentration arguments on orthogonal groups, leading to: Theorem 3.3. QK satisfies with
2=22 Cxp(-1CD,(11+))) 1)3),
for a universal constant
Crucially, the bound in Thm. summation over sub-dimensions, and concentration depending on the D, These theoretical insights support experimental findings in Section 5.3 where QK requires higher target dimensions (D) to achieve performance comparable to AFFD. Because space constraints the proofs are included in Appendix C
3.3 is less favorable than that of Thm. 3.2 This is due to the -rooooo
Expanding the Utility of Sketching Algorithms
In this Section expand the usage of sketching to other applications.
Improving the search for the intrinsic dimension. The intrinsic dimension (Dirr) [14], is the minimum dimension (D) of random subspace (V) where SGD yields least 90% of the full model's performance on a target metric (7.). Sketching algorithms, more efficient than FastFood, already accelerate the search for iint [17]. Our memory-efficient algorithms enable to investigate scenarios where Dint may approach the model dimension. For instance, [16] applied FFD to fine- tune generative language models but capped the target dimension (D) 500k due memory limits. We propose novel search algorithm that estimates Dint in single training run. Current methods rely on multiple runs across potential Dint values. streamline, consider a binary search approach (assuming Dint is a power of 2) starting from initial guess (Dmin) up to the model parameter count (N). This would require [log2 Dmas runs. Instead, we propose single training run where D increases progressively. The heuristic: a fixed computational budget (c steps) should yield an expected improvement of at least 8. We start with D Dmin- If, after steps, the target metric's improvement less than 8 or hasn't reached 730, we double D. This yields an estimate D* within a factor of 2 of Dint. A subsequent run with D D verifies if this factor can be eliminated. See Appendix B for Python code. In Sec. 5.4, we apply this approach to pre-trained language models on classification and generative tasks. We find that Dint < N for classification, but for the generative task. D depends on the choice of T and can approach the parameter count N. This finding significant, as
While 90% is common, this configurable parameter.

---

challenges prevailing assumptions about intrinsic dimensionality in such models [4] 17. These studies primarily focused on classification tasks, obtaining generalization bounds like ((vDint) (αI Eq. 4]. [17 Sec. 3] where Dint was presumed much smaller than the model dimension. Our results suggest a need for non-linear projection algorithms to achieve lower intrinsic dimensionality for generative tasks.
Scaling eigenvalue estimation Investigating the Hessian's spectrum and eigenvectors often relies on memory-boute iterative algorithms. The Arnoldi algorithm, as used [12). requires full re- orthogonalization for stability. This necessitates storing vectors as large as the model itself for each iteration: constraint that limits the number of estimable eigenvalues (e.g., Krylov subspaces of dimension = 90 in [2].. Inspired by the theoretical work of [25]. we propose sketching to address this memory bottleneck. With sketching dimension of 10", we can readily construct Ik-dimensional Krylov subspaces for pre-trained language models like BART, Roberta, and GPT-2L. This represents a breakthrough because the method [12 would demand an impractical 3TB of storage for GPT-2L alone. Our technique enables exploration of conjectures (2 12 ( about Hessian structure in the context of fine-tuning large language models. These conjectures are elaborated in the experimental results on eigenvalue estimation (Sec. 5.5). Crucially, we find that Hessian spectra in these models may deviate significantly from behaviors observed in smaller networks.
5 Experiments
To thoroughly evaluate the proposed sketching methods, we present a comprehensive set of exper- iments. First, we highlight the limitations of existing TDA scaling strategies (Sec. 5.2). Next, we dissect the impact of specific design choices on our sketches (Sec. (53)). We then introduce and validate an algorithm for intrinsic dimension estimation, enabling computational savings (Sec. (.4) and show-casing that the intrinsic dimensionality of generative tasks can be large. Finally, we apply our techniques explore the evolution of the Hessian spectrum during pre- pre-trained language model fine-tuning (Sec. (5.5)
5.1 How we define the Training-Data Attribution score.
In TDA there are different ways to measure the similarity score between two examples x and y. In our experiments opt for the TDA score defined as VoT.(0, Vo =) because of its simplicity, allowing us to iterate on multiple algorithms and layer selection schemes, and being a building block for more complicated methods. While high correlation with full gradient dot products may not be the definitive measure of long-term TDA performance time range and building block of more computationally intensive methods like TRAK (1). For example, in the short time range, gradient dot products correlate with loss changes and are relevant to select examples for error correction 233. Evaluating with the LDS from TRAK would introduce more hyper-pacc parameters and considerable more computation: one would need at least 50 models fine-tuned on different subsets of the data; moreover, TRAK itself relies on accurate gradient sketches, as measured by dot products, as the basic building block as TRAK demonstrates that "preserving inner products to sufficient accuracy results in gradient-descent system that approximately preserves the same evolution as the one corresponding to model re-training" [19. C.2).
19 233, it is a practical metric in the short
5.2 Shortcomings of previous Training-Date Attribution scaling strategies.
Past work 20 28
tackles the memory bottleneck of Training- Attribution by calculating
gradients restricted specific layer and potentially applying Dense Sketch. Here, we demonstrate that layer selection distorts both influence scores and eigenvalue estimates, while Dense Sketches exhibit poor scaling characteristics. These findings align elosely with the first subsection of Sec. 3 Furthermore, advice on layer selection lacks consistency: (20 promotes the last layer, whereas [28] supports using Token Embeddings in NLP tasks. We demonstrate the distortion caused by layer selection on influence scores (the inner product of two gradients). Considering 2 pairs of points (x., z), we compute the Pearson correlation r between the TDA score VoL xr) Lo, estimated using a layer-specific gradient and the ground truth based on the full gradient. We adopt the setup of [6] generative task fine-tuning GPT-2 on the WikiText- 103 dataset (BART and zsRE results in Appendix A). Our findings indicate the unreliability of layer selection (Table I. Correlations
7

---

Table Layer selection results in unreliable estimates for influence scores and eigenvalue estimation. The best correlation with ground truth influence scores does not exceed 90% and is quite low for most layers; the relative emror in eigenvalue prediction is always least 20%.
MODEL
LAYER
R
EIG. ERR.
GPT-2
TOK. EMB.
0.16
0.72
GPT-2
1
0.75
0.24
GPT-2
2
0.89
0.31
GPT-2
3
0.90
0.19
GPT-2
0.89
0.24
GPT-2
5
0.78
0.37
GPT-2
0.38
0.40
Table 2: Dense projections on the layers do not scale; for each layer we report the wall time for the maximum dimension that does not result in OOM.
LAYER
GPU WALL (MS)
k (MAX VAL BEFORE OOM).
TOK. EMB.
31.2
32
LAYER
25.1
128
LAYER 2
24.4
128
LAYER 3
23.9
128
LAYER
23.1
128
LAYER
22.4
128
LAYER
21.7
128
with ground truth rarely exceed 90% and are significantly lower for most layers. In contrast, AFJL achieves 98% correlation with compact D 213 Extending the analysis to Hessian-based influence functions by looking into eigenvalue estimation emphasizes the shortcomings of layer selection. We aim to compute the top 10 eigenvalues of both the full Hessian and those restricted to chosen layer. Even with potential differences in magnitude and location, layer selection could still be valuable if the true eigenvalues were approximated well by applying an orientation- perreerrinn-rssrviing linear map R² + RR to those computed for particular layer. However, this is not the case, with the relative mean absolute error at 20% on the best layer (Table Finally, layer selection coupled with dense projections faces severe scalability limitations. Setting D 21 within the same setup highlights this issue. Limited memory forces us to divide the computation into D/k projections to RR* where Out-of-Memory error. For token
is the smallest power of 2 enabling dense projection without embeddings, we find & = 32, whereas other layers require k = 128 on a V100 GPU. Projecting to R takes 31ms for token embeddings, resulting in total of ~ 4s to project R., Other layers require ~ 24ms per projection to R$ and ~ 0.77s to project to R., see Table for per-layer breakdown. Finally, alternative approach to dense projections is materialize dense matrices on the fly in chunks [19] this has two substantial disadvantages: (1) the runtime scales linearly with the target dimension (memory is traded off with compute), and (2) specialized kernels necessary for efficient implementation, with unclear applicability to TPUs; we include a demonstration these limitations Appendix A
5.3 Analyzing the Impact of Design Choices
In this section, we analyze the impact of the design choices presented in Sec. 3 both sketching quality and performance. Regarding sketching quality, we find that small values of D often lead remarkably accurate reconstructions of influence scores. For TDA scores, achieving correlation
r 2 0.95 requires the following dimensions: FJL. FFD, AFFD: D = 200, AFJL: D 212; QK: D 244 To reach 2 0.99, simply increase each dimension by factor of While memory limitations prevented sketching HVPs with FJL and FFD for eigenvalue estimation, the other algorithms scaled well. We achieved relative mean absolute error below 5% with the following dimensions: AFFD: D = 240, AFJL: D 222 QK: D 213 Note that QK requires larger D, consistent with the theoretical comparisons Theorems 3.3 and 3.2 (Sec. Regarding perfomance, outline here the key findings and refer the reader to Appendix A for comprehensive analysis
3.3
8

---

Table 3: Wall-time T and peak memory usage M comparison on gradient sketches for GPT-2.
Removing look-ups is crucial for TPU performance and decreasing GPU memory utilization.
AT.GO
GPU (VT00)
TPU (V2)
(MS)
M (GB)
T (MS)
M (GB)
FJL
123
6.9
8997
3.0
AFFD
205
3.2
134
2.8
FFD
197
4.1
8694
AFJL
116
2.9
89
2.7
QK
82
1.7
64
1.1
Table Speed-ups (ratio R. of the slowest wall-time the fastest one) corresponding changing design choice (c.g. implicit to explicit or HN to the FFT.).
ALGO
GPU (V100) (R)
TPU (v2) (R)
ALGO
GPU (VI00) (R)
TPU (v2) (R)
IMPLICIT - EXPLICIT
HN
FFT
FJL
1.96
1.20
AFFD
1.98
1.00
AFFD
1.81
2.07
AFJL
1.64
1.00
FFD
1.76
1.10
HN
AFJL
1.64
1.67
AFFD
1.45
1.00
QK
1.41
2.45
AFJL
1.44
1.00
across design choices. First, removing look-ups significantly reduces wall-time on TPUs, while on GPUs it substantially lowers peak memory usage (reductions of 2.4x for FJL to AFJL, and 1.3x for FFD to AFFD), , see Table Second, explicit sketching consistently provides substantial conjecture this is due to the fact that implicit sketching results in more memory accesses. Finally while modifying the pre-conditioner doesn't 't affect TPU performance, significantly improves GPU performance, see Table under the headings HN FFT and H - Q.
speed-ups (see Table
5.4 Estimating the intrinsic dimension.
Our experiments evaluate the efficiency and accuracy of our intrinsic dimension estimation algorithm (presented Sec. ( We consider two experimental setups: classification, where we fine-tune Roberta on SNLI with accuracy the target metric; generation, where we fine-tune BART on XSUM for text summarization, using Rougel and Rouge2 for evaluation. We employ three projection algorithms: FFD, AFJL (FFT-based variant), and QK. We first validate algorithm consistency by demonstrate that the algorithm produces consistent estimates of the intrinsic dimension across multiple runs; repeat each search experiment with three random seeds obtaining estimates within factor of (Appendix B). We then verify that the estimated D* accurately represents the intrinsic dimension Dint by fine-tuning the model in D D /0-/22iimesss 12-dimensional subspace and ensuring that the target metric remains below the Tyo threshold (details Appendix BB. We point out that selecting good search parameters 5c was relatively straightforward by observing target metric improvement during full fine-tuning. Regarding computational efficiency, our approach requires significantly fewer steps than binary brute-force search. For instance, the XSUM search required only twice the steps of a full fine-tuning run and the same amount of steps for the classification task. We finally look how the intrinsic dimension varies between classification and generation tasks. For classification, our findings are consistent with prior work [1]. were all algorithms yield metric-dependent: for 1 FFD and QK estimate D 225 while AFJL yields D = 224, for Rouge2, Dint equals the full model dimension N 227 both cases, however, the intrinsic dimension is close the full model dimension: this challenges assumptions about intrinsic dimension in the context of generative tasks. While prior studies [7 focused on classification with generalization bounds the form O(vDint), our results indicate that: generative tasks may exhibit higher intrinsic dimension and generative models may require new non-linear projections to uncover significantly lower intrinsic dimensions.
Dint 233 N. For generation we first observe that the results

---

0.4
64
S002
(a)
(b)
Figure 2: left: ratio (RNEG) the absolute value of the top negative to the top positive eigenvalue; right: ratio R of the n-th largest positive cigenvalue the largest positive cigenvalue. We define outliers when > 20%, motivated by [12] Fig.2]. Higher-resolution versions for printing can be found in Appendix A These results disprove conjectures on the Hessian structure, see Sec. 5.5
5.5 Hessian Analysis During Pre-trained Language Model Fine-Tuning
In this section, examine the Hessian's 's evolution during pre-trained language model fine-tuning. Using sketching, as proposed in Section 4 we estimate the Hessian's eigenvalues and eigenvec- tors. We fine-tune Roberta and GPT-2L on SNLI classes) and BART on XSUM, employing AFFD, with target dimension D = 220 to construct Ik-dimensional Krylov subspace. This substantially extends the work of [2] who considered smaller models and smaller subspaces (90- dimensional Krylov subspace). Spectrum estimation is performed every 1k steps, including initial steps (0, 10, 50, 100, 150, 200, 250,500 500). Our goal is to see if observations from previous studies with smaller networks hold in this context:
Obs1: Negative eigenvalues gradually disappear during training 21 12
outlier eigenvalues emerge for K-class classification, with the gradient aligning to their corresponding subspace (9] 21 3.. Moreover, these outliers, which hinder optimization, stem from training without normalization [2]
Obs2: K
We find that these obervations don't 't fully translate to pre-trained language model fine-tuning. Re- garding Obs1, we compute the ratio (RNEG) of the absolute values of the top negative and positive eigenvalues; RNEG shows inconsistent behavior across models (Fig 2 (a)) Roberta maintains a higher RNEG, while GPT-2L and BART see diminish over time. Regarding Obs2, outlier cigenvalues don't strictly adhere the K 1 rule. Roberta has more outliers (6), and the gradient-outier alignment is less pronounced (27% Roberta, 35% GPT-2L, 8% BART) compared to smaller networks [22 9]. See Fig. 2 (b. Moreover, outliers emerge despite layer normalization.
6 Conclusions and Limitations
In this work, we have dissected the theoretical and practical limitations of existing gradient sketching techniques when applied to modern neural networks and accelerators. Our analysis motivated the design of novel sketching algorithms, for which we established theoretical guarantees; additionally, we exposed limitations the theoretical underpinnings of the Fastfood transform. These methods, along with refined intrinsic dimension estimation and Hessian eigenvalue computation, provide an efficient toolkit for model analysis. We successfully apply this toolkit to pre-trained language models, revealing the need to rethink layer-selection based influence functions, the high intrinsic dimensionality generative task, and the deviation LLMs' Hessian spectra from what observed in smaller networks. While in Sec.5.4) we exhibit an example of a generative task with a large intrinsic dimension, we leave an in-depth study for future work. We tested the efficiency of our sketching algorithms with Transformers in Sec. 5.3 but results might vary for other model architectures.
10

---

Acknowledgements
We thank Jonathan Heek for helping with optimizing code for TPUs and Katja Filippova for feedback on the draft.
References
[1] Armen Aghajanyan, Sonal Gupta, and Luke Zettlemoyer. Intrinsic dimensionality explains the effectiveness of language model Fine-Tuning. In Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference Natural Language Processing (Volume I: Long Pupers), pages 7319-7328, Online, August 2021. Association for Computational Linguistics.
[2] Nir Ailon and Bernard Chazelle. nearest neighbors. SIAM Journal on Computing, 2009.
The fast caanssonn-indensrauss transform and approximate
[3] Gerard Ben Arous, Reza Gheissari, Jiaoyang Huang, and Aukosh Jagannath. High-dimensional sgd aligns with emerging outlier eigenspaces, 2023.
[4] Krzysztof Choromanski and Francois Fagan. Fast nonlinear embeddings via structured matrices, 2016.
[5] Andrew Engel, Zhichao Wang, Natalie S. Frank, Ioana Dumitriu, Sutanay Choudhury, Anand Sarwate, and Tony Chiang. Faithful and efficient explanations for neural networks via neural tangent kernel surrogate models, 2024 (also accepted in ICLR24).
[6] Jillian Fisher, Lang Liu, Krishna Pillutla, Yejin Choi, and Zaid Harchaoui. Influence diagnostics under seff-concordance. Francisco Ruiz, Jennifer Dy, and Lan-Willem van de Meent, editors, Proceedings of The 26th International Conference Artificial Intelligence and Statistics, volume 206 Proceedings of Machine Learning Research, pages 10028-10076. PMLR, 2023.
[7] Roger Grosse, Juhan Bac, Cem Anil, Nelson Elhage, Alex Tamkin, Amirhossein Tajdini, Benoit Steiner, Dustin Li, Esin Durmus, Ethan Perez, Evan Hubinger, Kamilé Lukošiūté, Karina Nguyen, Nicholas Joseph, Sam McCandlish, Jared Kaplan, and Samuel R. Bowman. Studying large language model generalization with influence functions, 2023.
[8] Han Guo, Nazneen Rajani, Peter Hase, Mohit Bansal, and Caiming Xiong. FastIF: Scalable influence functions for efficient model interpretation and debugging. In Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, pages 10333-10350, Online and Punta Cana, Dominican Republic, November 2021. Association for Computational Linguistics.
[9] Guy Gur-Ari, Daniel A. Roberts, and Ethan Dyer. Gradient descent happens in tiny subspace, 2018.
[10] William B Johnson and Joram Lindenstrauss, Extensions of lipschitz mappings into a hilbert space, 1984.
[11] Pang Wei Koh and Percy Liang. Understanding black-box predictions via influence functions. In Doina Precup and Yee Whye Teh, editors, Proceedings of the 34th International Conference Machine Learning, volume 70 of Proceedings of Machine Learning Research, pages 1885-1894. PMLR, 06-11 Aug 2017.
[12] Shankar Krishnan, Behrooz Ghorbani, and Xiao Ying. investigation into neural net opti- mization via hessian eigenvalue density. In ICML, 2019.
[13] Quoc Le, Tamas Sarlos, and Alex Smola. Fastfood approximating kernel expansions in loglinear time. In 30th International Conference on Machine Learning (ICMI., 2013.
[14] Chunyuan Li, Heerad Farkhoor, Rosanne Liu, and Jason Yosinski. Measuring the intrinsic dimension of objective landscapes. In International Conference on Learning Representations, 2018.
11

---

[15] Fanghui Liu, Xiaolin Huang, Yudong Chen, and Johan A K Suykens, Random features for kernel approximation: A survey on algorithms, theory, and beyond. IEEE Transactions on Puttern Analysis and Machine Intelligence, April 2020.
[16] Haokun Liu, Derek Tam, Mugeeth Mohammed, Jay Mohta, Tenghao Huang, Mohit Bansal, and Colin Raffel. Few-shot eeffettr-efficient fin-tuning is better and cheaper than in-context learning. In Advances in Neural Information Processing Systems, 2022.
[17] Sanae Lotfi, Marc Anton Finzi, Sanyam Kapoor, Andres Potapczynski, Micah Goldblum, and Andrew Gordon Wilson. PAC-bayes compression bounds so tight that they can explain generalization. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, 2022.
[18] Michael W. Mahoney. Lecture notes on randomized linear algebra. Arxiv, 2016.
[19]
Sung Min Park, Kristian Georgiev, Andrew Ilyas, Guillaume Leclerc, and Aleksander Madry, Trak: Attributing model behavior at scale, 2023 (also accepted in ICMI.23).
[20] Garima Pruthi, Frederick Liu, Satyen Kale, and Mukund Sundararajan. Estimating training data influence by tracing gradient descent. In Hugo Larochelle, Marc' Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan-Tier Lin, editors, Advances in Neural Information Processing Systems 33: Annual Conference Neural Information Processing Systems 2020, NeurlPS 2020, December 6-12, 2020, virtual, 2020.
[21] Levent Sagun, Utku Evci, V. Ugur Guney, Yann Dauphin, and Leon Bottou. Empirical analysis of the hessian of over-parametrized neural networks, 2018.
[22] Tamas Sarlos. Improved approximation algorithms for large matrices via random projections. In 2006 47th Annual IEEE Symposium on Foundations of Computer Science (FOCS pages 143-152, 2006.
[23] Andrea Schioppa, Katja Filippova, Ivan Titov, Polina Zablotskaia. Theoretical and practical perspectives what influence functions do, 2023 (also accepted in Neurips23).
[24] Andrea Schioppa, Polina Zablotskaia, David Vilar Torres, and Artem Sokolov. Scaling up influence functions. In AAAI-22, 2022.
[25] William Swartworth and David P Woodruff. Optimal eigenvalue approximation via sketching. STOC 2023: Proceedings of the 55th Annual ACM Symposium on Theory of Computing, June 2023.
[26] Roman Vershynin. High-Dimensional Probability. Cambridge University Press, 2018.
[27] David P Woodruff. Sketching as tool for numerical linear algebra, November 2014.
(28) Chih-Kuan Yeh, Ankur Taly, Mukund Sundararajan, Frederick Liu, and Pradeep Ravikumar. First is better than last for language data influence. In Advances Neural Information Process- ing Systems 36: Annual Conference Neural Information Processing Systems 2022, NeurlPS
2022, February 2022.
12

---

Contents
Appendix: Additional Experimental Results
13
A.1
Additional results on layer selection
13
A.2
Quality of sketches for inner products and eigenvalue estimation
13
A.3
A closer look at FJL vs AFJL
14
A.4
Compute cost of sketching gradients
16
A.5
Compute cost of sketching HVPs
18
A.6
Search for the intrinsic dimension
20
Higher resolution version of Fig.
22
A.8
Comparison to on-the-fly dense random projections
22
Appendix: Implementation details
23
B.
Libraries, Compute resources and implementation of FFD and FJL.
23
B
HVP in Implicit vs Explicit Form
24
B.3
Step-by-step didactic implementation
24
B.4
Sketching and model parallelism)
28
B.5
Algorithm for searching the intrinsic dimension
30
B.6
Hyper-parameters for Sections 5. and
31
B.7
Hyper-parameters for Sec. 5.4
B.8
Hyper-parameters for Sec. 5.5
Appendix: Theory
31
Definition of Higher order sketches
Guarantees on distorting distances.
31
Definition of the Halsh-Hadamard transform.
31
Failure of concentration for FFD.
32
C.5
Comparison to Kronecker products in
32
C.6
Concentration result for QK: Theorem
3.3
32
Concentration result for AFFD
35
A
Appendix: Additional Experimental Results
A.1 Additional results on layer selection
Table we include the full results of layer selection GPT-2 and BART: in Sec, 5.2 we restricted the discussion to GPT-2 because of space constraints; for the purpose of this experiment we consider the setup of [6] NL.P tasks which consists fine-tuning GPT-2 on the WikiText-103 dataset and BART and zsRE.
A.2 Quality of sketches for inner products and eigenvalue estimation
For each algorithm, we report in The minimal value loge D necessary to reach Pearson correlation > 0.9r with the ground truth when estimating inner products of gradients using sketches. As expected from the worse concentration bound, QK requires larger dimension. We conjecture that the fact that FFD is effective might be due the gradient distribution giving O-measure to the inputs
13

---

Table 5: Most of the time layer selection results in unreliable estimates for influence scores and eigenvalue estimation. The best layer correlation with ground truth influence scores does not exceed 90% and is quite low for most layers. The relative error in eigenvalue prediction is always least
20%.
MODEL
LAYER
R
FIG. ERR.
GPT-2
TOK. EMB.
0.16
0.72
GPT-2
0.75
0.24
GPT-2
2
0.89
0.31
GPT-2
3
0.90
0.19
GPT-2
4
0.89
0.24
GPT-2
0.78
0.37
GPT-2
6
0.38
0.40
BART
TOK. EMB.
0.54
0.20
BART
DEC.
0.62
0.39
BART
DEC. 2
0.88
0.43
BART
DEC. 3
0.73
0.28
BART
DEC. 4
0.91
0.28
BART
DEC. 5
0.84
0.19
BART
DEC.
0.70
0.18
BART
ENC.
0.41
0.50
BART
ENC. 2
0.45
0.59
BART
ENC. 3
0.56
0.48
BART
ENC. 4
0.71
0.40
BART
ENC.
0.91
0.46
BART
ENC. 6
0.89
0.16
Table For each algorithm the minimal value log, D necessary to reach Pearson 7 > æ where 0.9(5, 9) for estimating inner products of gradients.
ALGO
0.95
0.98
0.99
FJL
12
AFFD
10
12
13
AFJL
12
13
15
QK
14
16
FED
12
that FFD would fail to sketch (Theorem 3.1 For AFFD, AFJL and QK we report Table 7 the minimal value of log, D necessary to reach relative mean absolute error err < when estimating top 10 eigenvalues of Hessian.
A.3
closer look at FJL vs AFJL
In Sec. 5.3 we pointed out that the PJ''s wall time and memory usage increase with the target dimenston We compare the peak memory usage and the wall time of FJL with that of AFJL on the inner product task of Sec. (.3 GPU (V (V100), see Figures FJL's cost significantly increases with and does not scale beyond 220
Table For each algorithm the minimal value of log, D necessary reach relative error err where = 0.2, 0.1, 0.05 in reconstructing the top 10 eigenvalues.
ALGO
err < 0.2
< 0.1
0.05
AFFD
10
10
AFJI
12
12
QK
13
13
13
14

---

12
FJL
Peak memory (GB)
AFJL
10
AFJL + FFT
8
6
4
2
8
10
12
14
16
18
20
22
D (log2)
Figure 3: Peak memory usage comparing FJL with AFJI. Results on GPU (V100); for FJI results
with D > 220 are not reported as there were cettof--eeory enrors.
FJL
160
AFJL
Wall time (ms)
140
AFJL + FFT
120
100
80
8
10
12
14
16
18
20
22
D (log2)
Figure 4: Wall time comparing FJL with AFJL. Results on GPU (V100); for FJLL results with D > 220 are not reported as there were oft-of-Memory errors.
15

---

A.4 Compute cost of sketching gradients
In Table we report the compute costs of sketching gradients the setup of Sec.
16

---

Table 8: Compute costs of sketching gradients in the setup of Sec. 5.3 I=1 denotes that a method is implicit, F=1 that the FFT used pre-conditioner instead of the Walsh-Hadamant Transform, and Q=1 that random orthogonal matrices decomposed as Kronecker product are used as a pre- conditioner The lowest time and memory costs are in blue and bold and the highest ones in red and italic. nan indicates result that was not computed because of XLA compilation errors. For this benchmark we considered target dimension D the range [2 222). GPU V100, TPU TPUv2.
ALGO
GPU WALL (MS)
GPU MEM (GB)
TPU WALL (MS)
TPU MEM (GB)
FJI.
0
0
0
123
6.9
8997
3.0
FJL
0
1
0
85
6.1
9271
4.4
FJL
0
0
244
6.7
10958
3.3
FJL
164
6.1
NAN
NAN
AFED
0
205
3.2
134
2.8
AFFD
O
O
1
142
3.2
134
2.8
AFED
O
0
104
4.2
134
2.8
AFFD
O
454
3.2
278
3.0
AFFD
237
3.2
278
3.0
AFFD
126
4.2
278
3.0
FFD
O
197
4.1
8694
4.3
FED
140
5.6
NAN
NAN
FED
352
4.2
9600
4.7
FFD
239
4.8
NAN
NAN
AFJI
0
116
2.9
89
2.7
AFJL
81
2.9
89
2.7
AFJL
0
3,8
89
2.7
AFJL
3,0
149
2.8
AFJL
0
121
3.1
149
2.8
AFJL
0
88
4.1
149
2.8
QK
0
0
82
1.7
64
1.1
QK
116
1.7
157
1.4
17

---

A.5 Compute cost of sketching HVPs
In Table we report the compute costs of sketching HVPs in the setup of Sec.
18

---

Table 9: Compute costs of HVP using sketching. I=l denotes that a method is implicit, F=1 that the Fourier Transform is used instead of the Walsh-Hadamard Transform and Q=1 that random orthogonal matrices decomposed as a Kronecker product are used instead of Hadamard matrices. The lowest time and memory costs are in blue and bold and the highest ones in red and italic. GPU is V100, TPU is TPUv2.
ALGO
Q
GPU WALL (MS)
GPU MEM (GB)
TPU WALL (MS)
TPU MEM (GB)
AFFD
0
0
0
397
2.8
158
2.7
AFFD
0
0
183
2.8
158
2.7
AFFD
111
3.3
158
2.7
AFFD
O
0
627
3.1
268
3.6
AFFD
O
1
238
3.1
267
3.6
AFED
O
137
3.7
268
3.0
AFJL
230
2.7
108
2.6
AFJL
121
2.7
108
2.6
AFJI.
O
0
92
3.3
108
2.6
AFJI
0
0
345
3.0
191
3.5
AFJL
150
3,0
190
3.5
AFJL
111
191
3.5
QK
0
0
118
1.6
82
1.3
QK
0
146
1.9
161
1.7
19

---

A.6 Search for the intrinsic dimension
In Table 10 we report the values of D across the seeds used in each search experiment. We see good agreement within a factor of 2. For SNLI the target accuracy to exceed was T90 = 80.1% for XSUM the value of Rougel exceed was 36.6 while that of Rouge2 was 15.69. In the case of Rouge2 compressing BART in half would lead to a search with D = 226; in that case the final values of Rouge2 did exceed 14.4, so stayed well-below the required threshold that defines the intrinsic dimension Dint using 90% target value of the metric obtained by fine-tuning the full model.
20

---

Table 10: Values of D' returned by the search the intrinsic dimension Dins using different seeds. This shows the stability of our algorithm which doubles the dimension the fine-tuning subspace after some compute budget if the target metric has not improved enough.
TASK
METRIC
PROJECTION ALGORITHM
D
SNLI
ACCURACY
FFD
2 (14,13113)
SNLI
ACCURACY
AFJL
144,1,,33)
SNLI
ACCURACY
QK
213,13,133
XSUM
ROUGE1
FFD
2 262,26,225
XSUM
ROUGEL
AFJL
(24,24,25)
XSUM
ROUGE1
QK
2 (28,2,266)
21

---

BARTIXSUMI
ROBERTA(SNL)
3.0
RNEG
2.0
1.5
1.0
0.5
0.0
2000
4000
6000
8000
10000
12030
14000
16000
18000
20000
Step
Figure ratio (RNEG) of the absolute value of the top negative the top positive eigenvalue
1.0
ROBERTAISNLI)
BARI (XSUM)
GPT-2L(SNLI)
0.8
0.6
0.4
0.2
0.0
20
80
100
Figure 6: ratio R of the n-th largest positive eigenvalue to the largest positive eigenvalue. We define outliers when R 20%, motivated by [2 Fig.2]
A.7
Higher resolution version of Fig.
Figs. 5 and 6 are high-resolution versions of Fig. 2
A.8
Comparison to on-the-fly dense random projections
[19 proposes to materialize the full random projection on-the-fly in chunks (we will refer to this sketching method as TRAK, even though the full TRAK attribution method in [19] is more involved). This leads to two significant drawbacks: (1) the run-time scales linearly with the target dimension (memory traded off with compute). and (2) specialized kernels are necessary for efficient implementa- tion, with unclear applicability to TPUs. Our attempts to implement a competitive version using pure JAX were unsuccessful due to the lack of control over memory allocation and placement. We have included a plot (Figure 7 and table (Table [1) demonstrating the linear runtime growth of TRAK compared to the constant runtimes of AFFD and QK. The implementation challenges of TRAK are further highlighted by the fact that our Triton implementation does not outperform the original CUDA
22

---

16000
TRAK[CUDA]
14000
TRAK[Triton]
wall time (ms)
12000
AFFD[ours]
10000
QK[ours]
8000
6000
4000
2000
0
12
13
14
15
16
17
18
D (log2)
Figure 7: Our methods exhibit constant wall time with respect the target dimension D. In contrast, TRAK's runtime increases with the target dimension. Efficient implementation of dense random projections with recomputed projectors is non-trivial; compare the performance difference between TRAK/CUDA] and TRAK[Triton]. TRAK[CUDA] utilizes the CUDA kernel provided by the original TRAK authors 19
Table Wall-time (ms) Comparison: Our Methods vs. On-the-fly Dense Projections (TRAK) using a V100 GPU. TRAK requires custom kernels and is thus restricted to GPU computation. Our methods exhibit constant runtime with respect to the target dimension, whereas TRAK's runtime increases substantially as the target dimension grows.
loga D)
TRAK/CUDA]
TRAK[TRITON]
AFFD[OURS]
QK[OURS]
12
359
879
213
84
13
594
2885
207
87
14
1078
6073
207
91
15
2047
12900
215
82
16
3969
219
82
17
7812
210
91
18
15500
209
90
kernel released by the TRAK authors, underscoring the difficulty of efficiently implementing random projections in chunks.
B Appendix: Implementation details
B.1 Libraries, Compute resources and implementation of FFD and FJL
We use Jax and Hugging Face libraries; experiments in Sec. 5.2 were carried out using one GPU V100 or TPUv2 (8 cores). Experiments in Sec. 5.4 used 2 V 100s in the classification setting and A100s the generation setting. Experiments Sec. 5.5 used 2 1000s. Experiments were performed on cloud infrastructure and the virtual machines employed for each experiment had at most 32GB of RAM. We used the Jax profiler tool extract information about peak memory usage and wall time. We checked that under multiple runs the wall time estimates reported by the profiler tool stay within 10% relative error, while the peak memory usage does not significantly change. For FJL and FFD the lookup operation is implemented as in previous work [1 storing the permutations or entries to sample using a vector. Regarding permutations there an important implementation detail: if one wants the implementation of FFD to give the same results in implicit and explicit mode the
23

---

permutation used needs to be inverted between the two setups. For FJL, in implicit mode the lookup operation needs be transposed and is implemented an XLA scatter-add.
B.2
HVP in Implicit vs Explicit Form
The implicit implementation of the HVP is easy; one changes the loss to be defined on the image of projection & as follows:
Limplici (w) = [(0 + w)"w),
(10)
and then one computes the HVP at the origin w = 0. For the explicit form of the HVP. one takes vector in the image of 4 and lifts the vector o in the parameter space; then one computes the standard HVP for x² v and applies & to the result obtaining again vector in the image of Φ.
B.3 Step-by-step didactic implementation
We describe step-by-step didactic implementation in Jax and Flax.
First, we rely on the following modules.
import jax
import jax. numpy as jnp from typing import Sequence, Tuple from flax. core import apply from flax. core import init from flax. core import nn import scipy import functools
The core subroutine is applying pre-conditioners using Kronecker products. One can use the Jax' einsum operation to write it in few lines of code.
Listing 1: Kronecker Implemetation
def konpute_kronecker.eo :..]: """Computes shapes to decompose Kronecker products. ...
(+, dimension: int) Tuple[int,
For realistic use cases, bump it up, e.g. 1024
max max_ block_sizee 32
n dimension
Divide into block sizes shape = []
while n > 1:
shape append (ni max_block_size)
n //= nax_block_size
shape reverse return
def kronecker.product
*, x: inp.ndarray matrices: Sequence (jnp ndarray]
) -> jnp ndarray:
"""Performs kronecker product using Jax einsum. ".*
shape tuple (map (lambda x. shape [0], matrices))
y = x. reshape (shape)
num num dims (shape) Einsum iterative implementation.
for i, m in enumerate (matrices))
dims = ('.roi) for j in range (numm)im))
dimss f'lil num dims 1)

---

out_dims y_dims. rep str (num_dimms + 1). 1) operands f' dims dims]->(e out__immss y = jnp einsum (operands, y, m) return y. flatten())
The next step is implementation of QK which quite straightforward.
Listing 2: QK Implemetation
def q q_init (rng, shape, dtype=jnp float32): Random orthogonal matrix initializer: ...
jax. random. normal (rng, shape=shape)
9, jnp linalg. qr (x, mode= complete return q astype (dtyyp)
farget_dim D in paper
input_dim N in paper
vec is the N-vector to sketch to a D-vector. def (scope, vec, target_dim, input_dim):
Implements QK. .."
shapes compute_ hcaeeeeccrcceeeesseeeeeeeeeeeeee (imna
sigma = jnp sqrt
params = []
for i, in enumerate
p = scope. scope. prope.param(. pf (i)', qiiit, shape=(s, s))
params append
return sigma kronecker_product (x=vec, matrices=params) [: target_dim]
For using QK in implicit mode need to transpose QK; we illustrate how this can be carried out in Flax:
Listing 3: Transpose of QK
The tranpose QK vec is a D-vector to lift to an N-vector.
def qk_transpose (scope, vec, target_dim, input_dim)
shapes_1 compute_kroneckee shaapee (dimension-input_drrr shapes_2 = compute_
sigma = jnp. sqrt haweeerrrrocceeessseeeeeee
params
for i, (s_1, 2_2) in enumerate (zip (caa((sapes___ shapes_ __
p = scope. param params append (p [p.T[[:__])
qiinit, shape=(s_1, s_1))
return sigma kromecker_product (x=vec, matrices-params)
The implementation of other algorithms is more challenging: include implementation of AFFD; it is easy to figure out the implementations of the others (note that for FFD one needs to start with the implementation of the transpose because FFD defined as random feature generator).
Listing 4: Implementation of AFED
def uhit_hhaamarr (rng, shape: Tuple[int, int], permute_col: bool) -> Sequence [jnp- ndarrayl:
.." ""eeate randomly permuted hadamard matrices. ..*
matrix jnp. array linalg. aadaape [0])
pi jax. random. penitommpemmuuaionncc.c (rng, shape [0])
if permute_coll
matrix = matrix.ati at[:, ()]
else: # permute rows
25

---

matrix attrix.a((pi. :] getget( (
return matrix
airget_dim D in paper
input_dim = N in paper
vec is the N-vector to sketch to D -ecccr...
def ascoe,, vec, target_dim, dim, input_dim): """Implement" AFFD. a.n
wwweeekooooceeeeeeeeeeee (dinemsion-inputtddii inn sigma inp.sqra (inpputt.iii//taret...iii
shapes
def init_bet (rng, shape, dtype=jnp. int32):
return jax. random choice (rng, jnp jnp. ar-ay 11, dtype='int32* 'nt32 )
shape)
b = scope param('B', init_ber, (input_dim,)) vec vec _l_ppaams [] for i, s in (oumerate(shapes))
h_1 = scope param ______ init_hadamard, shape=(s, s),
permute_col=Falss
h h_1_params append (h_
vec =
matrices=h p-rramms) / jnp sqrt (
input_dim)
def init_gauss (rng, shape, dtype=jnp. float32): return jax. random. normal (rng, shape=shape, dtype=dtype)
g scope param('G', init_gauss, (input_dim,))
vec = vec
h_2 ________________________________ params [] for i, in (eaeraatesshapess. h_2 scope. param (f H (_
haddmadd shape=(s, s),
permute_col=True
h_ 2_p_pas ap..pard
vec kronecker_productteee input_dim dim)
matrices-h.2-parames
sqrt
vec = sigma vec [: target_dim]
return vec
To transpose AFFD, we just need to reverse the above steps and transpose application of the Hadamard matrices.
Listing 5: Transpose of AFFD
The tranpose of AFFD vec is D-vector to lift to an N-vector.
def affd_transpose (scope, vec, target_dim, input_dim)
shapes_1 = compute
shapes_2 compute_kronecker_r Atroneckersshape (aanapsssdinsisiniiipp.dded sigma jnp. sqrt aaanupt__diittaagget.dd (iineessiiiiiiittttrreeeooo
srapure_kkrneeckr_sshpeee
h_2 _araassms [] for i, 1s_1, s_2) in enumerate (aineraateeipssaappppee shapes_2)) 2 h_2 = scope. param (f' init_ init_hadamard, shape=(s_1, s_1),
26

---

permute__ol=Truee)
h 2_2_mas (bp..[s]t_2
vec = kronecker.produut matrices-h.2_params)) jnp. sqrt ( input_dim)
def init_gauss (rng, shape, dtype=jnp. floatt2)
return jax. random normal (rng, shape=shape, dtype=dtype)
= scope param('G', init_gauss, (input_dim.)
vec = vec
h par-mas []
for i, s in mnumerate(shapes_1
h_1 = scope param (f "_1_(ii'' init init_hadamard, shape=(s, s),
permute_col=Fall
h _ar_ap_aam____ append (h_1.T)
vec = kronecker_produc (x=vec, matrices-h.._parmms)) / saptsqrtt input_dim)
def init_ber (rng, shape, dtype=jnp. int32)
return jax. random choice jnp. array array 1], dtype= int32'
)
shape)
= scope param( init_ber, (input_dim,))
vec = vec
vec
sigma
return vec
We now turn to a didactic implement of sketching gradients a loss functions in implicit and explicit mode. We first make assumption about the signature of loss and sketching functions
def loss_ fn (model_params batch) Loss function signature. .... pass
def sketch_ fn (saatchh_paaams, vec): .." Sketch function signature. """
pass
def transpose_sketeeeeeeeeessseeeee ftt.ssseetccchtsteeeeeeeeeee params, vec):
"""Trrsspos of sketch_fn fn signature. pass
Then here's how one can implement explicit and implicit gradient sketching in a few lines of code.
Listing 6: Implementation of Implicit and Explicit Gradient Sketching
def explicit_grad__kete sketch (model (model_params, params, sketch parcm_s batch) .." PPrrforms an explicit gradient sketch. ...
grad = jax. (loss_ fn) (model batch) return sketch_fn (ssecchhhsskttccaaaaans grad)
def
(modee sketch_params, batch, target_dim):
"""Performs implicit gradient sketch. """
def inner_loss_fee omega = transposed sketch_ fnricchhfffsskk (ssketch_params omega) return loss (ww.ellmmmmllammm... + omega, batch)
27

---

omega jnp. zeros ((target_ dim,)
grad jax .rad (inner loss_ fnss_nn (omega)
return grad
Here's how one can do the same for the sketched HVP.
Listing Implementation of Implicit and Explicit HVP Sketching
Note tangent__aaraee is D-dimensional vector
def explicit_hvp_skeeee cheesphsshhhhee (model_params tangent_params, sketch_params, batch):
"""Performs an explicit HVP sketch. """
tangent_params
hhsppppsee_kkecccc..hcccccc...pp fn
tangent_params)
loss_ = functools. parctaalsppartilll (loss_fn, batch=batch) grad_fn
fn = jax. gax. (1oss_)
hvp = jax. (grad fn, (model_paranl ,), (tangent_params, [1]
return sketch_ (wwttaffsssththaaaaaaaaee
hvp)
Note tangent_params is a D-dimensional vector
def sketch_params, batch, target_dim); """Performs an implicit HVP sketch.
(model tangent_params,
def transposedssketch.....
omega transposed sketch_ fn tsketcc_pprras omega)
return loss fn (mrael_params omega, batch)
omega = jnp. zeros ( ((tagget_dim,)
loss_ functools, partial (inner_loss_fi
grad_ fn jax (loss_)
hvp = jax. jvp (grad_fn, (omega,), (tangent_pprams )) [1] return hvp
B.4 Sketching and model parallelism
We take the case of AFFD, and show how the single device code may be lifted code employing model parallelism.
First, we rely the additional modules.
from jax. sharding import NamedSharding
from jax. experimental import shard_map
from jax. sharding import Mesh
from jax. swwrding import PartitionSpec as P
from jax. skarding import NamedSharding as NS from jax import lax
from jax import tree_util as tu
import numpy as np
We then define the device mesh; we assume 8 cores with 4-way model parallelism and 2-way data parallelism.
mesh = Mesh (
np. array (
j jax devices ()) reshape 42, (' data" 'model' ,))
We now lift initialization and application of the FlaX modules to model-parallel code:
28

---

Listing 8: Lift AFFD to model parallel code
target_dim D in paper input_dim = N in paper
vec is the N-vector to sketch to a D-vector.
def part_ fppt_en()ytree):
""""ariitions each parameter on the last dimension. "."
def inner_pant (p)
out = (None,) (d.ndim-1) ('model',) return P (
return tu. tree_ (inner pytree)
def init_affd_mp (scope, vec, target_dim, input_dim): ""IIitializzs AFFD for model-paralle code.
# bind dimensional arguments to make jax tracer happy with
affd_init functools panctools.partall. init(affd), target_dim= dim==argget_dimm imput_dim=input_ddim
params_shape jax. haap.evalashappaffffdiiiiitt... rng, vec) params_part part_fn(paarmm__hha
We need to redefine the input_dim because the code
# is now executed on each model partition.
fffd_init_fn sunctools.partie init(affd), target_dim=tee carget_dimstarget_ddin // mesh. shapel model '])
def init_fn(rng) vec):
# different rng on each model slice rng = jax. fold_in (rng, lax. axis_index (' model ')) out, params
"ffd_init initiiitttiinnnn (rng, vec)
# The vector output is fully replicated and we need to sum on the 'model' partitions
out = lax. pax.psumoout, axis_name= model
return out, params
return shard_map mapp init_fn,
mesh-mesh, (P part_f((vec))
out_specs=(P (None,), params-part), ) (rng, vec)
def apply_affd_md (params, vec, target_dim, input_dim) ""AAplles AFFD for model ridel-parallel code. ".n
We need to redefine the input_dim because the code is now executed on each model partition. affd_apply_fn aanctools.partin apply(affd) target_dim=-arget_e input_dim=input.cd // mesh. shapel' model 'J)
def hmply_fn(param:: vec):
out fffd_applyy_fmppprr faprrams, vec)
The vector output is fully replicated and we need sum # on the 'model' partitions
out = lax. po..ssuuoout axis_names model') return out

---

return shard_map maard_map( apply_fn, mesh=mesh, in_specs=part_fnt (tne-paar__n((ppaass vec)), Sessspes=PPNontt
) (params, vec)
We now illustrate how to use the previous code.
We create the vector on single device and partition it on the model axis. rng = jax. random rRNdom.RRNKKeyy (0)
target_ _ii 128
input_dim 1024
jax. random. normal (jax. random.... fold_in 1), shape=(input_din
.))
x_mp = jax. deviic_eeut (x, NS (Sesh P('model'
affd x __mp, _m, jax random. fold_in 3), x_mp, target_dim, input_dim)
mp mesh (init_ (
Consistency check
affd_ _______p_2222 mesh (aspllyy_affd__.)
affd_parame pwramsaaappppppppp x_mp, target_ dim, input_dim)
assert affd_x____________ _ffd_x_mp_2_ ump_2.shape == affd_ x ___p_____ shape inp.linalg norm (affd_x_mp_1
affd_x_mp
B.5 Algorithm for searching the intrinsic dimension
Our algorithm for searching the intrinsic dimension is in Listing 9 Without loss of generality we assume that the target metric needs to be maximized, e.g. for the loss one might use the negative loss.
Listing 9: An algorithm that searches the intrinsic dimension
def finetune (model_params, D_max: int, d: int, c: int): ""nFineuuee function signature.
Finetunes for steps in D_max dimensional subspace but zeros out the last D_max components of the gradient. Returns the updated model_params.
pass
def evaluate (model_params) ""Elll function signature.
pass
def gearch_intrinsic_cinessss
model_params, D_min: int, D_max: int, tau_target: float, int, delta: float):
model_params: initial model parameters.
D_min: start value for the intrinsic dimension.
D_max: maximum allowed value of the intrinsic dimension. tau_target: desired target metric.
c: number finetuning steps in which we expect improvement. delta: minimum expected improvement. ..n
D_min
tau_old = evaluate (model_paraml
while True:
model_params finetune (wwdel_ D_max, d, c)
30

---

tau_new evaluate if tau_new tau_target:
return else if tau_new tau_old < delta:
d = 2 if d D_max: raise SauueError("D_nne exceeded") tau_old tau_new
B.6 Hyper-parameters for Sections 5.2 and 5.3
Models were trained with the released code from [6] then checkpoints were converted to Jax for benchmarking the sketching algorithms.
B.7
Hyper-parameters for Sec. 5.4
Roberta was fine-tuned with a batch size of 32 10k steps with Adam and constant learning rate of 2 10-5. For the search algorithm 9 the learning rate was increased to 10-1, 5 = 0.1 and c 2k steps.
BART was fine-tuned with a batch size of 32 for 20k steps with Adam and a constant learning rate of 2 10-5, For the search algorithm 9 the learning rate was increased to 10 10-4, 8 Rougel 0.5, Rouge2 = 0.5 and = 2k steps; the total number of steps was increased to 40k.
B.8 Hyper- parameers for Sec. 5.5
We consider the SGD optimizer as in previous work 12 the batch size was 32, the learning rate set 10-5
C Appendix: Theory
C.1 Definition of Higher order sketches
For higher-order derivatives of L, one consider sketches of operators, For example the Hessian vector product is the operator HVP RN + RR given by HVP(u) = V²)(O) i.e. i.e. HVP(u);-
L()uu. A sketch of the Hessian vector product then obtained follows: S(O)(e)
@(HVP($7 u)) where u E RD so that we obtain an operator mapping RD - RDD Sketches of matrices extensively studied [22] to speed-up evaluation of matrix products. Extending the HVP-example, for an operator O mapping R&N to RR=N the transpose & is applied to the k input indices and & is applied the output indices obtain mapping S(O) R&D RaD via:
N
(8(
Σ
Σ
II
& Pigit Oipmiiitt
II
+ ja Uirnie (11)
-1=
C.2 Guarantees on distorting distances.
[10] can leverage the equation about concentration of the sketched norm (T) to prove that, given Al points in N the distances between their sketches in R are distorted by at most multiplicative factor 1 The point is that concentration arguments establish, for the 5 appearing in []. bound of the form (cxx((=e2 3²)): one can thus apply ( to the M(M-1) differences between points by requiring that Vlog is sufficiently large; this is considerable gain replacing bound in terms of 2 with one involving log M.
By the method of Moanson-Lindenstrann
8
C.3 Definition of the Walsh-Hadamard transform.
The Fastfood paper 13 defines H without scaling, so it is not an isometry; however we follow the definition with scaling as in Wikipedia that HN an isometry. Specifically, N needs to be power of 2; then for i < N let (ii) denote the vector of Os and Is and of length log, N, representing (A()))))
i in its binary form; then (Hx) )J
where (,) denotes the inner product.
31

---

C.4 Theorem C.1. There are some inputs x such that (z) does not satisfy ①
Failure of concentration for FFD.
Proof. Intuitively, the problem with FFD is that transposition fails to apply the pre-conditioner to some bad inputs. Recall that the FFD, when used for sketching, gets transposed because it is applied in implicit form, i.c generating random features that perturb the model parameters. The transposed operation of concatenation gives rise to sum; specifically, given unit vector E R., we decompose it into N/ D-Dokss of size D, denoting the b-th block by 2b. Then
N/D
FFD = Σ b=1
(12)
then choose x such that all blocks 2b = 0 for b > 1 and æn is such that (x,) 61, the first coordinate vector in R.,. Then
FFD(2)) o³gी
(13)
where 91 is the first entry of G then must have a² and [FD(())]2 cannot concentrate around 1 because si has unit variance.
C,5 Comparison to Kronecker products in 17
proposes two projection operators. The first one is
Ps = (1 R1 R2 1)
(14)
where I is vector of ones in R N and R3 is Gaussian of shape D vN that the memory of Po is O(VNN) The second one is
Ps
Q1
Q2.
(15)
where Q, is Gaussian shape VD x VN so that the memory cost of Pg is O(VDVN), Our QK proposal is more general it calls for more general Kronecker decomposition
Q = a Q()) 2 Q (2) x ... @ (
(16)
where Q(+) is of shape D, B,, II, D (reconstruction
the target dimension) and II, B; N
(reconstruction of the model dimension). So if we choose K = 2, D; DD and B, = VN we recover 8' Strictly speaking, P' is different from Pg. but one might reconstruct by averaging two of our Q's (16). both defined with with K 2: indeed, we select Q where Q(1) is the vector T/VN in (14) and Q(2) is of shape 1) VN; then select Q2 where Q(2) is T/vN and Q(1) of shape D VN; in other words, get Pe we restrict the sampling of one factor to / I/VN. Note that memory-wise our approach is more efficient than 177 as we allow K > 2; memory cost i= B, D,) which, choosing B, 2 A and D; = A' allows for a memory cost O(AA' log N). An implementation difference with [17 is that we sample from the special orthogonal group: we sample a Gaussian of shape D, B, and obtain Q(+) using the QR QR-decomposition.
O
C.6 Concentration result for QK: Theorem 3.3
Theorem C.2. Consider the projection algorithm QK where Q decomposes as Q(2) .... ((K) where Q(3) has shape Dr x Be; then
P 1 P < 110(z)112
D V (a+t)llllz)) 2 21-2 1-2 (pp))
(17)
In particular, as long as each D, is sufficiently large one still obtains a concentration result; the price to pay for the Kronecker product decomposition is that the concentration probability is dampened by the number of factors as ((1 + 1/1 1)² 2 ($)².
32

---

Proof. Step 1: Reduction to the case of applying single factor in the Kronecker product
representation of Q. We assume that Q decomposes as Q @.... Q(K) where Q(9) has shape D; x B5; if we reshape z into tensor of shape [B,,... Bx) indexed by (a1...., ,.,,), ) the output (xx) can be represented as a tensor of shape (D1,'.. (D)) indexed by (b,,-. bb)):
Q())br.rn.ba .K
Σ
ox Fay.aaa
(18)
We look into applying (8) step time. We reshape x shape (B1; B2 Bk) ob- taining a matrix x! we then contract it with Q over a₁ to obtain a matrix y of shape
(Di, B2 ... Bx). To apply Q we need to first reshape ( to shape (Di, B2, Bx), then transpose the first and second indices, and reshape it to a matrix X of shape (B2, D3 ... -x); we can then contract with 02.00 over 0g obtain matrix y of shape (D2, D2 ... Bx) ). should be clear how this procedure can be continued for each E (3.... K}. Assume that for each
we can prove that the Frobenius norm (which is the no- if we reshape it to be vector) of by.c concentrates around that of ( b..c up to a multiplicative factor / D
(i)
P =)) 12 ||r S V B,
/ Pa + + 2 21-1 8,;
(19)
by conditional independence of the matrices Q on one another we get that
K II (1 -,),
P(V D I 1=112 - 11Q())1l2 S
D I +
(20)
where we used I D, D and II -h
B, =.
Step 2: Using concentration of measure on the orthogonal group. The entries of Q are not independent because of the orthogonality requirement and the fact that the rows need to have norm equal to We will employ measure concentration without independence; as reference for notation and theorems we use [26]. From [26] 2.5.2] we recall the definition of the sub-Gaussian norm of real-valued random variable X as:
| lly inf (t 0: exxp(X2/e) s 2):
(21)
obtaining a bound on |X| 1l equivalent to a concentration inequality:
P(X > 1) < 2exp(-<t//1X11])
(22)
for a universal constant > O. Note that (2)) can be sampled on the orthogonal group SO(B,) by restricting to the first D,-rows in the case in which D; < B, (by sampling from (BB) and changing in case the sign of one of the last B, D, rows to ensure the determinant is 1), while the result we B; because then Q(i) is full-rank. We now invoke the concentration of measure for SO(B,) 26 5.2.7]; if f. SO( - Lipschitz:
are proving is trivial if D;
(If(C()) Ef(Q())) S &clliileap VB,
(23)
where C is a universal constant and the Lipschitz constant || ll computed using the Frobenius norm on the tangent space. We now define:
( bi x())))) 7
=
=
Σ
(24)
b,.
which has derivative:
af(Q0)
Ee
1( (1)
bie b.eeaagi aiie
(25)
aQuiac
|
33

---

the Cauchy-Schwaad inequality implies that:
(26)
aQu.a.
|
from which it follows that Il// I may rescale by constant; if we rescale it so that ||0]||2 then all the norms of the intermediate
one assumes that || 1/2 S 1. Note that to derive 20 we
results ||2,, ||Y( [/2 are most 1 because the matrices Q(+) are orthogonal. We have thus
established
lI(())) - Ef( S C
(27)
VB.
Step 3: Replacing ()) with something simpler estimate. A drawback of (27 is that the term E)( not easy estimate. However, using a symmetry argument, it is easy estimate )) indeed the D; variables (x are identically distributed and D, = B, one would get an isometry; so
b,.c
D B:
))
(28)
So we would like to replace (f( with v /E(((())) the intuition why this would work is that concentration around the mean is equivalent to concentration around the median; so Ff((())
concentrates around the median Mú; as f (()()) is non-negative, the median of (f() is M and this variable concentrates both around the mean and the median, and we have a closed form for the
mean (28) To make this more precise, by the fact that concentration around the mean is equivalent to
concentration around the mean (see 226 5.1.13]), we have
C VB,
(I(()())) - Millaaa S
(29)
where the constant C' might have changed but is universal. We then just need to show that (VE(((Q0))) - M, is O(1/vB.). By the triangle inequality:
W B(((QU))) M. S EIS(Q() - MP
(30)
and we can compute the right hand side with the layer cake decomposition:
v EII(())) M1 ( P()(QQ))) - M, z u) du
(31)
and we apply the concentration inequality (22) to the right hand side to get
(1, 2exp(-ūB,u) du
(II(QO)) M22 <
(32)
which implies that the right hand size is O(1 Bi). We have thus established
P |Irr - VV. ( 12 21) S 2exxp(CCB..).
(33)
We now deduce
19 conditional that it holds for j = 1, 1,... , _i- --1 so that we may assume that
(i 1) 1|2
2
1.
then we may take = E
D
( [ and get 19 with
2
B,
8, = 2 -4CC))))
(34)
Step 4: Choosing the e.. We just aim for 6, to be equal and that II. + &)) = + E; this
is achieved by letting (1 ()++) 1) - 1. In this case we may lower bound II, (1 - 6,) by
1 - 2, (D0o44CD((1++))//// 1)
34

---

C.7
Concentration result for AFFD
Theorem C.3. Consider sketching with AFFD; then
(11-=))l=lle & |FFDD=))] & (1 +=]1=11))
2 8.
(35)
where for each 81 we have
N
D log
8 $ 81 + exp
-C6²
+ exp
-Ce²
(36)
2 log
2N 0,
2N 01
Note that the first two terms on the right hand side (36) can be made arbitrarily small (especially as N is typically very large and 51 will affect C' in the third term); thus the effective bound for 8 is of the form d I exp (-Ce²
log No
Proof. We will again use the notation and conventions from 26 Let us recall the definition of AFFD:
((x) Ro(a H2 G. H1 B(x));
(37)
without loss of generality we will assume that ||0|| = 1.
Step 1: Using the pre-conditioner H1 distribute the mass of Note that H1 is an N N- dimensional matrix with entries of the form + .N each entry of the vector H1 B(x) is of the
form Σ, abb. where the [br are independent Bernoulli and 8; = +1; applying Hoeffding's
VN
inequality 26
Th2.2 each entry of H; (zz) obtain that:
too 2N exp
t2x 2
P [H1
2
(38)
VN
intuitively, the norm of each entry H1 (x) cannot become much larger than a multiple of its variance VN' this is the purpose of using pre-conditioner to distribute the mass of x.
Step Decomposing (xxx)
We now let H, B(x) so that o(2) a RR G, u); we let II be the permutation associated with rearranging the columns of H2 so that
(x)) 0²
Σ
(39)
in (39h we decompose the effect of the diagonal and the off-diagonal terms obtaining
D
N
| (0))
- Σ
9z
+
ΣΣ
Σ
Han Man9
(40)
12
We now use the first term Ti to compute
D
ET
o² Σ
(41)
N
j=1
where we used the fact the components of G2 have unit variance; as H1 B(r) is an isometry, to have
ET we just need to set N
D
Step 3: Concentration for VT. If we regard VTi as a function (G.), we have
1/2
N
f(G.)) =
Σ
9
u
(42)
1=1
35

---

conditional on the good event E goood that || Ba))llo S fx. this function is Loc -Lipschitz in the
VN
VN
1 no-norm of Gu; applying concentration for the Gaussian measure on N 26 Sec. 5.2.2] we get that conditional Egood
P ((VTi- -17--1221 S exp (-c N E²
(43)
Step 4: Concentration for [22) We would like to claim that T2 is small with high probability. This is the second point in which we use the pre-conditioner H; the intuition is that applying the cre-conditooner H to G2 further reduces the finite sample correlation of the rows of the resulting matrix. Let us rewrite T2 as follows:
D N
T:
H. M.) Maa Hono) uybk9)9ki
(44)
D
B3
so we have reduced T2 to a bi-linear form 6.2.1] we have
Σj, Ba by the Hanson-Wright inequality
26
5²
P(T2 2 S exp
-CCmin
(45)
B
|B||s
where ||]|| is the Frobenious norm of B and [|||| is the spectral nomn. Let look By,8 conditional on Egood:
D Byd Σ H, Hanny
(46)
Recall now that D and N are both powers of 2 and the definition of H in Sec. C.3 if A(j)
denotes the binary vector, of length log2 N,, representing an integer S N we have Hi,r(())
1 VN (-I) where (a, b) denotes the inner product of two vectors of length log, N. We thus obtain the bound:
1 N
Ba
Σ (-
(47)
D
is 0 unless the vectors A (II()) and (II(()) agree in the first logg D entries, otherwise varying S D we can always find two terms in the sum that cancel each other by flipping the parity of in the first slot where the vectors A(T(G)) and (TT(k)) differ. So for each j. there are at most 27 possible k-s such that [Byal is non-zero; moreover, as the sum DL (-1)(4()) (-))))AAAA))))))) A(I()))+
the sum --
most 1) in absolute value, we get
N N
4 N
D
4
B
B
<
(48)
D2
No
D
note that a bound on the spectral norm |BIls is trivial from 47 as we can just take the maximum of the absolute values of the |By which is bounded by f Given the stronger bound for |B]|s and
that < 1, the minimum in the exponential in (45) is achieved by the term involving || and we thus obtain:
P([ 2)) S exp (-C=²
(49)
1I
Step 5: Picking up constants. Conditional on Egood and on T 2 we have:
+ T2 1| VT
V T 1 1|+ 2/72/
(50)
We bound the right hand side of (50) by 38 if Egood and the concentration inequalities for VTV and T2 hold, Thus, by decreasing the constant C' by factor 9, we have
P ((N Ti + T2 - 11 2 8) S 8
(51)
36

---

where
+ exp (-c
N ?
D
5 = 2N exp
x
exp
-Ce²
(52)
2
14
having fixed a small ò1, we set too
/ 2 log 2N
we get
8 s 81 + exp (-ce-
N
(-Ce²
D
+ +exp
(53)
2 log 81
2N
2N 4log 61
Comparison with [4] It seems plausible that Step 4 could be carried out with arguments similar to those of [4] Lem. 16, Lem. 17] by analyzing the chromatic number of the P-system associated with H; however think the method that uses the Hanson- Hanson-Wright Wright inequality is more simple for this specific case.
37

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer:
[Yes]
Justification: the theoretical justification is provided in Sec. and Sec. the experimental evidence is provided Sec. 5.3
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
Limitations are discussed in Section
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g.. independence assumptions, noiseless settings, model well-specification. asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
38

---

Theorems are stated with the full assumptions in the paper, proofs are provided As the proofs can be lengthy, the main proof ingredients (e.g. concentration for quadratic forms, analysis via explicit sketching) are explained in the paper.
Justification:
in Appendix C
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proof' in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: The complete experiment setup is described in Appendix B Python code of Jax implementations of the algorithms used is also provided in Appendix
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer this question will not perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm.
(b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way access this model for reproducing the results way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path reproducing verifying the results.
5. Open access to data and code
39

---

Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [No]
Justification: All datasets and models used are public and can be obtained from HuggingFace. Python code to implement the proposed algorithms (in Jax) provided Appendix B
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips... cc/
public/guides http:///ooeeeaaaaaiiooonooooooaa
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to to reproduce the results, See the NeurIPS code and data submission guidelines (https
//nips.cc/publl wttp:cccppppccccceecc.c carolliccguieee www.dedeeeooeebmmeeeeeee www.hibmmsssooPooo ssionPolicy for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: The main paper provides the core ideas for each experimental setup and full details reproduce the results are available Appendix B
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [No]
Justification: Experiments in Sec. 5.44 require multiple runs to check the consistency and accuracy of our estimation method for the intrinsic dimension. As the target quantity is discrete, we have opted to report the values across the random seeds in Appendix(A instead of using error bars. For wall time estimates we checked that on different runs the relative errors stay withing 10%, while estimates did not change for the peak memory usage (discussion about using the Jax profiler in Appendix B).
Guidelines:
The answer NA means that the paper does not include experiments.
40

---

The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables figures symmetric error bars that would yield results that are out range (e.g. negative rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: See Appendix B
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics httpp:://ooriisscccpppiie http://wwwrrarccoououu/c/actaaaa
Answer: [Yes]
Justification:
We followed the guidelines.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is a special consid- eration due to laws or regulations their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification:
We do not foresee societal impacts for this work.
Guidelines:
41

---

The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g.. gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
No releases of generative models or datasets.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification:
For the experiments in Sec. 5.2 we used the setup of (6 which has been cited.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset used and, possible, include a URL
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
42

---

If assets are released, the license, copyright information, and terms of use in the thcode has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
package should be provided. For popular datasets, paperswi
latasets.paperssi
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided. this information not available online, the authors are encouraged to reach out to the asset's s creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
No new assets are released.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moodl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
No crowdsourcing, human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then much detail as possible should be included in the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country institution) were obtained?
Answer: [NA]
Justification:
No crowdsoureing, no human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
43

---

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
44


Reviews:
[{"review_id": "ZGQZmNLe60", "reviewer": "Reviewer_kKVg", "summary": "Gradient information is useful for various tasks such as training data attribution and intrinsic dimension analysis, but often suffers from huge compute/memory costs, limiting its practical utility. This paper focuses on several popular sketching approaches (e.g., FJL, FFD), points out lookup-based memory as the major bottleneck in enabling efficient computations on modern accelerators, and proposes the improvement using the Kronecker products with theoretical guarantees. Their scalable sketching algorithms are applied to downstream applications including training data attribution, intrinsic dimension estimation, and Hessian analysis, and provide various useful insights, some of which are contradictory to existing common beliefs.", "questions": "1. I wonder how easy it is to combine proposed sketching algorithms with (data-parallel) distributed computing. While I appreciate general efficiency improvements in AFJL and AFFD, I believe distributed computing becomes necessary at some point for scaling these analyses to large-scale networks and datasets. Can you enable distributed computing simply by wrapping the model or do you need special implementation tricks?\n\n2. How are ground-truth eigenvalues are obtained in Table 1? If those eigenvalues are not exact (e.g. approximated using some iterative methods), then being close to approximate eigenvalues do not necessarily imply better accuracy?", "limitations": "The authors have appropriately discussed limitations.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "I agree that gradient information is useful in understanding important characteristics of neural networks, and sketching is a promising approach to reduce the cost of storing/computing gradient information. Indeed, (random) low-rank gradient projection has been studied in several prior research (e.g., Arnoldi IF, TRAK), with similar downstream applications in mind. That being said, I believe one of the main contributions of this work is in improving efficiency of existing sketching algorithms including FJL and FFD. Their solution of using the Kronecker product is, in my opinion, a smart strategy. Furthermore, their theoretical analyses are very interesting and make their arguments more rigorous. With scalable sketching algorithms, they were able to scale various analyses that were only applicable to small-scale networks to much larger networks, and provided many useful insights. For instance, many people have previously observed that full fine-tuning often achieves better performance than LoRA fine-tuning in LLMs, and the intrinsic dimension analysis in this paper can experimentally corroborates such observation. Overall, I believe the proposed algorithms would be valuable tools for understanding various phenomena in neural networks at scale.", "weaknesses": "In overall, I believe this paper is a strong paper *if* I limit the scope to gradient sketching. However, I am confused by several arguments.\n\n1. TDA (Sec 5.1): The authors claimed that layer-selection for TDA is inadvisable by pointing out the low (Pearson) correlation with the full gradient dot product between (x, z) pairs. While a high correlation with full gradient dot products could indicate accurate sketching, it doesn't necessarily imply better TDA performance. There are better options for TDA evaluation such as linear data modeling score from TRAK, mislabel detection, data subset selection, brittleness tests, etc. \n\n2. Sketching for HVP: While the authors cited matrix sketching for dealing with higher-order gradient sketching, I am unsure if theoretical guarantees apply in this case. For instance, if we adopt the Fisher information approximation of the Hessian, we can understand the Hessian as the gradient covariance matrix. If the vector in HVP is some sort of gradients, we can understand HVP, especially in influence functions, as a whitening operation that makes all components equally important. In this case, information loss from sketching is not negligible, and the general sketching argument (ie Eq(1)) doesn't really hold? \n\n3. Clarity: I find notations in the paper to be a bit confusing. For instance, in Eq. (2, 3, 6), they used respectively P, F, \\Phi to indicate the sketching process.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the insightful comments, taking the time to read the paper, and finding the insights provided by our sketching algorithms valuable.\n\n**Answers to Specific Questions**\n\n1. We focused on models that don't require partitioning weights across devices. For larger models, sketching can be applied to individual partitions. We commit to including an example of lifting single-device code to multi-device code in JAX in Appendix B. Specifically, if the JAX model is sharded using jit with sharding specifications, shard-map can lift sketching algorithms to a sharded version.\n\n2. We follow [11, Krishnan et al.], as exact eigenvalues are infeasible due to model size. This estimation is accurate, with error [11, eq(7)] decreasing exponentially in the number of iterations.\n\n**Clarifications related to weaknesses**\n\nWe hope these replies help to clarify points raised in the weaknesses.\n\n1. ``` TDA (Sec 5.1): The authors  ... brittleness tests, etc.```:  We agree that high correlation with full gradient dot products doesn't measure TDA performance in the long time range; however it is a practical metric in the short time range and a building block of more computationally intensive methods like TRAK. For example, in the short time range, gradient dot products correlate with loss changes and are relevant to select examples for error correction [Schioppa et al. https://arxiv.org/pdf/2305.16971 ]. Evaluating with LDS from TRAK would introduce more hyperparameters and computation (>50 models on different subsets of the data);  TRAK itself relies on good gradient sketches measured by dot products [TRAK, https://arxiv.org/pdf/2303.14186, C.2], and the TRAK authors  make the following point: *as long as we preserve the inner products to sufficient accuracy, the resulting system has approximately the same evolution as the original one [they refer to tracing the training dynamics, our note]. This justifies replacing the gradient features with their random projections*\n\n2. ```Sketching for HVP: ... (ie Eq(1)) doesn't really hold?```: Theoretical guarantees for matrix sketching apply here. The Hessian has limited bulk [11, Krishnan et al.], exploited by Arnoldi-based influence functions [21, Schioppa et al. ]. Sketching provides guarantees for approximating matrix spectrums [22, Swartworth et al.] and solving linear systems [20, Sarlos, Thm 12].\n\n3. ```Clarity: I find notations... to indicate the sketching process.```: Would replacing P, F, and $\\Phi$ with $\\Phi$ throughout the paper address this concern?", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "sXkJIPmmi2", "reviewer": "Reviewer_NhaV", "summary": "This paper introduces new methods for sketching high-dimensional gradients and HVPs. These are important building blocks for tools like training data attribution and Hessian spectrum analysis. Their methods introduces both empirical and theoretical improvements over prior methods, and are used to demonstrate new insights about properties of pre-trained LMs.", "questions": "- Would benefit from better exposition of the different methods in Section 3. Some type of figure would be more useful to get intuition for differences between different methods (it's a bit hard to keep track with so many different symbols). Happy to revisit if there's a newer version.\n- Some missing text, e.g., there seems to be no transition at the start of S4", "limitations": "none", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- Clearly written and organized\n- Addresses a practically relevant and timely technical problem (projecting high-dim gradients and HVPs) \n- Provides interesting new insights (e.g., that intrinsic dimension is not that low for LMs)", "weaknesses": "Nothing major, but:\n- Misses some important related work ([1], [2])\n\t- but gives good coverage otherwise\n- Gradient inner product is not a good TDA estimate (many papers have found this now, e.g, [1][3])\n\t- Estimating inner product is fine, but be clear that's the goal then\n\n[1] TRAK: Attributing model behavior at scale https://arxiv.org/abs/2303.14186\n[2] Faithful and Efficient Explanations for Neural Networks via Neural Tangent Kernel Surrogate Models https://arxiv.org/abs/2305.14585\n[3] Training Data Attribution via Approximate Unrolled Differentiation https://arxiv.org/abs/2405.12186", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank the reviewer for their insightful feedback and positive assessment of our work. We commit to include [1 & 2] suggested in the review to the Related Work.\n\n**Answers to Questions**\n\n```Would benefit from  ... there's a newer version.```:  We acknowledge the need for improved clarity in Section 3. While we cannot upload a revised paper during the rebuttal phase, we have included a proposed diagram (Figure 2) in the rebuttal PDF. We believe this visual aid will significantly enhance understanding of the various methods. We would greatly appreciate it if you could take a moment to review the diagram.\n\n```Some missing text, ... at the start of S4```: Thank you for pointing this out. We commit to ensure that a smooth transition is added at the beginning of Section 4.\n\n**A clarification on Weaknesses**\n\n```Gradient inner product is not a good TDA estimate```: We agree that high correlation with full gradient dot products alone does not guarantee optimal TDA performance in the long time range. However, it serves as a practical metric for short-term time evaluations and is a foundational component in more computationally demanding methods like TRAK. For instance, in the short time range, gradient dot products correlate with loss changes and are useful for example selection in error correction [Schioppa et al. https://arxiv.org/pdf/2305.16971 ]. Notably, TRAK itself relies on accurate gradient sketches as measured by dot products [TRAK, https://arxiv.org/pdf/2303.14186, C.2], with the authors stating: *as long as we preserve the inner products to sufficient accuracy, the resulting system has approximately the same evolution as the original one.* If layer selection is employed for computational efficiency, any adverse impact on the estimation of gradient inner products would have consequences for TDA even in the long run. We commit to revise the paper to make it more transparent that our method goal is to estimate inner products.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "OrnbexPI9x", "reviewer": "Reviewer_9kYE", "summary": "The authors present a framework for scalable gradient sketching and HVP sketching. They introduce three algorithms AFFD, AFJL, QK and provide guarantees for sketching. The paper focuses on three applications: training data attribution, intrinsic dimension computation, and Hessian spectra analysis. These are all implemented with a focus on pre-trained language models.", "questions": "*  Which components of the 6 contributions in the introduction are the main contributions of the paper?\n* Just above equation (2), “A sketch of the HVP can be obtained as…”: Is this a contribution of the paper, or is there a reference to this available?\n* In practice, which of the algorithms should one use and why?\n* Could the authors define what the meaning of an influence score is for this paper?", "limitations": "The authors focus on scaling their approach to GPUs and TPUs and adequately talk about the memory constraints and improvements made by each approach. No code was included with the paper.", "rating": 5, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "* The scalable approaches to sketching are well-motivated and interesting. It seems important to figure out useful sketching routines that scale to high-dimensional parameter spaces that we are seeing in pre-trained language models.\n* The theory appears to be a strength of the paper. It appears that there are two approaches taken in the paper, the direct approach that implements a gradient in the original parameter space and the implicit approach that implements a gradient in the sketched dimension. It is interesting to see the comparison made between these two different approaches.\n* The introduction of the FFT as a faster pre-conditioner appears novel.", "weaknesses": "* Presentation is a weakness of the paper. Since the paper is proposing three different algorithms along with three different applications, it makes it challenging for a reader to follow the contributions and comparisons being made. For example, there seems to be two separate topics being discussed: 1) New theory/algorithms for scalable sketching, and 2) Implications for Pre-trained large language models. Additionally 2) aims to cover three different insights including layer selection, intrinsic dimension, and LLM Spectra. It is a challenge to cover all of the aforementioned points in enough detail in a 9 page paper. \n* Table 1 is not clear how it is related to the proposed algorithms. In particular, TDA (section 5.1) just writes in bold “Our findings indicate the unreliability of layer selection (Table 1)”. Layer selection does not seem to be well defined, nor related to  AFFD, AFJL, QK. \n* Table 3 is also unclear. Which algorithm should one use in practice?\n* Minor: Figure 1 is hard to read and is too small. A single sentence in the caption telling the reader why these plots are significant (rather than just a description) would help link it back to the contributions.\n* Minor: Related to presentation is the use of citations and acronyms. It seems like many acronyms are not defined, e.g. AFFD, AFJL, and QK. Also, citations should not read “While [11] developed…”. It should be “ While X et al. 2024 [11] developed…”. There is a command for this in Latex.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for taking the time to read the paper, give feedback on the presentation and appreciating the theoretical part of the paper.\n\n**Answers to Questions**\n\n```Which components of the 6 contributions in the introduction are the main contributions of the paper?```: The main contributions are 1-3 (lines 45-54). The other 3 bullet points in the intro (59-68) are novel insights enabled by contributions 1-3.\n\n```In practice, which of the algorithms should one use and why?``` The choice depends on the target dimension and hardware, with AFFD excelling in accuracy for smaller dimensions, and QK/AFJL offering speed advantages on TPUs and GPUs, respectively. Let us look at this more closely. If the target dimension $D$ is constrained (<4k), we recommend AFFD (see Sec 5.2) for its superior accuracy. For larger $D$, QK on TPU and AFJL (with the FFT preconditioner) on GPU V100 are the fastest options. Refer to Table 2 for QK on TPU performance. For AFJL, the vanilla version takes 82ms (Table 2), and the FFT preconditioner provides a 64\\% speed-up on top of that (Table 3). We believe our analysis offers a valuable set of adaptable options rather than a single ``best'' algorithm.\n\n```Could the authors define what the meaning of an influence score is for this paper?``` We define influence score as the gradient inner product $\\nabla_{\\theta}L(\\theta,x) \\cdot \\nabla_{\\theta}L(\\theta,z)$ (l279). We focus on this for three reasons:\n\n1. Practicality: In the short term, it correlates with loss changes when adjusting the weighting of specific data points [Schioppa et al., https://arxiv.org/pdf/2305.16971 ].\n1. Foundation for Advanced Methods: It serves as a building block for methods utilizing Hessian pre-conditioners or relying on gradient sketches across\nmultiple model runs [TRAK, https://arxiv.org/pdf/2303.14186 ].\n1. Clarity: The definition is straightforward and avoids additional hyper-parameters to tune (e.g., number of models or size of removed datasets in [TRAK]).\n\n\n**Clarifications related to weaknesses**\n\nWe hope to address here additional concerns raised in the review. \n\n```Minor: Figure 1 is hard to read and is too small ... There is a command for this in Latex.```: We commit to incorporate these suggestions to improve clarity and presentation.\n\n ```Since the paper is proposing three different algorithms along with three different applications... It is a challenge to cover all of the aforementioned points in enough detail in a 9 page paper.```: We acknowledge the challenge of covering this material comprehensively in a 9-page paper. We believe there is value in presenting both (1) new theory/algorithms and (2) their applications to LLMs. Our strategy has been to split theory (Sec 3) from applications (Sec 4 & 5); each subsection of Sec 5 corresponds to a different application, so the reader can directly jump to the parts they are more interested in. We also aimed to provide sufficient references for interested readers to delve deeper. Additionally, we have included supplementary material in the Appendix. We welcome suggestions for further content to enhance the Appendix.\n\n```No code was included with the paper.```:  Appendix B contains a step-by-step tutorial with Python code.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "TISJ6ssl6y", "reviewer": "Reviewer_xaAC", "summary": "The paper proposes efficient sketching algorithms designed to address memory constraints in large-scale models. The key idea is to eliminate the multiplication of a dense projection matrix multiplication (large matrix materialization) found in the existing sketching algorithms such as in FJL and FastFood. The authors provide theoretical guarantees for their proposed algorithms: AFFD, AFJL, and QK. They showcase the advantages of these methods in three different contexts: training data attribution, intrinsic dimension estimation, and eigenvalue estimation of the Hessian.", "questions": "I have also included suggestions in addition to questions.\n\n- The authors should be more specific when discussing Training Data Attribution (TDA) methods. Rather than using the broad term \"TDA,\" they should explicitly mention that their work focuses on gradient-based TDA methods, such as TracIn [3] and GradDot [4]. TDA is a wide-ranging concept that encompasses various approaches, some of which do not align with the description provided in line 22.\n- In the introduction, the authors state that random projections have been implemented using dense matrices for TDA. Several works perform random projections without explicitly materializing dense matrices. The same holds for the experiment section in 5.\n- The statement “memory constraints limited their investigation of ID in generative tasks to 500k” in line 98 is unclear. What does ID stand for?\n- In line 103, the authors claim that their work makes influence function computations more efficient. Please clarify this point, considering that influence functions typically require inverse Hessian-vector product (IHVP) estimation, which isn't directly addressed in Section 5.1. Is this claim primarily based on improving methods like the Arnoldi iteration?\n- In Section 5.1, the conclusion that \"layer selection coupled with dense projections faces severe scalability limitations\" relative to AFFD, QK, and AFJL is not clear. Consider including tabular results (similar to Table 5) in the main paper to support this claim.\n- Do you anticipate any challenges in applying these techniques to even larger language models?\n\nMinor Comments\n\n- Line 129: Replace \"D\" with \"$D$\" for consistent mathematical notation.\n- Line 150: Clarify that $B$ and $H_D$ now have dimensions $D \\times D$.\n- Line 172: Specify the relevant section when mentioning \"unacceptable early TPU results\" for clearer cross-referencing.\n- Line 227: Add a period at the end of the sentence.\n- In the case of TDA, the rankings between the gradient dot products are important. Are there any reasons why authors do not use the Spearman correlation?\n- In line 311, it would be helpful if the authors provided a direct explanation of why explicit sketching provides substation speed up.\n- Figure 1 has low resolution, making it difficult to read when printed. \n\n[3] Pruthi, G., Liu, F., Kale, S., & Sundararajan, M. (2020). Estimating training data influence by tracing gradient descent. Advances in Neural Information Processing Systems, 33, 19920-19930.\n\n[4] Charpiat, G., Girard, N., Felardos, L., & Tarabalka, Y. (2019). Input similarity from the neural network perspective. Advances in Neural Information Processing Systems, 32.", "limitations": "The paper would benefit from a more comprehensive discussion of limitations.", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "- The paper addresses an important topic that is likely to be of interest to the NeurIPS community, with potential applications across various domains.\n- The authors provide a well-motivated problem statement, clearly articulating the need for efficient sketching algorithms.\n- The results presented are comprehensive and interesting, covering multiple aspects of the proposed algorithms and demonstrating their effectiveness across several use cases.\n- While the authors have not provided code with the submission, the appendix contains sufficiently detailed information to enable the reproduction of the results.", "weaknesses": "- While the paper is generally well-written, there are several areas that require improvement: a) Some notations are not clearly explained. b) Certain statements lack precision, potentially leading to ambiguity (specific examples are provided in the detailed comments section).\n- The authors do not sufficiently address the limitations of the proposed approach. The brief mention in the conclusion does not adequately address the potential drawbacks or constraints of the method (e.g., the limitations of QK).\n- The paper would benefit from a more comprehensive comparative analysis, such as the sketching algorithm used in TDA [1, 2].\n\n[1] Park, S. M., Georgiev, K., Ilyas, A., Leclerc, G., & Madry, A. (2023). Trak: Attributing model behavior at scale. arXiv preprint arXiv:2303.14186.\n\n[2] Xia, M., Malladi, S., Gururangan, S., Arora, S., & Chen, D. (2024). Less: Selecting influential data for targeted instruction tuning. arXiv preprint arXiv:2402.04333.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank the reviewer for their careful reading of our paper and the many valuable suggestions to improve the presentation. We also appreciate the reviewer's acknowledgment that this work is likely to be of interest to the NeurIPS community, with potential applications across various domains.\n\n**Answers to Questions**\n\n```The authors should be more specific  ... description provided in line 22.```:  We acknowledge the need for greater specificity when discussing Training Data Attribution (TDA) methods. We commit to revise the manuscript to explicitly mention that our work focuses on gradient-based TDA methods, such as TracIn [3] and GradDot [4]. We recognize that TDA is a broad concept encompassing various approaches, and we commit to clarify that our work is not applicable to all TDA techniques. We commit also to clarify that we consider methods built on estimating the gradient (e.g., TracIn, GradDot) or quantities derived by preconditioning the gradient with the inverse Hessian ([10, Koh et al.], [21, Schioppa et al.]).\n\n```In the introduction, the authors state that random projections ... experiment section in 5.```: We acknowledge that some works perform random projections without explicitly materializing dense matrices. However, the methods we are aware of (e.g., TRAK [\nhttps://arxiv.org/pdf/2303.14186, B.1]) still need to temporarily materialize chunks of the dense projection on-the-fly. This approach has two substantial disadvantages: (1) runtime grows linearly in the target dimension (they trade-off memory with compute), and (2) there is a need for specialized kernels for efficient implementation, and it's thus unclear how to implement this on TPUs. We have included a plot of runtimes (Figure 1) and a table (Table 1) in the rebuttal PDF to illustrate these points. We emphasize that implementation difficulty is a drawback of these methods, as evidenced by our Triton implementation of TRAK not outperforming the original CUDA kernel released by the TRAK authors. We also emphasize that in the case of using pure JAX we were not able to get a satisfactory implementation because one cannot control how the XLA compiler will deallocate these temporary arrays or their placement in the memory hierarchy. On the other hand, our methods can be written in pure JAX and have constant runtime in the\ntarget dimension.\n\n```The statement ... What does ID stand for?```: *ID* stands for Intrinsic Dimension (line 90). We meant that [15, Liu et al.] could only work with a target sketching dimension <= 500k, which prevented them from searching for the true value of the intrinsic dimension (ID), as we demonstrate on a summarization task where ID approaches the model dimension.\n\n```In line 103, ...  like the Arnoldi iteration?```:  Let us start with an observation from [7, Guo et al; Sec 3, eq(4)]: by symmetry of the Hessian, the iHVP can be either applied to the train or test points; however, given that the train data is much bigger, the iHVP should be applied just to the test point. Therefore the cost of an iHVP is limited only to test points. Now, our claim about improving influence function computations is based on two different approaches:\n\n1. *Sketching the output of an iterative iHVP solver applied to the gradient.* This doesn't change the iHVP computation itself, but makes storage and the search for influential examples more efficient.\n1. *Applying the Arnoldi iteration [21, Schioppa et al.] to the sketched Hessian, then using the sketched eigenvalues and eigenvectors to apply their method.* This is more efficient as it operates on vectors of dimension $D$ (sketching dimension) instead of the original model dimension $N$. Moreover, as we demonstrate in 5.4,\nwe can scale the Arnoldi iteration well beyond to what was done in [11, Krishnan et al.] and [21, Schioppa et al.].\n\nIn both cases storage is reduced from $N$ to $D$; in the second case the step of applying $k$ Arnoldi projectors can be performed very efficiently if $kD$ is sufficiently small to fit in the GPU/TPU memory as the computation can be parallelized across projectors.\n\n```In Section 5.1 ... to support this claim.```: We commit to add this table.\n\n```Do you anticipate ... larger language models?```: We focused on models that don't require partitioning weights across devices. For larger models, sketching can be applied to individual partitions. We commit to include an example in Appendix B demonstrating how to lift single-device code to multi-device code in JAX using shard-map to create sharded versions of sketching algorithms.\n\n```In the case of TDA, ... Spearman correlation?```: We use Pearson correlation because we are comparing sketched or layer-selected dot products against full dot products, where a linear relationship is expected. Spearman correlation is more appropriate when the relationship is assumed monotonic but not necessarily linear.\n\n**Other**\n\nWe commit to incorporate the suggested minor fixes and appreciate the reviewer pointing them out.  Regarding limitations, we are open to suggestions for expanding the discussion in Section 6 and would welcome specific areas the reviewer believes warrant further elaboration. At the moment we can think of emphasizing that QK might require higher values of the target dimension than AFFD.\n\n**Additional Clarifications Regarding LESS**\n\nWe provide further clarification on the comparative analysis with LESS (TRAK is discussed above)\n\nLESS builds upon the sketching of TRAK by applying it on top of LoRA. However, as noted by reviewer kKVg, *full fine-tuning often outperforms LoRA fine-tuning in LLMs, a finding corroborated by our intrinsic dimension analysis*. Additionally, LESS focuses on data selection, which is orthogonal to the core insights of our work.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We extend our gratitude to the reviewers for their insightful comments and suggestions. We particularly appreciate the encouraging feedback, acknowledging the practical relevance and timeliness of our work, as well as its potential interest to the NeurIPS community.

In this rebuttal, we would like to highlight two key points: the use of gradient dot products and a comparison to the dense projection implementation in TRAK [ https://arxiv.org/pdf/2303.14186 ].

**Gradient Dot Products**: We concur that high correlation with full gradient dot products may not be the definitive measure of long-term TDA performance; however it is a practical metric in the short time range and a building block of more computationally intensive methods like TRAK. For example, in the short time range, gradient dot products correlate with loss changes and are relevant to select examples for error correction [Schioppa et al. https://arxiv.org/pdf/2305.16971 ]. Evaluating with LDS from TRAK would introduce more hyperparameters and computation (>50 models on different subsets of the data); TRAK itself relies on accurate gradient sketches, as measured by dot products, and the authors emphasize that preserving inner products to sufficient accuracy results in a gradient-descent system that approximately preserves the same evolution
as the one corresponding to model re-training [TRAK, https://arxiv.org/pdf/2303.14186, C.2].

**Random projections like in TRAK**: While TRAK avoids materializing the full random projection, it still requires the temporary materialization of chunks of the dense projection. This leads to two significant drawbacks: (1) runtime scales linearly with the target dimension (memory traded off with compute), and (2) specialized kernels are necessary for efficient implementation, with unclear applicability to TPUs. Our attempts to implement a competitive version using pure JAX were unsuccessful due to the lack of control over memory allocation and placement. We have included a plot (Figure 1) and a table (Table 1) in the rebuttal PDF demonstrating the linear runtime growth of TRAK compared to the constant runtimes of AFFD and QK. The implementation challenges of TRAK are further highlighted by the fact that our Triton implementation does not outperform the original CUDA kernel released by the TRAK authors, underscoring the difficulty of efficiently implementing random projections in chunks.

Meta-review: The work is interesting, novel, albeit limited to gradient sketching and perhaps addresses a sub-community of NeurIPS. With this, all reviewers (and the AC) appreciate the contribution. The work have had an elaborate discussion between the authors and reviewers and we are looking forward to see their results reflected in the final submission.
Decision: Accept (poster)