Efficient Sketches for Training Data Attribution and Studying the Loss Landscape

Authors: andrea  schioppa
Venue: NeurIPS_cc_2024_Conference
Published: 2024-04-30 15:32:52

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.

Main paper: 
 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), eigenvalue estimation, and the computation of the intrinsic dimension. While previous studies on the intrinsic dimension have employed the Fastfood Transform to mitigate memory demands, we demonstrate its theoretical limitations for sketching and its persistent memory bottlenecks on modern accelerators. To resolve these issues, we propose novel sketching algorithms designed for modern hardware and underpinned by robust theoretical guarantees. Our experiments on 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, 11] and Hessian eigenvalue estimation [12] offer powerful insights into neural network behavior. TDA requires storing vectors of the same dimensionality N as the network's parameters for each training point, scaling linearly with dataset size (O(N T ) for T training points). Similarly, numerically stable Hessian eigenvalue estimation algorithms demand repeated Hessian-vector product (HVP) computations and storage, scaling with network size and iterations (O(N T ) 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 or HVPs into lower-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 a target dimension D) which necessitate the restriction of gradients to a subset of layers. Current TDA scaling methods require layer selection, as demonstrated with BERT in [8] , where only 15M parameters were used out of 110M. Beyond the computational cost of dense matrices, our experiments show that layer selection introduces substantial distortion in the estimation of TDA scores (Sec. 5.2). Recent work on neural network geometry [14] suggests using the Fastfood Transform for gradient sketching (abbr. FFD [13] ) with O(N ) memory; however, as the Fastfood Transform is a 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.1). 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:1. Scalable Gradient Sketching for Modern Accelerators and Networks: We introduce algorithms (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. 3.1). Our analysis further indicates a dimensionality reduction advantage for AFFD over QK, a 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. 4).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, so their usage should be avoided (Sec. 5.2). 2. High Intrinsic Dimension: In contrast to assumptions drawn from classification-based studies, 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 [14, 1, 17] . 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, 3, 12] .Paper organization Sec. 2 provides a survey of relevant research in the field, contextualizing our contributions. Sec. 3 introduces our novel sketching algorithms. We begin with necessary background material, analyze design choices, and provide a step-by-step implementation tutorial in Appendix B. Sec. 4 outlines our proposed techniques for efficient intrinsic dimension search and Hessian eigenvalue computation. Sec. 5 describes our experimental setup: subsections are aligned with Sections 3 and 4 to enhance the connection between theory and empirical results 
related work: 
 Sketching Sketching algorithms have been extensively studied (see surveys [27, 18, 15] ). 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 era 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 to bi-linear forms (Thm. 3.2) and the special orthogonal group (Thm. 3.3) . Recent work on PAC bounds [17] leveraged Kronecker-product 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 a sketching algorithm, absent in [17] . Additionally, we demonstrate that the Kronecker structure used in [17] is not essential for performance gains, highlighting that the true bottleneck in FFD and FJL is memory access. For a comprehensive comparison with [17] , please refer to Appendix C. Finally, to support our eigenvalue estimation, we leverage the theoretical guarantees outlined in [25] .The concept of intrinsic dimension (ID) offers a valuable metric for understanding the complexity of learning tasks. Originally employed to analyze loss landscapes [14] , intrinsic dimension has expanded into the study of language models. [1] 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 [16] , 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 ≤ 500k, which prevented searching for the true value of ID as we demonstrate on a 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-data attribution remains a 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, 20] . Furthermore, we offer eigenvector sketching to enhance methods like the Arnoldi iteration [24] . [19] has proposed to use dense random projections by materializing them in chunks and on-the-fly; drawbacks of this approach are: (1) the lack of the scalability in the target sketching dimension and (2) the need of hardware-dependent 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 [5] .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 a 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 3 -dimensional) for models like GPT-2L (770M parameters). This advancement significantly surpasses the memory constraints of the method utilized by [12] : 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, 9] within the context of pre-trained language models. 
design principles for efficient sketching algorithms: 
 In this section, we explore a design space for more efficient sketching algorithms. To establish a 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, we demonstrate that FFD lacks the theoretical foundation required for sketching. Relevant experimental findings are presented in Sections 5.2 and 5.3.A D-dimensional sketch of the gradient of a real-valued function L(θ) (θ ∈ R N ) is a random projection of the gradient ∇L ∈ R N to R D .To ensure this projection preserves essential geometric properties, the random projection operator Φ : R N → R D must, with high probability, concentrate the norm of Φ(x) around the norm of x. Mathematically, for each ε and δ there exists a large enough target dimension D(ε, δ) so that for D ≥ D(ε, δ):EQUATIONThis concept generalizes to sketching higher-order derivatives (see Appendix C). For our purposes, consider the Hessian vector product operator HVP : R N → R N defined as HVP(u) = ∇ 2 L(θ)(u) A sketch of the HVP can be obtained as v → Φ(HVP(Φ T v)), v ∈ R D . This sketch defines a linear mapping R D → R D [25] . While a simple Dense Sketch (using a random D × N Gaussian matrix) ensures the norm property, it has O(DN ) memory and O(DN 2 ) compute requirements. The FJL algorithm [2] addresses the compute cost with a sparser projection matrix:EQUATIONwhere: B ∈ R N ×N is a diagonal matrix with B i,i = ±1 (random signs); H N is the Ndimensional Walsh-Hadamard transform (Appendix C); G s is a sparse R D×N Gaussian matrix with (in-expectation) Θ(D log 2 M ) non-zero entries (M is a parameter). The H N • B component preconditions sparse inputs x for which (1) might otherwise fail. Implementing G s efficiently presents challenges on modern hardware and frameworks like Jax that lack native sparsity support.FFD: Implicit Gradient sketching The FFD transform, introduced by [13] in the context of kernel machines, provides a computationally efficient way to approximate high-dimensional feature maps. As a random feature generator [15] , FFD constructs high-dimensional random features (∈ R N ) from a lower-dimensional input vector (u ∈ R D ). Given u ∈ R D , FFD concatenates N D vectors of the form:EQUATIONwhere B and H D are as in (2) (and are both D × D-matrices), Π is a permutation matrix, G v is a diagonal D × D-matrix with i.i.d. standard Gaussian entries, and σ F is a normalization constant (see [13] ; 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 a map R D → R N , direct sketching of a gradient is not possible. To address this, [14] perform what we call an Implicit Gradient Sketch:EQUATIONthis formulation effectively applies the transpose of FFD to the gradient. While [13] establish certain properties of FFD, a complete proof of its suitability as a random feature generator is missing. Additionally, whether the FFD satisfies sketching guarantees (1), remains to be fully investigated (we will address it in Thm. 3.1).Explicit sketches. In light of equation ( 4), it's natural to consider whether a direct sketch of the gradient could be achieved using a map Φ : R N → R D . We define this as an Explicit Gradient Sketch:EQUATIONThis approach offers flexibility. For random feature generation methods like FFD, Φ needs to be implemented as 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-based tutorial on transposing sketching algorithms. Which one is the right approach? Intuitively, implicit sketches may seem more efficient since they avoid direct gradient materialization. However, as we'll demonstrate in Section 5.3, explicit sketches surprisingly offer significant performance advantages. For optimal accelerator usage, we limit factors to roughly 1024 columns. Building on this, we modify (3). In the resulting AFFD algorithm (6), we remove Π, use row-permuted factors in H π1 N , and column-permuted factors in H π2 N :EQUATIONwhere R D denotes restriction to the first D coordinates. While all matrices are now N × N , efficient Hadamard implementations avoid large matrix materializations (see Appendix B for our Jax code). We further introduce AFJL (7) , where H π2 N is removed:EQUATIONThis can be seen as replacing G s in FJL with a diagonal Gaussian matrix. Generalizations with multiple (G v , H π1 N ) pairs are possible but weren't explored due to the vanilla version's strong results. Empirical results on TPUs (Table 3 ) show the crucial impact of our changes. AFJL achieves a 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 20 on GPUs.Alternative pre-conditioners. To address the smoothing of sparse inputs in sketching algorithms, [2] introduced H N in (2) . The theoretical basis lies in the Heisenberg uncertainty principle, leveraging the Fourier Transform on a 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 H N with the FFT offers significant advantages. Experimentally, it reduces wall time by 62% on GPUs (Tab 4). Inspired by the Kronecker product structure enabling efficient H N implementation, we propose another generalized pre-conditioner, Q. This random N × N orthogonal matrix has a Kronecker decomposition of K independent orthogonal matrices of sizes {B i × B i } K i=1 , sampled according to the Haar measure on SO(B i ). Our approach allows direct application of Q without the additional diagonal matrix B, and offers a 40% wall time reduction on GPUs (Table 4 ). This Q pre-conditioner has the potential to unlock broader optimizations within sketching algorithm design as discussed in the next subsection.Direct usage of the pre-conditioner Q. Inspired by the design of the pre-conditioner Q, we introduce a novel sketching algorithm, QK. This ablation explores the direct use of Q to transform the input, potentially leading to improved efficiency and memory usage. QK is defined as:EQUATIONHere, Q is a random D × N -orthogonal matrix with a Kronecker product decomposition of K independent orthogonal matrices of sizes {D i × B i } K i=1 . Each factor is generated by sampling from SO(B i ) according to the Haar measure and restricting to the first D i 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 1 illustrates AFFD, AFJL and QK with diagrams.We now delve into the theoretical 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 (1).Note that this limitation is specific to FFD and not inherent to implicit-mode algorithms. This distinction is important: the explicit-mode formulation allowed to simplify the theoretical analysis to prove Thm. 3.1. Next, we establish a theoretical guarantee for our AFFD algorithm. Compared to [2] there is added complexity as independence arguments cannot be used for G v ; we thus apply the Hanson-Wright inequality for quadratic forms to prove: Theorem 3.2. AFFD satisfies (1) with for a universal constant C and for any δ 1 > 0.EQUATIONFinally, we analyze QK. Its structure allows for concentration arguments on orthogonal groups, leading to:Theorem 3.3. QK satisfies (1) with δ = 2 i exp(-4CD i ((1 + ε) 1/K -1) 2 ),for a universal constant C.Crucially, the bound in Thm. 3.3 is less favorable than that of Thm. 3.2. This is due to the 1/K-root, summation over sub-dimensions, and concentration depending on the D i . These theoretical insights support our experimental findings in Section 5.3, where QK requires higher target dimensions (D) to achieve performance comparable to AFFD. Because of space constraints the proofs are included in Appendix C. 
expanding the utility of sketching algorithms: 
 In this Section we expand the usage of sketching to other applications.Improving the search for the intrinsic dimension. The intrinsic dimension (D int ) [14] , is the minimum dimension (D) of a random subspace (V ) where SGD yields at least 90% of the full model's performance on a target metric (τ1 ). Sketching algorithms, more efficient than FastFood, already accelerate the search for D int [17] . Our memory-efficient algorithms enable us to investigate scenarios where D int may approach the model dimension. For instance, [16] applied FFD to finetune generative language models but capped the target dimension (D) to 500k due to memory limits. We propose a novel search algorithm that estimates D int in a single training run. Current methods rely on multiple runs across potential D int values. To streamline, consider a binary search approach (assuming D int is a power of 2) starting from an initial guess (D min ) up to the model parameter count (N ). This would require log 2 N Dmin runs. Instead, we propose a single training run where D increases progressively. The heuristic: a fixed computational budget (c steps) should yield an expected improvement of at least δ. We start with D = D min . If, after c steps, the target metric's improvement is less than δ or hasn't reached τ 90 , we double D. This yields an estimate D * within a factor of 2 of D int . A subsequent run with D = D * /2 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 D int N for classification, but for the generative task, D int depends on the choice of τ and can approach the parameter count N . This finding is significant, as it challenges prevailing assumptions about intrinsic dimensionality in such models [14, 1, 17] . These studies primarily focused on classification tasks, obtaining generalization bounds like O( [17, Sec. 3] ) where D int 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.√ D int ) ([1, Eq. 4],Scaling eigenvalue estimation Investigating the Hessian's spectrum and eigenvectors often relies on memory-bound iterative algorithms. The Arnoldi algorithm, as used in [12] , requires full reorthogonalization for stability. This necessitates storing vectors as large as the model itself for each iteration: a constraint that limits the number of estimable eigenvalues (e.g., Krylov subspaces of dimension 90 in [12] ). Inspired by the theoretical work of [25] , we propose sketching to address this memory bottleneck. With a sketching dimension of 10 6 , we can readily construct 1k-dimensional Krylov subspaces for pre-trained language models like BART, Roberta, and GPT-2L. This represents a breakthrough because the method in [12] would demand an impractical 3TB of storage for GPT-2L alone. Our technique enables exploration of conjectures [21, 12, 9] 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. 
experiments: 
 To thoroughly evaluate the proposed sketching methods, we present a comprehensive set of experiments. 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. 5.3). We then introduce and validate an algorithm for intrinsic dimension estimation, enabling computational savings (Sec. 5.4) and show-casing that the intrinsic dimensionality of generative tasks can be large. Finally, we apply our techniques to explore the evolution of the Hessian spectrum during pre-trained language model fine-tuning (Sec. 5.5).In TDA there are different ways to measure the similarity score between two examples x and y. In our experiments we opt for the TDA score defined as ∇ θ L(θ, x) • ∇ θ L(θ, z) 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 [19, 23] , it is a practical metric in the short time range and a building block of more computationally intensive methods like TRAK [19] . For example, in the short time range, gradient dot products correlate with loss changes and are relevant to select examples for error correction [23] . Evaluating with the LDS from TRAK would introduce more hyper-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 a gradient-descent system that approximately preserves the same evolution as the one corresponding to model re-training" [19, C.2].Past work [20, 28, 8] tackles the memory bottleneck of Training-Data Attribution by calculating gradients restricted to a specific layer and potentially applying a 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 closely 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 12 pairs of points (x, z), we compute the Pearson correlation r between the TDA score ∇ θ L(θ, x)•∇ θ L(θ, z) estimated using a layer-specific gradient and the ground truth based on the full gradient. We adopt the setup of [6] : a 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 1 with ground truth rarely exceed 90% and are significantly lower for most layers. In contrast, AFJL achieves a 98% correlation with a compact D = 2 13 . 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 a 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-preserving linear map R 1 → R 1 to those computed for a particular layer. However, this is not the case, with the relative mean absolute error at 20% on the best layer (Table 1 ). Finally, layer selection coupled with dense projections faces severe scalability limitations. Setting D = 2 12 within the same setup highlights this issue. Limited memory forces us to divide the computation into D/k projections to R k where k is the smallest power of 2 enabling a dense projection without an Out-of-Memory error. For token embeddings, we find k = 32, whereas other layers require k = 128 on a V100 GPU. Projecting to R k takes ∼ 31ms for token embeddings, resulting in a total of ∼ 4s to project to R D . Other layers require ∼ 24ms per projection to R k and ∼ 0.77s to project to R D , see Table 2 for a per-layer breakdown. Finally, an alternative approach to dense projections is to 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 are necessary for efficient implementation, with unclear applicability to TPUs; we include a demonstration of these limitations in Appendix A.In this section, we analyze the impact of the design choices presented in Sec. 3 on both sketching quality and performance. Regarding sketching quality, we find that small values of D often lead to remarkably accurate reconstructions of influence scores. For TDA scores, achieving correlation r ≥ 0.95 requires the following dimensions: FJL, FFD, AFFD: D = 2 10 ; AFJL: D = 2 12 ; QK: D = 2 14 . To reach r ≥ 0.99, simply increase each dimension by a factor of 8. While memory limitations prevented sketching HVPs with FJL and FFD for eigenvalue estimation, the other algorithms scaled well. We achieved a relative mean absolute error below 5% with the following dimensions: AFFD: D = 2 10 ; AFJL: D = 2 12 ; QK: D = 2 13 . Note that QK requires larger D, consistent with the theoretical comparisons in Theorems 3.3 and 3.2 (Sec. 3). Regarding perfomance, we outline here the key findings and refer the reader to Appendix A for a comprehensive analysis 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 3 . Second, explicit sketching consistently provides substantial speed-ups (see Table 4 ); we conjecture this is due to the fact that implicit sketching results in more memory accesses. Finally ,while modifying the pre-conditioner doesn't affect TPU performance, it significantly improves GPU performance, see Table 4 under the headings H N → FFT and H N → Q.Our experiments evaluate the efficiency and accuracy of our intrinsic dimension estimation algorithm (presented in Sec. 4). We consider two experimental setups: classification, where we fine-tune Roberta on SNLI with accuracy as the target metric; generation, where we fine-tune BART on XSUM for text summarization, using Rouge1 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; we repeat each search experiment with three random seeds obtaining estimates within a factor of 2 (Appendix B). We then verify that the estimated D * accurately represents the intrinsic dimension D int by fine-tuning the model in a D = D * /2-dimensional subspace and ensuring that the target metric remains below the τ 90 threshold (details in Appendix B). We point out that selecting good search parameters (c, δ) was relatively straightforward by observing target metric improvement during full fine-tuning . Regarding computational efficiency, our approach requires significantly fewer steps than binary or 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 at how the intrinsic dimension varies between classification and generation tasks. For classification, our findings are consistent with prior work [1] , were all algorithms yield D int = 2 13 N . For generation we first observe that the results are metric-dependent: for Rouge1, FFD and QK estimate D int = 2 25 while AFJL yields D = 2 24 ; for Rouge2, D int equals the full model dimension N ∼ 2 27 . In both cases, however, the intrinsic dimension is close to the full model dimension: this challenges assumptions about intrinsic dimension in the context of generative tasks. While prior studies [14, 1, 17] focused on classification with generalization bounds of the form O( √ D int ), 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. In this section, we examine the Hessian's evolution during pre-trained language model fine-tuning. Using sketching, as proposed in Section 4, we estimate the Hessian's eigenvalues and eigenvectors. We fine-tune Roberta and GPT-2L on SNLI (3 classes) and BART on XSUM, employing AFFD, with a target dimension D = 2 20 , to construct a 1k-dimensional Krylov subspace. This substantially extends the work of [12] who considered smaller models and smaller subspaces (90dimensional Krylov subspace). Spectrum estimation is performed every 1k steps, including initial steps {0, 10, 50, 100, 150, 200, 250, 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] . • Obs2: K -1 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 [12] .We find that these obervations don't fully translate to pre-trained language model fine-tuning. Regarding 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 it diminish over time. Regarding Obs2, outlier eigenvalues don't strictly adhere to the K -1 rule. Roberta has more outliers (6) , and the gradient-outlier alignment is less pronounced (27% Roberta, 35% GPT-2L, 8% BART) compared to smaller networks [12, 9] . See Fig. 2 (b ). Moreover, outliers emerge despite layer normalization. 
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 in 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 of a generative task, and the deviation of 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. Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, 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 to reproduce the results are available in Appendix B.Guidelines:• The answer NA means that the paper does not include experiments.• The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them. • The full details can be provided either with the code, in appendix, or as supplemental material. 
experiment statistical significance: 
 Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?Answer: [No] Justification: Experiments in Sec. 5.4 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.• 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 a 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 as 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. • If 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). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)? Answer: [NA] Justification: No releases of generative models or datasets. Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a 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 to 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 a 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 is used and, if 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.• If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. • If this information is not available online, the authors are encouraged to reach out to the asset's creators. 
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/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc. • The paper should discuss whether and how consent was obtained from people whose asset is used. • At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file. 
crowdsourcing 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 as details about compensation (if any)?Answer: [NA] Justification: No crowdsourcing, no human subjects.Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects. • Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as 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. 
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 an equivalent approval/review based on the requirements of your country or institution) were obtained?Answer: [NA] Justification: No crowdsourcing, no human subjects.Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects. • Depending on 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.• We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to 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. 
 
 Appendix: In Table 5 we include the full results of layer selection of 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] : NLP tasks which consists in fine-tuning GPT-2 on the WikiText-103 dataset and BART and zsRE.For each algorithm, we report in Table 6 the minimal value of log 2 D necessary to reach a Pearson correlation > 0.9x with the ground truth when estimating inner products of gradients using sketches.As expected from the worse concentration bound, QK requires a larger dimension. We conjecture that the fact that FFD is effective might be due to the gradient distribution giving 0-measure to the inputs that FFD would fail to sketch (Theorem 3.1). For AFFD, AFJL and QK we report in Table 7 the minimal value of log 2 D necessary to reach a relative mean absolute error err < x when estimating the top 10 eigenvalues of the Hessian.In Sec. 5.3 we pointed out that the FJL's wall time and memory usage increase with the target dimension D. We compare the peak memory usage and the wall time of FJL with that of AFJL on the inner product task of Sec. 5.3 on GPU (V100), see Figures 3, 4 In Table 8 we report the compute costs of sketching gradients in the setup of Sec. 5.3. In Table 9 we report the compute costs of sketching HVPs in the setup of Sec. 5.3. In Table 10 we report the values of D * across the 3 seeds used in each search experiment. We see good agreement within a factor of 2. For SNLI the target accuracy to exceed was τ 90 = 80.1%; for XSUM the value of Rouge1 to 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 = 2 26 ; in that case the final values of Rouge2 did not exceed 14.4, so it stayed well-below the required threshold that defines the intrinsic dimension D int using a 90% target value of the metric obtained by fine-tuning the full model. 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 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 7 ) and a table (Table 11 ) 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 We use Jax and Hugging Face libraries; experiments in Sec. 5.2 were carried out using one GPU V100 or a TPUv2 (8 cores). Experiments in Sec. 5.4 used 2 V100s in the classification setting and 2 A100s in the generation setting. Experiments in Sec. 5.5 used 2 A100s. 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 to 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 a 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 is an important implementation detail: if one wants the implementation of FFD to give the same results in implicit and explicit mode the permutation used needs to be inverted between the two setups. For FJL, in implicit mode the lookup operation needs to be transposed and is implemented as an XLA scatter-add.The implicit implementation of the HVP is easy; one changes the loss to be defined on the image of the projection Φ as follows:EQUATIONand then one computes the HVP at the origin ω = 0. For the explicit form of the HVP, one takes a vector ν in the image of Φ and lifts it to the vector Φ T ν in the parameter space; then one computes the standard HVP for Φ T ν and applies Φ to the result obtaining again a vector in the image of Φ.We describe a 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 functoolsThe core subroutine is applying pre-conditioners using Kronecker products. One can use the Jax' einsum operation to write it in a few lines of code. The implementation of other algorithms is more challenging; we include an 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 is defined as a random feature generator). We now turn to a didactic implement of sketching gradients of a loss functions in implicit and explicit mode. We first make an assumption about the signature of loss and sketching functions def loss_fn ( model_params , batch ): """ Loss function signature .""" pass def sketch_fn ( sketch_params , vec): """ Sketch function signature .""" pass def transpose_sketch_fn ( sketch_params , vec): """ Transpose of sketch_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_sketch ( model_params , sketch_params , batch ):""" Performs an explicit gradient sketch .""" grad = jax.grad( loss_fn )( model_params , batch ) return sketch_fn ( sketch_params , grad)def implicit_grad_sketch ( model_params , sketch_params , batch , target_dim ): """ Performs an implicit gradient sketch .""" def inner_loss_fn ( omega ): omega = transposed_sketch_fn ( sketch_params , omega ) return loss_fn ( model_params + omega , batch ) omega = jnp. zeros (( target_dim ,)) grad = jax.grad( inner_loss_fn )( omega )return gradHere's how one can do the same for the sketched HVP.Listing 7: Implementation of Implicit and Explicit HVP Sketching # Note tangent_params is a D-dimensional vector def explicit_hvp_sketch ( model_params , tangent_params , sketch_params , batch ): """ Performs an explicit HVP sketch .""" tangent_params = transposed_sketch_fn ( sketch_params , tangent_params ) loss_ = functools . partial (loss_fn , batch = batch ) grad_fn = jax.grad( loss_ ) hvp = jax.jvp(grad_fn , ( model_params ,) , ( tangent_params ,)) [1] return sketch_fn ( sketch_params , hvp) # Note tangent_params is a D-dimensional vector def implicit_hvp_sketch ( model_params , tangent_params , sketch_params , batch , target_dim ): """ Performs an implicit HVP sketch .""" def inner_loss_fn ( omega ): omega = transposed_sketch_fn ( sketch_params , omega ) return loss_fn ( model_params + omega , batch ) omega = jnp. zeros (( target_dim ,)) loss_ = functools . partial ( inner_loss_fn ) grad_fn = jax.grad( loss_ ) hvp = jax.jvp(grad_fn , (omega ,) , ( tangent_params ,)) [1] return hvpWe take the case of AFFD, and show how the single device code may be lifted to code employing model parallelism.First, we rely on the additional modules.from jax. sharding import NamedSharding from jax. experimental import shard_map from jax. sharding import Mesh from jax. sharding import PartitionSpec as P from jax. sharding import NamedSharding as NS from jax import lax from jax import tree_util as tu import numpy as npWe then define the device mesh; we assume 8 cores with 4-way model parallelism and 2-way data parallelism. mesh = Mesh( np. array ( jax. devices ()). reshape (2 ,4) , ( ' d a t a ' , ' m o d e l ' ,))We now lift initialization and application of the FlaX modules to model-parallel code: 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.4Roberta was fine-tuned with a batch size of 32 for 10k steps with Adam and a constant learning rate of 2 × 10 -5 . For the search algorithm 9 the learning rate was increased to 10 -4 , δ = 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 -4 ; δ Rouge1 = 0.5, δ Rouge2 = 0.5 and c = 2k steps; the total number of steps was increased to 40k.B.8 Hyper-parameters for Sec. 5.5We consider the SGD optimizer as in previous work [12] ; the batch size was 32, the learning rate set to 10 -5 .C.1 Definition of Higher order sketches For higher-order derivatives of L, one can consider sketches of operators. For example the Hessian vector product is the operator HVP : R N → R N given by HVP(u) = ∇ 2 L(θ)(u), i.e. HVP(u) i = j ∂ 2 i,j L(θ)u j . A sketch of the Hessian vector product can then obtained as follows: S(O)(v) = Φ(HVP(Φ T v)) where v ∈ R D so that we obtain an operator mapping R D → R D . Sketches of matrices were extensively studied [22] to speed-up evaluation of matrix products. Extending the HVP-example, for an operator O mapping R kN to R sN the transpose Φ T is applied to the k input indices and Φ is applied to the output s indices to obtain a mapping S(O) : R kD → R sD via:(S(O)v) l1•••ls = N t1•••ts=1 N i1•••i k =1 D j1•••j k =1 s β=1 Φ l β ,t β • O i1•••i k ;t1•••ts • k α=1 Φ jα,iα • v j1•••j k . (11)By the method of Johnson-Lindenstrauss [10] one can leverage the equation about concentration of the sketched norm (1) to prove that, given M points in R N , the distances between their sketches in R D are distorted by at most a multiplicative factor 1 ± ε. The point is that concentration arguments establish, for the δ appearing in (1), a bound of the form δ = O(exp(-ε 2 β 2 )): one can thus apply (1) to the M (M -1) 2 differences between points by requiring that β √ log M is sufficiently large; this is a considerable gain replacing a bound in terms of M 2 with one involving log M .The Fastfood paper [13] defines H N without scaling, so it is not an isometry; however we follow the definition with scaling as in Wikipedia so that H N is an isometry. Specifically, N needs to be a power of 2; then for i ≤ N let ∆(i) denote the vector of 0s and 1s and of length log 2 N , representing i in its binary form; then (H N ) i,j = 1 √ N (-1) ∆(i),∆(j) , where , denotes the inner product.C.4 Failure of concentration for FFD.Theorem C.1. There are some inputs x such that FFD(x) does not satisfy (1).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.e. generating random features that perturb the model parameters. The transposed operation of concatenation gives rise to a sum; specifically, given a unit vector x ∈ R N , we decompose it into N/D-blocks of size D, denoting the b-th block by x b . ThenEQUATIONthen choose x such that all blocks x b = 0 for b > 1 and x 1 is such that H(x 1 ) = e 1 , the first coordinate vector in R D . ThenEQUATIONwhere g 1 is the first entry of G v ; then we must have σ 2 = 1 and FFD(x) 2 cannot concentrate around 1 because g 1 has unit variance.[17] proposes two projection operators. The first one isEQUATIONwhereI is a vector of ones in R √ N and R i is Gaussian of shape D × √ N so that the memory cost of P ⊕ is O(D √ N ).The second one isEQUATIONwhere Q i is Gaussian of shape√ D × √ N so that the memory cost of P ⊗ is O( √ D √ N ).Our QK proposal is more general as it calls for a more general Kronecker decompositionEQUATIONwhereQ (i) is of shape D i × B i , i D i = D (reconstruction of the target dimension) and i B i = N (reconstruction of the model dimension). So if we choose K = 2, D i = √ D and B i = √N we recover P ⊗ . Strictly speaking, P ⊕ is different from P ⊗ , but one might reconstruct it by averaging two of our Q's (16), both defined with with K = 2: indeed, we select Q 1 where Q (1) is the vector I/ √ N as in (14) and Q (2) is of shape D × √ N ; we then select Q 2 where Q (2) is I/ √ N and Q (1) is of shape D × √ N ; in other words, to get P ⊕ we restrict the sampling of one factor to I/ √ N . Note that memory-wise our approach is more efficient than [17] as we allow K > 2; memory cost is O( K i=1 B i D i ) which, choosing B i A and D i 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 i × B i and obtain Q (i) using the QR-decomposition. Q (1) ⊗ • • • ⊗ Q (K)where Q (i) has shape D i × B i ; thenEQUATIONIn particular, as long as each D i 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/K -1) 2 ( ε K ) 2 .Proof.Step 1: Reduction to the case of applying a single factor in the Kronecker product representation of Q. We assume that Q decomposes asEQUATIONWe now look into applying (18) one step at a time. We reshape x to shapeEQUATIONa1,c ; we then contract it with Qb1,a1 over a 1 to obtain a matrix Yb1,c of shape(D 1 , B 2 • • • B K ). To apply Q (2)b2,a2 we need to first reshape Y (1) to shape (D 1 , B 2 , • • • , B K ), then transpose the first and second indices, and reshape it to a matrix X(2) a2,c of shape (B 2 , D 1 B 3 • • • B K ); we can then contract it with Q (2)b2,a2 over a 2 to obtain a matrix Y(2) b2,c of shape (D 2 , D 1 B 3 • • • B K ).It should be clear how this procedure can be continued for each i ∈ {3, • • • K}. Assume that for each i we can prove that the Frobenius norm (which is the l 2 -norm if we reshape it to be a vector) ofY (i) bi,cconcentrates around that of X (i) bi,c up to a multiplicative factor Di Bi :EQUATIONby conditional independence of the matrices Q (i) on one another we get thatEQUATIONwhere we usedK i=1 D i = D and K i=1 B i = N .Step 2: Using concentration of measure on the orthogonal group. The entries of Q (i) are not independent because of the orthogonality requirement and the fact that the rows need to have l 2 -norm equal to 1. We will employ measure concentration without independence; as a reference for notation and theorems we use [26] . From [26, 2.5.2] we recall the definition of the sub-Gaussian norm of a real-valued random variable X as:EQUATIONobtaining a bound on X ψ2 is equivalent to a concentration inequality:EQUATIONfor a universal constant c > 0. Note that Q (i) can be sampled on the orthogonal group SO(B i ) by restricting to the first D i -rows in the case in which D i < B i (by sampling from O(B i ) and changing in case the sign of one of the last B i -D i rows to ensure the determinant is 1), while the result we are proving is trivial ifEQUATION)where C is a universal constant and the Lipschitz constant f Lip is computed using the Frobenius norm on the tangent space. We now define:EQUATIONEQUATIONwhich has derivative:EQUATIONthe Cauchy-Schwartz inequality implies that:EQUATIONfrom which it follows that f Lip ≤ 1 if one assumes that X (i) 2 ≤ 1. Note that to derive (20) we may rescale x by a constant; if we rescale it so that x 2 = 1 then all the norms of the intermediate results X (i) 2 , Y (i) 2 are at most 1 because the matrices Q (i) are orthogonal. We have thus establishedEQUATION)Step 3: Replacing Ef (Q (i) ) with something simpler to estimate. A drawback of ( 27) is that the term Ef (Q (i) ) is not easy to estimate. However, using a symmetry argument, it is easy to estimate E(f (Q (i) )) 2 ; indeed the D i variables c (Y(i)bi,c ) 2 are identically distributed and if D i = B i one would get an isometry; soEQUATION)So we would like to replace Ef (Q (i) ) with E(f (Q (i) )) 2 ; the intuition why this would work is that concentration around the mean is equivalent to concentration around the median; so Ef (Q (i) ) concentrates around the median M i ; as f (Q (i) ) is non-negative, the median of (f (Q (i) )) 2 is M 2 i ; 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 [26, 5.1.13] ), we haveEQUATIONwhere the constant C might have changed but is universal. We then just need to show that| E(f (Q (i) ) 2 ) -M i | is O(1/ √ B i ).By the triangle inequality:EQUATIONand we can compute the right hand side with the layer cake decomposition:EQUATIONand we apply the concentration inequality (22) to the right hand side to getEQUATION)which implies that the right hand size is O(1/ √ B i ). We have thus establishedEQUATIONWe now deduce (19) conditional that it holds for j = 1, • • • , i -1 so that we may assume that X (i-1) 2 ≥ 1 2 ; then we may take t = ε i Di Bi X (i) 2 and get (19) withEQUATION)Step 4: Choosing the ε i . We just aim for ε i to be equal and that i (1 + ε i ) = 1 + ε; this is achieved by letting ε i = (1 + ε) EQUATIONwhere for each δ 1 > 0 we haveEQUATION)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 δ 1 will affect C in the third term); thus the effective bound for δ is of the form δ ≤ exp -Cε 2 D log 2 N .Proof. We will again use the notation and conventions from [26] . Let us recall the definition of AFFD:EQUATIONwithout loss of generality we will assume that x 2 = 1.Step 1: Using the pre-conditioner H 1 to distribute the mass of x. Note that H 1 is an N × Ndimensional matrix with entries of the form ± EQUATIONintuitively, the norm of each entry of H 1 B(x) cannot become much larger than a multiple of its variance 1 √ N : this is the purpose of using a pre-conditioner to distribute the mass of x.Step 2: Decomposing Φ(x) 2 2 . We now let u = H 1 • B(x) so that Φ(x) = σR D (H 2 • G v • u); we let Π be the permutation associated with rearranging the columns of H 2 so thatEQUATIONin (39) we decompose the effect of the diagonal and the off-diagonal terms obtainingEQUATIONWe now use the first term T 1 to compute σ:EQUATIONwhere we used the fact the components of G v have unit variance; as H 1 B(x) is an isometry, to have ET 1 = 1 we just need to set σ = N D .Step 3: Concentration for EQUATIONEQUATION)Step 4: Concentration for |T 2 |. We would like to claim that T 2 is small with high probability. This is the second point in which we use the pre-conditioner H; the intuition is that applying the pre-conditioner H to G v further reduces the finite sample correlation of the rows of the resulting matrix. Let us rewrite T 2 as follows:EQUATIONso we have reduced T 2 to a bi-linear form j,k B j,k g j g k ; by the Hanson-Wright inequality [26, 6.2.1] we haveEQUATIONwhere B F is the Frobenious norm of B and B S is the spectral norm. Let us look at B j,k conditional on E good :EQUATIONRecall now that D and N are both powers of 2 and the definition of H in Sec. C.3: if ∆(j) denotes the binary vector, of length log 2 N , representing an integer j ≤ N , we haveH i,Π(k) = 1 √ N (-1) ∆(i),∆(Π(k)), where a, b denotes the inner product of two vectors of length log 2 N . We thus obtain the bound:EQUATIONthe sum D i=1 (-1) ∆(i),∆(Π(k))+∆(Π(j)) is 0 unless the vectors ∆(Π(j)) and ∆(Π(k)) agree in the first log 2 D entries, otherwise varying i ≤ D we can always find two terms in the sum that cancel each other by flipping the parity of i in the first slot where the vectors ∆(Π(j)) and ∆(Π(k)) differ. So for each j, there are at most N D possible k-s such that |B j,k | is non-zero; moreover, as the sum D i=1 (-1) ∆(i),∆(Π(k))+∆(Π(j)) is at most D in absolute value, we getEQUATIONnote that a bound on the spectral norm B S is trivial from (47) as we can just take the maximum of the absolute values of the |B j,k | which is bounded byt 2 ∞ N .Given the stronger bound for B S and that ε 1, the minimum in the exponential in (45) is achieved by the term involving B 2 F and we thus obtain:EQUATIONStep 5: Picking up constants. Conditional on E good and on √ T 1 ≥ 1 2 we have:EQUATIONWe can bound the right hand side of (50) by 3ε if E good and the concentration inequalities for √ T 1 and T 2 hold. Thus, by decreasing the constant C by a factor 9, we haveEQUATIONQuestion: Does the paper provide open access to the data and code, with sufficient instructions to 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) is provided in Appendix B.Guidelines:• The answer NA means that paper does not include experiments requiring code. • Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.While 90% is common, this is a configurable parameter. 


Decision: 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.

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: 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.

1. 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. 

2. 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? 

3. 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.
questions: 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?

2. 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: limitations: The authors have appropriately discussed limitations.
paper_summary: paper_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.
strength_weakness: strength_weakness: Nothing major, but:
- Misses some important related work ([1], [2])
	- but gives good coverage otherwise
- Gradient inner product is not a good TDA estimate (many papers have found this now, e.g, [1][3])
	- Estimating inner product is fine, but be clear that's the goal then

[1] TRAK: Attributing model behavior at scale https://arxiv.org/abs/2303.14186
[2] Faithful and Efficient Explanations for Neural Networks via Neural Tangent Kernel Surrogate Models https://arxiv.org/abs/2305.14585
[3] Training Data Attribution via Approximate Unrolled Differentiation https://arxiv.org/abs/2405.12186
questions: 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.
- Some missing text, e.g., there seems to be no transition at the start of S4
limitations: limitations: none
paper_summary: paper_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.
strength_weakness: strength_weakness: * 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. 
* 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. 
* Table 3 is also unclear. Which algorithm should one use in practice?
* 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.
* 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.
questions: questions: *  Which components of the 6 contributions in the introduction are the main contributions of the paper?
* 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?
* In practice, which of the algorithms should one use and why?
* Could the authors define what the meaning of an influence score is for this paper?
limitations: 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.
paper_summary: paper_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.
strength_weakness: strength_weakness: - 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).
- 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).
- The paper would benefit from a more comprehensive comparative analysis, such as the sketching algorithm used in TDA [1, 2].

[1] Park, S. M., Georgiev, K., Ilyas, A., Leclerc, G., & Madry, A. (2023). Trak: Attributing model behavior at scale. arXiv preprint arXiv:2303.14186.

[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.
questions: questions: I have also included suggestions in addition to questions.

- 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.
- 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.
- The statement “memory constraints limited their investigation of ID in generative tasks to 500k” in line 98 is unclear. What does ID stand for?
- 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?
- 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.
- Do you anticipate any challenges in applying these techniques to even larger language models?

Minor Comments

- Line 129: Replace "D" with "$D$" for consistent mathematical notation.
- Line 150: Clarify that $B$ and $H_D$ now have dimensions $D \times D$.
- Line 172: Specify the relevant section when mentioning "unacceptable early TPU results" for clearer cross-referencing.
- Line 227: Add a period at the end of the sentence.
- 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?
- In line 311, it would be helpful if the authors provided a direct explanation of why explicit sketching provides substation speed up.
- Figure 1 has low resolution, making it difficult to read when printed. 

[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.

[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: limitations: The paper would benefit from a more comprehensive discussion of limitations.

Comments: title: Thanks for the Rebuttal
Response
Reply
comment: In light of the additional results and the comments above, I will increase my score. I am still concerned with the overall presentation, but I am hopeful the authors will be able to adjust this.
> 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.

Thanks. This would be helpful.

> 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.

Thanks for the answer.

> 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].

I agree with the authors that the theoretical guarantees for sketching and apply, **if** we separately consider the Hessian and the gradient. However, when sketching is applied to both the Hessian and train/test gradients simultaneously in influence computations, there would be an inherent information loss due to the whitening effect of the inverse Hessian---all components in gradients become equally important from the whitening effect of the inverse Hessian.

> 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

This is true if we only consider the short horizon. In the long horizon, however, small components that contribute minimally to naive dot product may still have a non-negligible influence on the overall training dynamics. 

> Would replacing P, F, and  with  throughout the paper address this concern?

I believe this would improve the readability.

Considering all these, I am willing to increase my score to 6. Thanks again for your comment.
I thank the authors for their reply and acknowledge reading their response. As the authors mentioned, it would be helpful to be more explicit about the TDA in the motivation of the paper (and give a more detailed description of why this also applies to quantities derived by preconditioning the gradient), add explicit comparisons to Trak's random projection (as the author did for the general rebuttal and the response to my review), and address other details the authors mentioned they will fix/add in the next revision of the manuscript. In general, it seems that most reviewers, including myself, had trouble following the notations used in the paper, which made the presentation of the work on the weaker side. Relatedly, while minor, I agree with reviewer 9kYE that it was difficult to follow all the acronyms in the paper (e.g. AFFD, AFJL, and QK), and it would be helpful to fix issues with \citep vs \citet. It would make the paper more clear if the authors could address these issues. Given these issues will be resolved, as the authors promised, I have increased my score (but it was difficult for me to give a higher score since the current process does not allow the reviewers to see the final manuscript).