Title: Efficient Sketches for Training Data Attribution and Studying the Loss Landscape
Venue: NeurIPS 2024
Accepted: yes
Keywords: Deep Learning

Language Models

Interpretability

Training Data Attribution

Loss Landscape
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

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.

1 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 ( NT ) 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 ( 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 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 ( ND ) 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:

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.

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

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:

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

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

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

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

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

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

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

Dense sketches and FJL 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 ( ε, δ ) :

This 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:

where: B ∈ R N × N is a diagonal matrix with B i,i = ± 1 (random signs); H N is the N -dimensional 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:

where 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 :

this 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 :

This 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. Table 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 s and Π introduce a lookup-based memory bottleneck on modern accelerators. This hinders performance, as seen in unacceptable early TPU results (compare Tab. 3). To address this, we propose randomizing the pre-conditioner H N . Efficient implementations of H N leverage the fact that it can be decomposed using Kronecker products; specifically, H AB = H A ⊗ H B , which allows a recursive multiplication by H N in O ( N log N ) -time 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

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 :

where 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:

This 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:

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

Theoretical results 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

x.vectortosketch

H:Hadamard

H:Hadamard

transform

AFFD

个

Gaussian vector

b:randlomSigns

H:Hadamard

<AFJL:drop thesecond Hadamard

transform

+

Reshapetomatrix

QK(2-clim case;Q1acts across

columns,Q2acrossrows)

Q1

Q2

for a universal constant C and for any δ 1 > 0 .

Finally, we analyze QK . Its structure allows for concentration arguments on orthogonal groups, leading to:

Theorem 3.3. QK satisfies (1) with

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.

4 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 glyph[ceilingleft] log 2 N D min glyph[ceilingright] 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 glyph[lessmuch] 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 ( √ D int ) ([1, Eq. 4], [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.

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 glyph[similarequal] 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.

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

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

5.2 Shortcomings of previous Training-Data Attribution scaling strategies.

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

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 ∼ 31 ms for token embeddings, resulting in a total of ∼ 4 s to project to R D . Other layers require ∼ 24 ms per projection to R k and ∼ 0 . 77 s 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.

5.3 Analyzing the Impact of Design Choices

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 .

5.4 Estimating the intrinsic dimension.

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 glyph[lessmuch] 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 .

(a)

(b)

4.0

BART(XSUM)

1.0

ROBERTA(SNLI)

ROBERTA(SNLI)

BART(XSUM)

GPT-2L(SNLI)

GPT-2L(SNLI)

3.0

0.8

RNEC

2.0

R

0.6

1.5

0.4

1.0

0.2

0.5

0.2

0.0

0.0

2000

4000

6000

8000

10000

Step

12000

1400016000

1800020000

40

n

60

80

100

5.5 Hessian Analysis During Pre-trained Language Model Fine-Tuning

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.

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

References

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 on Natural Language Processing (Volume 1: Long Papers) , pages 7319-7328, Online, August 2021. Association for Computational Linguistics.

Nir Ailon and Bernard Chazelle. The fast johnson-lindenstrauss transform and approximate nearest neighbors. SIAM Journal on Computing , 2009.

Gerard Ben Arous, Reza Gheissari, Jiaoyang Huang, and Aukosh Jagannath. High-dimensional sgd aligns with emerging outlier eigenspaces, 2023.

Krzysztof Choromanski and Francois Fagan. Fast nonlinear embeddings via structured matrices, 2016.

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

Jillian Fisher, Lang Liu, Krishna Pillutla, Yejin Choi, and Zaid Harchaoui. Influence diagnostics under self-concordance. In Francisco Ruiz, Jennifer Dy, and Jan-Willem van de Meent, editors, Proceedings of The 26th International Conference on Artificial Intelligence and Statistics , volume 206 of Proceedings of Machine Learning Research , pages 10028-10076. PMLR, 2023.

Roger Grosse, Juhan Bae, Cem Anil, Nelson Elhage, Alex Tamkin, Amirhossein Tajdini, Benoit Steiner, Dustin Li, Esin Durmus, Ethan Perez, Evan Hubinger, Kamil˙ e Lukoši¯ ut˙ e, Karina Nguyen, Nicholas Joseph, Sam McCandlish, Jared Kaplan, and Samuel R. Bowman. Studying large language model generalization with influence functions, 2023.

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.

Guy Gur-Ari, Daniel A. Roberts, and Ethan Dyer. Gradient descent happens in a tiny subspace, 2018.

William B Johnson and Joram Lindenstrauss. Extensions of lipschitz mappings into a hilbert space, 1984.

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 on Machine Learning , volume 70 of Proceedings of Machine Learning Research , pages 1885-1894. PMLR, 06-11 Aug 2017.

Shankar Krishnan, Behrooz Ghorbani, and Xiao Ying. An investigation into neural net optimization via hessian eigenvalue density. In ICML , 2019.

Quoc Le, Tamas Sarlos, and Alex Smola. Fastfood - approximating kernel expansions in loglinear time. In 30th International Conference on Machine Learning (ICML) , 2013.

Chunyuan Li, Heerad Farkhoor, Rosanne Liu, and Jason Yosinski. Measuring the intrinsic dimension of objective landscapes. In International Conference on Learning Representations , 2018.

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 Pattern Analysis and Machine Intelligence , April 2020.

Haokun Liu, Derek Tam, Muqeeth Mohammed, Jay Mohta, Tenghao Huang, Mohit Bansal, and Colin Raffel. Few-shot parameter-efficient fine-tuning is better and cheaper than in-context learning. In Advances in Neural Information Processing Systems , 2022.

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.

Michael W. Mahoney. Lecture notes on randomized linear algebra. Arxiv , 2016.

Sung Min Park, Kristian Georgiev, Andrew Ilyas, Guillaume Leclerc, and Aleksander Madry. Trak: Attributing model behavior at scale, 2023 (also accepted in ICML23).

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-Tien Lin, editors, Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual , 2020.

Levent Sagun, Utku Evci, V. Ugur Guney, Yann Dauphin, and Leon Bottou. Empirical analysis of the hessian of over-parametrized neural networks, 2018.

Tamas Sarlos. Improved approximation algorithms for large matrices via random projections. In 2006 47th Annual IEEE Symposium on Foundations of Computer Science (FOCS'06) , pages 143-152, 2006.

Andrea Schioppa, Katja Filippova, Ivan Titov, and Polina Zablotskaia. Theoretical and practical perspectives on what influence functions do, 2023 (also accepted in Neurips23).

Andrea Schioppa, Polina Zablotskaia, David Vilar Torres, and Artem Sokolov. Scaling up influence functions. In AAAI-22 , 2022.

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.

Roman Vershynin. High-Dimensional Probability . Cambridge University Press, 2018.

David P Woodruff. Sketching as a tool for numerical linear algebra, November 2014.

Chih-Kuan Yeh, Ankur Taly, Mukund Sundararajan, Frederick Liu, and Pradeep Ravikumar. First is better than last for language data influence. In Advances in Neural Information Processing Systems 36: Annual Conference on Neural Information Processing Systems 2022, NeurIPS 2022 , February 2022.

Contents

A Appendix: Additional Experimental Results

A.1 Additional results on layer selection

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.

A.2 Quality of sketches for inner products and eigenvalue estimation

For each algorithm, we report in Table 6 the minimal value of log 2 D necessary to reach a Pearson correlation > 0 . 9 x 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

for estimating inner products of gradients.

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.

A.3 A closer look at FJL vs AFJL

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: FJL 's cost significantly increases with D and does not scale beyond D = 2 20 .

12

FJL

B

AFJL

G

10

AFJL + FFT

memory

8

6

Peak r

4

8

10

12

14

16

18

20

22

D (log2)

FJL

160

AFJL

sw)

140

AFJL + FFT

Wall time

120

100

80

8

10

12

14

16

18

20

22

D (log2)

A.4 Compute cost of sketching gradients

In Table 8 we report the compute costs of sketching gradients in the setup of Sec. 5.3.

A.5 Compute cost of sketching HVPs

In Table 9 we report the compute costs of sketching HVPs in the setup of Sec. 5.3.

A.6 Search for the intrinsic dimension

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.

4.0

BART(XSUM)

ROBERTA(SNLI)

GPT-2L(SNLI)

3.0

G

RNE

2.0

1.5

1.0

0.5

0.2

0.0

2000

4000

6000

8000

10000

12000

14000

1600018000

）20000

Step

1.0

1

ROBERTA(SNLI)

BART(XSUM)

GPT-2L(SNLI)

0.8

0.6

R

0.4

0.2

0.0

0

20

40

60

80

100

n

Figure 6: ratio R of the n -th largest positive eigenvalue to the largest positive eigenvalue. We define outliers when R > 20% , motivated by [12, Fig.2]

A.7 Higher resolution version of Fig. 2

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

16000

TRAK[CUDA]

14000

TRAK[Triton]

3

12000

w)

AFFD[ours]

10000

QK[ours]

e

8000

6000

二

e

W

4000

2000

0

12

13

14

15

16

17

18

D (log2)

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

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 the projection Φ as follows:

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

B.3 Step-by-step didactic implementation

We describe a step-by-step didactic implementation in Jax and Flax.

First, we rely on the following modules.

The 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 next step is implementation of QK which is quite straightforward.

For using QK in implicit mode we need to transpose QK ; we illustrate how this can be carried out in Flax:

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

To transpose AFFD , we just need to reverse the above steps and transpose application of the Hadamard matrices.

' 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

Here's how one can do the same for the sketched HVP.

B.4 Sketching and model parallelism

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

We then define the device mesh; we assume 8 cores with 4-way model parallelism and 2-way data parallelism.

We now lift initialization and application of the FlaX modules to model-parallel code:

We now illustrate how to use the previous code.

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.

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 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 = 2 k 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 ; δ Rouge 1 = 0 . 5 , δ Rouge 2 = 0 . 5 and c = 2 k steps; the total number of steps was increased to 40k.

B.8 Hyper-parameters for Sec. 5.5

We consider the SGD optimizer as in previous work [12]; the batch size was 32, the learning rate set to 10 -5 .

C Appendix: Theory

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:

C.2 Guarantees on distorting distances.

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 .

C.3 Definition of the Walsh-Hadamard transform.

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

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

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

C.5 Comparison to Kronecker products in [17]

[17] proposes two projection operators. The first one is

where I 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 is

where 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 decomposition

where Q ( 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 glyph[similarequal] A and D i glyph[similarequal] 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.

C.6 Concentration result for QK: Theorem 3.3

Theorem C.2. Consider the projection algorithm QK where Q decomposes as Q (1) ⊗··· ⊗ Q ( K ) where Q ( i ) has shape D i × B i ; then

In 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 glyph[similarequal] ( ε 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 as Q (1) ⊗··· ⊗ Q ( K ) where Q ( i ) has shape D i × B i ; if we reshape x into a tensor of shape ( B 1 , · · · , B K ) indexed by ( a 1 , · · · , a K ) , the output Q ( x ) can be represented as a tensor of shape ( D 1 , · · · , D K ) indexed by ( b 1 , · · · , b K ) :

We now look into applying (18) one step at a time. We reshape x to shape ( B 1 , B 2 · · · B K ) obtaining a matrix X (1) a 1 ,c ; we then contract it with Q (1) b 1 ,a 1 over a 1 to obtain a matrix Y (1) b 1 ,c of shape ( D 1 , B 2 · · · B K ) . To apply Q (2) b 2 ,a 2 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) a 2 ,c of shape ( B 2 , D 1 B 3 · · · B K ) ; we can then contract it with Q (2) b 2 ,a 2 over a 2 to obtain a matrix Y (2) b 2 ,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) of Y ( i ) b i ,c concentrates around that of X ( i ) b i ,c up to a multiplicative factor √ D i B i :

by conditional independence of the matrices Q ( i ) on one another we get that

where we used ∏ K 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:

obtaining a bound on ‖ X ‖ ψ 2 is equivalent to a concentration inequality:

for 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 if D i = B i because then Q ( i ) is full-rank. We now invoke the concentration of measure for SO ( B i ) [26, 5.2.7]; if f : SO ( B i ) → R is Lipschitz:

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:

which has derivative:

the Cauchy-Schwartz inequality implies that:

from 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 established

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 ) b i ,c ) 2 are identically distributed and if D i = B i one would get an isometry; so

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 have

where 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:

and we can compute the right hand side with the layer cake decomposition:

and we apply the concentration inequality (22) to the right hand side to get

which implies that the right hand size is O (1 / √ B i ) . We have thus established

We 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 √ D i B i ‖ X ( i ) ‖ 2 and get (19) with

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 + ε ) 1 /K -1 . In this case we may lower bound ∏ i (1 -δ i ) by 1 -2 ∑ i exp( -4 CD i ((1 + ε ) 1 /K -1) 2 ) .

C.7 Concentration result for AFFD

Theorem C.3. Consider sketching with AFFD ; then

where for each δ 1 > 0 we have

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 :

without 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 × N -dimensional matrix with entries of the form ± 1 √ N ; each entry of the vector H 1 B ( x ) is of the form ∑ i s i b i x i √ N where the { b i } N i =1 are independent Bernoulli and s i = ± 1 ; applying Hoeffding's inequality [26, Thm.2.2.2] to each entry of H 1 B ( x ) we obtain that:

intuitively, 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 that

in (39) we decompose the effect of the diagonal and the off-diagonal terms obtaining glyph[negationslash]

We now use the first term T 1 to compute σ :

where 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 √ T 1 . If we regard √ T 1 as a function f 1 ( G v ) , we have

conditional on the good event E good that ‖ H 1 B ( x ) ‖ ∞ ≤ t ∞ √ N , this function is t ∞ √ N -Lipschitz in the l 2 -norm of G v ; applying concentration for the Gaussian measure on R N [26, Sec. 5.2.2] we get that conditional on E good

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:

glyph[negationslash]

so 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 have

where ‖ 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 :

Recall 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 have H 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:

the 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 get

note 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 by t 2 ∞ N . Given the stronger bound for ‖ B ‖ S and that ε glyph[lessmuch] 1 , the minimum in the exponential in (45) is achieved by the term involving ‖ B ‖ 2 F and we thus obtain:

Step 5: Picking up constants. Conditional on E good and on √ T 1 ≥ 1 2 we have:

We 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 have

where

having fixed a small δ 1 , if we set t ∞ = √ 2 log 2 N δ 1 , we get

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 we think the method that uses the Hanson-Wright inequality is more simple for this specific case.

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. 3 and Sec. 4; the experimental evidence is provided in Sec. 5.3.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in 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.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: Limitations are discussed in Section 6.

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 a 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 on 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 a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a 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 be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: Theorems are stated with the full assumptions in the paper; proofs are provided in Appendix C. 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.

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a 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

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

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

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If 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 is a specific model and empirical evaluation, it 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 is 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 a 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 all submissions 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 to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a 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 case authors are welcome to 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 to have some path to reproducing or verifying the results.

5. Open access to data and code

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

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) 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 run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) 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 a 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 in 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, 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.

7. 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 authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run 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 or 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 a 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 or figures symmetric error bars that would yield results that are out of range (e.g. negative error 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

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

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 as 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 it 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 https://neurips.cc/public/EthicsGuidelines ?

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 circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in 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:

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

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

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

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

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

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\n- 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.\n\nQUESTIONS:\nI 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.\n\nLIMITATIONS:\nThe paper would benefit from a more comprehensive discussion of limitations.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n* 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.\n\nWEAKNESSES:\n* 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.\n\nQUESTIONS:\n*  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?\n\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}

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

STRENGTHS:
- Clearly written and organized
- Addresses a practically relevant and timely technical problem (projecting high-dim gradients and HVPs) 
- Provides interesting new insights (e.g., that intrinsic dimension is not that low for LMs)

WEAKNESSES:
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:
- 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:
none
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-12

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nGradient 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.\n\nSTRENGTHS:\nI 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.\n\nWEAKNESSES:\nIn 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.\n\nQUESTIONS:\n1. 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?\n\nLIMITATIONS:\nThe authors have appropriately discussed limitations.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-13"}