Select or Project? Evaluating Lower-dimensional Vectors for LLM Training Data Explanations

TL;DR

This work tackles the impracticality of gradient-based instance explanations for large language models by comparing two dimensionality-reduction strategies: architecture-aware greedy selection of gradient components and geometry-preserving random projection. Using a novel retrieval-based benchmark, the authors demonstrate that a small, carefully chosen subset of layer components can outperform the full gradient and random projections in identifying training examples that influenced a given output, while also dramatically reducing computational cost. Key findings show that certain components (notably MLP gates/up projections) carry strong, task-relevant signals, whereas others (e.g., some attention keys/values) contribute little or even harm retrieval performance. The results argue for selecting informative architectural components over projecting the entire gradient, offering a practical path to scalable instance-based explanations with significant efficiency gains and broad applicability to data debugging and pruning tasks.

Abstract

Gradient-based methods for instance-based explanation for large language models (LLMs) are hindered by the immense dimensionality of model gradients. In practice, influence estimation is restricted to a subset of model parameters to make computation tractable, but this subset is often chosen ad hoc and rarely justified by systematic evaluation. This paper investigates if it is better to create low-dimensional representations by selecting a small, architecturally informed subset of model components or by projecting the full gradients into a lower-dimensional space. Using a novel benchmark, we show that a greedily selected subset of components captures the information about training data influence needed for a retrieval task more effectively than either the full gradient or random projection. We further find that this approach is more computationally efficient than random projection, demonstrating that targeted component selection is a practical strategy for making instance-based explanations of large models more computationally feasible.

Figures (6)

• Figure 1: Overview of the evaluation setup for the paraphrased setting. The gradient of a paraphrased sample $s_{p_i}$ is compared to gradients from its original counterpart $s_i$ and other candidates (e.g., $s_j$). The method succeeds at this retrieval task if the cosine similarity is highest for the original pair. The setup is identical for the generated setting.
• Figure 2: Component accuracy vs. parameter count. Each boxplot groups all components of a specific type, which share the same size. Left: Attention projections. Center: MLP components. Right: Embedding layer.
• Figure 3: Accuracy vs. layer depth for the model-generated setting. Each color represents a layer component.
• Figure 4: Greedy component selection (optimized for accuracy) vs. Random Projection. The x-axis shows the cumulative fraction of parameters utilized. In both settings, a small, greedily-selected subset outperforms the full gradient (dashed line) and random projection.
• Figure 5: Greedy component selection vs. random projection, optimizing for similarity to the full model gradient. While random projection quickly replicates the full gradient's geometry, this does not yield the best retrieval accuracy (cf. Figure \ref{['fig:greedy_vs_rp']}).