QKV Projections Require a Fraction of Their Memory
Authors: Malik Khalaf
Yara Shamshoum
Nitzan Hodos
Yuval Sieradzki
Assaf Schuster
Venue: ICLR 2026 Poster (2026)
Keywords: Memory Efficient Training
Pre-training
Finetuning
Approximate Matrix Multiplication
Compressed Activations
Primary area: foundation or frontier models, including LLMs
TL;DR: Significantly reduces QKV projection memory by leveraging Point-Approximate Matrix Multiplication (PAMM).

Abstract: The Multi-Head Attention mechanism is central to LLM operation, and multiple works target its compute and memory efficiency during training.
While most works focus on approximating the scaled dot product, the memory consumption of the linear projections that compute the $Q$, $K$, and $V$ tensors from the input $x$ is often overlooked.
To address this, we propose Point-Approximate Matrix Multiplication (PAMM), a novel tensor compression technique that compresses the activations of the $Q,K,V$ projections in attention layers by a factor of up to $\times 512$, effectively erasing their memory footprint, while achieving similar or better final perplexity. PAMM is fully composable with efficient attention techniques such as FlashAttention, making it a practical and complementary method for memory-efficient LLM training.

Published as a conference paper at ICLR 2026
QKV PROJECTIONS REQUIRE A FRACTION OF THEIR MEMORY
Malik Khalaf, Yara Shamshoum, Nitzan Hodos, Yuval Sieradzki, Assaf Schuster
Department of Computer Science Technion, Israel Institute of Technology
(malikkhalad fa yarlllaaaaarrrrrrreaeeaaeea yara-sh, hodosni tzan, syuvsie hwt.:/ssssrrcuuuuu.u.omomm
campus. technion....hiii.c.ccc
ABSTRACT
The Multi-Head Attention mechanism is central LLM operation. and multi- ple works target its compute and memory efficiency during training. While most works foeus on approximating the scaled dot product, the memory consumption of the linear projections that compute the Q, K, and V tensors from the input a is often overlooked. To address this, we propose Point-Approximate Matrix Multiplication (PAMM), a novel tensor compression technique that compresses the activations the Q, K, V projections in attention layers by a factor of up x512, effectively erasing their memory footprint, while achieving similar better final perplexity. PAMM is fully composable with efficient attention tech- niques such FlashAttention, making it practical and complementary method for memory-efficien LLM training.
Rows of
0 A B a
B
4
Σ lif(ss))
a,B, Rj
Rows afA
Compress
Approximate MacMuul
Figure 1: Illustration of Point-Approximate Matrix Multiplication (PAMM), PAMM approxi- mates the matrix multiplication O = A B for A E Bxx and B E Raxm in two stages. First, each row A, is represented by Å, defined as the closest point to A; on the line spanned by the best representative C,. Instead of storing the full matrix A. PAMM keeps only small set of k genera- tors E R together with an assignment mapping f IR.' and scaling factors
R Using these, PAMM approximates the product by first contracting into B E &x and then calculating Ô C B (equivalently, Ô AA" B). our setting, PAMM is applied approximate the gradi- ent of the projection matrices Q, K, V during backpropagation, e.g. VWQ = =XVQQ. Since the number of generating points {C, {C,)5=1 is typically very small, the memory required to store is drastically reduced.
INTRODUCTION
Large Language Models (LLMs) have become the foundation of modern NLP research, achieving state-of-the-al results across a wide range of benchmarks (Liu et al., 2019 Brown et al., 2020) et al. 2024a). However, as model sizes continue to increase, training them becomes increasingly

---

Published as a conference paper at ICLR 2026
challenging, especially when considering limited resources in terms of compute and memory. As result, numerous techniques have been developed to improve memory efficiency during training.
Since most LLMs are transformer- models (Vaswani et al. 2017). reducing the memory cost of the attention mechanism has been researched extensively, especially during training (Wang et al, 2020 Kitaev et al. 2020 Choromanski et al. 2022). These works often discuss extending context length, training with larger batch sizes, or reducing the time and memory complexity of the attention operation itself. However, when looking at the overall GPU memory consumption of the attention layer, component that is often overlooked becomes prominent: the memory requirement of the linear projection layers, which project the input x into the Q, V tensors used by scaled dot product attention. Since the imput æ is required for backward computation. it is saved by each layer during the forward pass, quickly accumulating to contribute as much as 20% of the total peak GPU memory required by the attention blocks.
High memory consumption by intermediate activations during training is not unique to attention projections. Recent work Shamshoum al. (2025) shows that activation memory far exceeds that of parameters or optimizer states in practical settings, since activations are the only component that scales with batch size and sequence length. This contrasts with the recent trend of focusing on
optimizer state compression to reduce memory consumption in LLMs (Zhao et al. 2024 Hao et al.,
2024 Muhamed et al. 2024 Zhang et al. 2024)."
These recent compression techniques use low-rank approaches that exploit redundancies in the hid- den dimension. In this work, we identify much more substantial source of redundancy: the quence dimension. We point out that activation tensors in transformer models are often highly redundant across tokens, due repeated patterns, padding, or local contextual similarity (Dao al.) 2022 Dao 2023 Krell et al. 2022 Dai et al. 2020; Yao et al.) 2022). This redundancy is ex- pected, as the dimensionality of each token representation is much smaller than the number of token embeddings in the activation tensor. We discuss this point in Section 3.1
To exploit this redundancy, we propose Point-Approximate Matrix Multiplication (PAMM), a sim- ple yet effective tensor compression method that stores only a small, representative subset of token activations to improve memory efficiency. The omitted tokens are approximated using scaled ver- sions of the stored subset. This process can be viewed as form of approximate clustering: rather than storing the full activations, we only save representative points. However, clustering token em- beddings on-the-fly is computationally prohibitive, as clustering algorithms' complexity is often quadratic in dataset size or worse (Jin & Han 2010). We show that, surprisingly. selecting represen- tative tokens at random is sufficient.
In this work, we apply PAMM to the Q, V projections of the attention block, motivated by recent works that show the activations of the attention layers in an LLM exhibit clustering behavior (Geshkovski et al.] (2024) and that the attention layer is much more amenable to compression than the FFN layers (Lv al. (024). We show PAMM is capable of reducing the memory consumption of these activations by a factor of x512, effectively eliminating their memory footprint during training. Importantly, PAMM leaves the forward pass and the gradients of other layers untouched, which accounts for its negligible impact on model performance..
PAMM is easily composable with most existing efficient training techniques, such as gradient check- pointing (Chen et al. 2016). low-rank adapters (Hu et al. 2021), 1), and most importantly, efficient potepddt techniques (Choromanski al., [022; Dao 2023). This makes PAMM practi- cal, easy-to-use plug-in for any model that uses scaled dot product attention.
In Section we introduce PAMM as general matmul approximation technique and present its com- plexity analysis and theoretical properties. In Section 4 we provide empirical evidence that PAMM drastically reduces the memory footprint of the Q. K, V projections while maintaining model per- formance, in both pretraining (4.2) and finetuning 4.3). We provide thorough throughput analysis in Section 4.4 and demonstrate that other compression methods don't scale nearly as well as PAMM
in Section 4.6
2

---

Published as a conference paper at ICLR 2026
2
RELATED WORK
Many recent works identify redundancies in the attention mecha- nism and exploit them to save time and memory during training. GQA (Ainslie et al.] (2023) is com- monly used in leading LLMs such as Mistral-7B (Jiang et al. 2023) and LLaMA (et al. 2024a to reduce time and memory consumption by sharing the key and value tensors across multiple queries. Linformer (Wang et al., 2020) achieves linear time and memory by applying learned linear pro- jections to the sequence dimension of the key and value matrices. Reformer (Kitaev et al.) 2020) leverages hashing to approximate attention in O(n log n) time. Performer (Choromanski et al. 2022] replaces softmax attention with positive random feature maps to kernelize the attention and achieve linear complexity. Longformer (Beltagy et al., 2020) combines these with a sliding-window local attention with a small set of global tokens, yiclding ((w)) complexity (where w is window size) that remains efficient for very long sequences. All of these works approximate the scaled-dot- product mechanism itself, and are therefore complementary to our work, which reduces the memory consumption of the linear layers that precede attention. In particular, (Choromanski et al. [022) ap- plies low-rank projections to the hidden embedding dimension, not the sequence dimension, which could be harmful model performance as it exploits less redundant axis.
Memory Efficient Attention
Low Rank Attention at Inference time Several works propose replacing the linear projections with low-rank approximations to reduce inference cost, including DeepSeek- V2 (et al., 2024b)
which introduced Multi- Latent Attention (MLA), Deepseek's 's Native Sparse Attention (NSA)
(Yuan et al. )025) and SinkLora Zhang 2024); these methods primarily target compression of the KV cache at inference. In this work, we focus on pretraining and finetuning, making these works mostly irrelevant to our discussion. ESPACE (Sakr & Khailany 2024) also employs a matrix- multiplication approximation, but with a fundamentally different objective: it introduces static pro- jection matrices during training to enable post-training weight compression. This approach does not reduce training memory and may even increase it, since the projection matrices must be stored. In contrast, PAMM directly reduces attention memory during training without modifying model weights or inference-time behavior.
Memory Efficient Training
Numerous works aim at improving the efficiency of LLM training without focusing on the attention mechanism, such as quantization methods (Pan al., 2022, Liu
et al.
2022 Dettmers et al.
2023 Xi et al.) 2024 Zhang et al. 2024), parallelism techniques (Huang
al.
2019| Shoeybi et al. 2020; Li al. 2014] Dean et al. 2012), mixed precision training (Mi-
cikevicius et al. 2017), and low rank approximation methods (Shamshoum et al. [2025] Zhao et al.)
2024 Hu et al. 2021 1 Lialin et al. 2023 Han et al. 2024). These approaches are either orthogonal
to PAMM can be applied in conjunction with it.
Approximate Backward Pass Various methods have been proposed to reduce memory consump- tion of LLM training by approximating the backward pass. (Adelman al. 202T Liu et al.,
2024) perform an approximate activation-gradient multiplication during backpropagation by ran- domly sampling rows and columns of the tensors. Approx-BP (Yang al. 2024 reduces memory during backpropagation by approximating nonlinearities and introducing memory memorr-sharinn normal- ization layers. CompAct (Shamshoum al. (025)) compresses the computation graph activations, gradients, and optimizer states using Gaussian random projections. Unlike CompAct, PAMM per- forms simple date-dependent compression that operates along the batch and sequence dimensions. This difference is key: by exploiting redundancy across the batch rather than the embedding di- mension, PAMM achieves substantially higher compression rates without degrading perplexity. We compare variations of these methods in Section 4.6
3 PAMM: POINT-APPROXIMATE MATRIX MULTIPLICATION
3.1 METHOD OVERVIEW
Consider a linear layer Z = X W, e.g. the Q projection of one layer of attention in an LLM. Here X E R nl.xn W E Raxm Z E BB.xm. where B is the batch size and L is the sequence length during training. For brevity, we denote the total number of tokens in a batch as b = BL. During backpropagation, receive the gradient of the loss w.r.t Z, denoted V Z. Backpropagation
3

---

Published as a conference paper at ICLR 2026
for this layer produces two gradients: VX = VZ W and VW - VZ. Thus, the input activations X must be stored during the forward pass for use in the backward computation. To mitigate this memory bottleneck, we introduce novel tensor compression technique, Point- Approximate Matrix Multiplication (PAMM). Instead of storing the full input X, PAMM stores a much smaller compressed representation and computes an approximate gradient VW during the backward pass as shown in 2
PAMM leverages the fact that the number of tokens in a training batch, b, is typically much larger than the hidden dimension n. For example, when pretraining LLaMA-1B, 2048. Consequently rank(X) S. n, meaning that the rows of X lie in much lower-dimensional subspace. In principle, one could represent using only n basis rows an 8x reduction in this example without any information loss. Accurately computing this aggressive reduction is costly, requiring a full Singular Value Decomposition, SVD(X) = UEVT, which is both tone-conuuiing and memory intensive, U Rxxn the same size X.
16384 and n =
PAMM provides an efficient alternative. We treat rows of X as points in II"" and compress X by finding closely matching points in IR" that generated from small subset of the original rows of X, used as the generating set, which we denote C. This allows us to only store the subset and two auxiliary vectors required generate the approximated X from C. The size of C' is k = b, where r is the compression ratio. Empirically, we find that extremely small values of r (as low as 1/512) maintain model performance while reducing the memory used by the input activation tensor to nearly zero.
We present PAMM for general matrix multiplication in Section
X bKi
W 3 RnK
Compress
Optimizer
Approximate Multiplication
VZ B
Figure 2: Overview training with PAMM. In standard linear layer, backpropagation multiplies the upstream gradient by the input X to compute the gradient with respect to the parameters VW X VZ. In our method, for the Wo, WK. Wv projections, we store compressed version of and approximate the gradient-input matrix multiplication with PAMM. Here, b denotes the total number of tokens in a training batch. Pseudocode provided in Appendix
3.2 ALGORITHM DESCRIPTION
Consider two matrices A Bexn and B E Rbxm their product O A B, and let A; E IR denote the i-th row of A. PAMM efficiently approximates O in two stages: (1) compressing A into a small set of generators and per-row coefficients, and (2) using this compressed representation compute an efficient approximation Ô. This is illustrated in Figure and full pseudocode is provided in Appendix A
Compressing A: Let {C,)}=1 denote of is generating points, or generators in R", which are the rows of matrix E R These can be any points "close enough" to the points in A. In our implementation, we obtain C by sampling & rows from A without replacement.
Given a point A,, we choose its representative as the closest point on the line spanned by any of the generators in C. For each generator Cj and point A,, denote the closest approximation point to
4

---

Published as a conference paper at ICLR 2026
A, on the line spanned by Cy
(A,, C3) ||C,1
A.,
C, c(,,))c,
Next, choose the closest approximation point A. A; its representative, with f(i) an assign- ment function: f(i) = arg min,I1AA -I,,A,,,,,le That for each point A, we choose the representative Á; = A1, = a(i, f()) Crt)) The assignment function f selects which generator Cj is used to generate A₁.
Note that a(i, j) csim(A,, Cj)-
L.Jle, where csim y) is the cosine similarity between x and [C
This points us to the following lemma, which we prove in Appendix B
Lemma 1. The generator representing A, is the one with the highest absolute cosine similarity to A, iee. f(i) arg max, Jesim(A,, Cj)l.
Lemma D alloow us to efficiently compute the assignment function f(i): we start by computing the cosine similarity matrix, csim(A, C) R which involves matmul AC and normalization over the rows. We then compute an arg max operation over the columns of (sim(A, C)) E h to obtain f(i).
Finally, to control the approximation quality, we introduce a tolerance value 2 and require:
|| A.ll2 S =||A,
(2)
We call this condition the neighborhood condition, as it requires the chosen representatives to lie an e-neighborhood around the points A,. the best representative we could find for A; doesn't fulfiil this requirement, simply drop A, in the final approximation. This is equivalent to choosing the zero vector E R" as A, representative or setting a; = O.
This allows us to replace the entire matrix A with set of generating points (C,) a list of indices = (f()) and list of coefficients a {a where a,
Approximate Matrix Multiplication: We wish to approximate O = -ATB. Given the com- pressed results of the previous stage C, f and a, the approximate matrix À 1gbx= can recon- structed as follows:
a; C5() A, satisfies Eq. 2
otherwise.
Reconstructing Ā would allow us to compute the direct estimate of the original product Ô ÅT B. However, this product would require as much compute as the original O. Instead, exploit the structure induced by the generators for efficiency. The approximate matrix multiplication can be written as a sum of the product of rank-one matrices A. E B=x1 and B, plxm, where each Å, can be expanded using the representatives in Eq.
Ô O
Σ
A B. Σ
C...
1=1
Since each generator Cj approximates several points, rewrite this sum
k Σ
Ô=
CJ
Σ :í((=jj
aaB
C.... By. #=1
As result, instead of reconstructing Ā R we can compute the matrix B Rkxm and compute Ō = C. B, which is much cheaper multiplication than the original O A B.
5

---

Published as a conference paper at ICLR 2026
Normalization factor 3: Since some
dropped a; 0. Ô underestimates the full sum in expectation. To mitigate this, introduce correction factor 3 into the computation of O such that [[[ O. Assuming the rows in A are drawn from some distribution A, PA over R" if we assume that row A, is kept its representative is close enough, e.g. A. Ai, then Ô can be written as:
M, f1, row i is kept. 0, otherwise.
Ô R 3
M.A Bi,
(4)
If 17 is the number of dropped rows, we have Pr[M, = 1| = b-n Hence in expectation:
b
b
b [IO] ~ 3 Σ EIMJA, = 3 - Σ A B. 3 70
(5)
Choosing 3 = b ensures [ = O. This 3 value is stored along with the rest of the compression b-
Note that requiring A; ~ A, implies using a small value for e, which means fewer points will satisfy the neighborhood condition, increasing 77-
outputs.
3.2.1 THEORETICAL GUARANTEES OF PAMM
Let Ze, = fi e (b]s(() = j} be the index set of all the rows A; whose representative Å, is generated by the point Cj. The union of these sets, I (
Zj is the set of all indices for which Å, satisfies the neighborhood condition, i.c. the indices of points that were not dropped. Denote these points as Ar, and all the dropped rows as Az. By construction of À we have
| - All
Σ
Σ
14
1.13
Σ
|1A,12
jE[LA iEI,
III
E Σ
Σ
|4
+
Σ
|| = = IA71
JE[LA]] FET,
As for the approximation error, due to submuliiplicatiity we get
Ô = (( - Â) B1} S [BI + 1|
This implies that the points without representatives could drastically hurt the approximation error. We analyze empirically the effect of the neighborhood condition on approximation error in Appendix
H
Next, we present a lemma for choosing the appropriate number of generators is to ensure all points in A are represented:
Lemma 2 (k bound under uniform sampling). Let A E gxxa and B E bbx Fix tolerance 0 < E < and a failure probability 0 < 6 < 1. Suppose we sample generators without replacement following uniform distribution the rows of A. The following bound on l is sufficient for the generating set C to fully cover A:
k
In
Pr[z [b]
1-6
>
Pmii
Here Pymin is the size of the smallest neighborhood of points close enough to any row A,, which represents the density of the data.
The lemma shows that as the total number of tokens in the batch = BL increases, we only need logarithmically more generators to capture the data distribution well. Mmai be shown to be proportional to b, i.e. b/Tmiix is approximately constant in b, which means that k > co In g for some constant cp. A proof of the Lemma 2 can be found in Appendix C including a proof that b/nmin is approximately constant in b.
This observation accounts for PAMM's strong performance even when & < In our experiments (Section 4h, compressing the activation by factor small 1/512 yields no loss in model per- formance and, in several cases, slightly improves the results when compared to training without PAMM.

---

Published as a conference paper at ICLR 2026
4 EXPERIMENTS
In this section, we evaluate PAMM on model pretraining (Section 4.2 and finctuning (Section 4.3 Throughout, PAMM applied to the Q. K, and V projections of every attention block in the traindd models. We demonstrate that even with an extremely small compression ratio down to 1 1/512
PAMM preserves
even improve performance across model sizes. In Section 4.4 we pro-
vide thorough evaluation of PAMM's effect on training throughput and runtime. In Section 4.6 compare with different compression algorithms,, such CompAct [Shamshoum et al. (2025)
and Uniform-CRS a sampling technique similar to (Adelman et al. [222]] Liu et al., 2024).. We demonstrate that these methods can't achieve the same memory reduction without degrading model performance significantly. All of the experiments were conducted on NVIDIA A 100 GPUs.
4.1 CHOOSING k AND IN PRACTICE
As discussed Section 3.1 PAMM leverages rank redundancy in the input tensor X used by the Q. K, and V projections. We parameterize the compression by ratio E (0, 1]: given b = BL rows for mini-batch size B and sequence length L we save = [r- generators. In pre- training. we push low 1/512; in some fine-tuning settings we reach k = 1, i.c. single generator C1. Also, we experimented with a wide range of с values, and found that setting + oc, Le., removing the neighborhood constraint altogether, gives the best performance. Another possible choice of is s O, which reduces PAMM to a Column-Row-Sampiee choice of column-row pairs, which we call Uniform-CRS. In Section 4.6 we delve deeper into these comparisons and the effect of &. We provide more visualizations of PAMM in Appendix H
algorithm with uniform
4.2 PRETRAINING
We evaluate PAMM pretraining by applying it the K,K, K, projections of LLaMA models TTovroo et al., 2023). from LLaMA-60M up to LLaMA-7B. All models were pretrained on the C4 dataset (Raffel et al., 2023) without data repetition. We provide full results for LLaMA-7B in Appendix proving its effectiveness in larger models.
Our experiment setup follows those of "Shamshoum et al. 2025 Zhao et al. (2024). where the learning rate is the only hyperparameter that tuned for each model size. For evaluation we report model perplexity at the end of training. For more details on the training experiments, we refer to Appendix D. and for stability analysis to Appendix
Poak Attenton on Kemory (ME)
PANN 1,101
307
10
371
10
2X1
Model Size
Model Size
(a) Pretraining perplexity across model sizes. We report validation perplexity for LLaMA-600 60M, 350M, and IB, trained with PAMM at different compression ratios, compared the full-rank baseline. Error hars black show minimal variance.
(b) Peak attention memory across model sizes. Memory corresponds to activations of all Q, K, V projection layers. PAMM reduces memory usage by over 97% all model sizes while preserving perplex- ity close to the full baseline.
Figure 3: Model perplexity and attention memory when pretraining LLaMA models C4 with PAMM. PAMM achieves massive memory reductions while maintaining or even increasing perplexity compared to the baseline. Results averaged over k=3 runs per configuration.
Figure 3 shows that applying PAMM to the Q, K, and V projections compresses the activations by a factor of 512x while leaving model quality essentially unchanged. For larger models, perplexity

---

Published as a conference paper at ICLR 2026
tually decreases relative to the full- memory baseline, suggesting that redundant rows in the attention inputs can hinder training. This observation supports the redundancy we anticipated, as discussed in Section 3.1 Exact measurement results are reported table in Appendix of the supplementary materials.
4.3 FINETUNING
For finetuning, we apply PAMM to the Q. K, V projections of RoBERTa-baid (Liu al. 2019) and evaluate on the GLUE benchmark (Wang et al. 2019). We follow the finetuning protocol of Shamshoum al. 2025
Zhao et al. 2024) and tune only the learning rate per task.
For all tasks we report performance by taking the average of 3 runs. Performance is measured F1 score for MRPC, Pearson's Correlation for STS-B, Matthew's correlation for CoLA, and classification accuracy for all other tasks. We measure the peak GPU memory consumption of all Q, K, projections the model in MB. For more details, refer Appendix G
Finetuning RoBERTa base on the GLUE benchmark with PAMM. We compare full fine tuning and PAMM with different compression rates. Reported memory in MB corresponds to the memory consumed by the activations of the Q. K. V projection layers. PAMM preserves the model's performance while reducing memory consumption by more than 97%.
Table I:
Memory (MB)
CoLA
STS-B
MRPC
RTE
SS12
MNLI
QNLI
QQP
Avg
Full Fine- Fine-Tuning
288
62.24
90,92
91,30
79.42
94.57
87.18
92.33
92.28
86.28
PAMM 1/128
6.75
63.17
90.90
92.12
78.5
93.84
86.64
92.17
91.579
86.11
PAMM 1/256
3.37
62.49
90.75
92.34
78.4
94.22
87.08
92.59
91.575
86.18
As shown in Table I PAMM drastically reduces the memory consumption of the Q, K, V pro- jections, by orders of magnitude, while achieving competitive performance across tasks and on average.
4.4 THROUGHPUT
In this section we provide empirical measurements for PAMM's effect on runtime and overall train- ing throughput. Overall, PAMM does not significantly reduce throughput, and the overhead is neg- ligible in the context of full model training, especially for larger models.
We measure throughput as the number of input tokens processed per second, per GPU, per full train- ing iteration. We use a batch size of 128K tokens averaged over 1000 training iterations, measuring throughput for PAMM with r 1/512 and baseline, on models of different sizes. For the models LLaMA-IB 1B and LLaMA-7B 7B we used DDP training with 8 GPUs, with the same global batch size of 128K tokens, meaning each GPU sees 16K tokens. Results are shown in Table [a The results show that PAMM's throughput degradation decreases significantly as the model size increases and remains below 2.7% for the IB model.
Table Throughput and runtime comparisons of models with and without PAMM.
(a) Throughput across model sizes. PAMM intro- duces minimal degradation, remaining below 12% for the 350M model and 2.7% for the IB model.
(b) Throughput comparison in tok/sec of forward (FP) and backward (BP) passes of LLaMA-IB. The final column reports runtime degradation rela- tive to baseline. The total runtime degradation very small.
PAMM (tok/sec)
Baseline (tok/see)
Throughput Degradation
Model Size
60M
122k
152k
19.74%
PAMM
Throughput Degradation
350M
181k
205k
11.71 %
Baseline
Total
1B
73k
75k
2.67 %
Forward Backward Total
247.6K
235.4K
4.92%
7B
20.05K
20.48K
2.1%
141.9K
138.3K
2.53%
88.4K
85.2K
3.61%
For LLaMA-IB we also report a detailed runtime breakdown of the specific operations PAMM in Appendix
8

---

Published as a conference paper at ICLR 2026
We that current implementation introduces only small runtime overhead (5.1% in FP, and 3.5% in BP). The additional compute required for PAMM is modest relative to the overall cost of training step. In short, this overhead is negligible in the context of full model training and quickly decreases with model size. An analysis of PAMM's time complexity provided in Appendix
4.5 ABLATION STUDY BATCH SIZE AND SEQUENCE LENGTH
We ran several experiments with different batch sizes and sequence lengths and measured model perplexity with and without PAMM. We used the LLaMA-60M model since we already demon- strated PAMM scales well to larger models. In all experiments, we used r = 1/512. Results are shown in Table 3 The results show that PAMM's perplexity is comparable with the baseline in all configurations, demonstrating that PAMM performs well regardless of the choice of batch size or sequence length.
Table 3: Training perplexity of Baseline PAMM across batch sizes and sequence lengths.
Batch Size
Sequence Length
Baseline Perplexity
PAMM Perplexity
Relative change
128
256
42.63
43.01
+0.8%
128
1024
37.47
37.03
-1.1%
256
256
37.61
37.25
-0.9%
256
512
33.32
33.73
+1,2%
512
128
37.38
36.43
-2.5%
512
256
30.97
32.46
+4.8%
512
512
29.16
30.30
+3.9%
4.6
ABLATION STUDY RANK AND E
In this section we answer two questions, using I.LMMA60MM model: (1) How does PAMM com- pare to other existing techniques that compress the activations of the Q, K, projections? (2) How does varying during training affect performance?
RAMM
opelor-05
Lnitarn
tgodorw1
Corppict
Epilorue
Baseline
Paplocty
12
30
14
3012
1050 1532
12
1/5
122
1/228
1012
(a) Benchmark comparison
(b) Effect of
Figure 4: (a) Compression technique comparison when applied to Q, K, V projections of LLaMA-6MM, with different compression rates. PAMM significantly outperforms other compres- sion methods and can utilize a very low r. (b) Effect of e on performance of LLaMA-60M with PAMM is equivalent to Uniform-CRS; when oc, PAMM doesn't 'tt enforce the neighborhood condition. We can see that choosing & = oo is the best option for all compression rates checked.
PAMM. When
We benchmark PAMM against CompAct (Shamshoum et al. 2025) and Uniform-CRi simple sub- sampling technique similar to (Adelman et 20011). which is equivalent to PAMM with ε as mentioned in Section 4.1 CompAct stores a low-rank sketch X = XP where E Raxk is random, & rb. and E PTP P| = I; gradients are unbiased but noisy. Uniform-CRS is the sim- plest form of PAMM: keeps only the sampled generators, discarding all others. Figure 4 shows perplexity versus memory shrinks. PAMM delivers up 512x memory savings with negligible

---

Published as a conference paper at ICLR 2026
perplexity change, whereas CompAct and Uniform-CRS suffer pronounced degradation, underscor- ing PAMM's superior memory -unntoyyy trade-off.
Finally, in order to choose the value of Results are reported in figure 4b The best accuracy obtained at E = oo, ic. when every row is represented by a generator; smaller values exclude rows and increase perplexity. The small gap between & = 1 and co indicates that attention inputs are already relatively clustered, which explains why dropping too many points harms performance.
measured model perplexity for different choices of E.
4.7 MULTI-MODAL AND PEFT COMPATIBILITY
We also finetune Pixtral- 12B Vision-Language Model (VLM) (Agrawal et al. 2024) the AID satellite image classification benchmark Xiia et al. 2017), with and without PAMM. This experi- ment demonstrates: (1) PAMM's scalability
larger models, (2) PAMM's compatibility with PEFT methods (I AL,RA), and (3) PAMM's applicability to multi-modal models. As we show in Table PAMM maintains performance while reducing non-negligible memory.
The model is given an image and textual instruction prompt asking it to produce final classifica- tion, standard in Visual Question Answering (VQA). Performance is evaluated using Fl averaged over the 30 possible classes. For the baseline, we finetune the model using LoRA with rank 128, applied to the language model and vision encoder. where LoRA(x) (Wo (x), Wo is the pretrained weight and A, B are learned low rank matrices. PAMM is combined with LoRA by com- pressing the input activation of the Q, K, as always, which is equivalent to applying PAMM to the A layer of the LoRA adapter. We note PAMM can also be applied to the B linear layer, but that would result in marginal memory gains as B stores a much smaller input activation.
Model
Macro F1
Weighted F1
Q, K, V Mem. Saved
Pretrained Bascline
0.51 0.9711 + 0.0028
0.54
+LoRA FT
0.9724 + 0.0030
0% 97. %
+LORA+PAMM 1/128
0.9727 + 0.0027
0.9737 + 0.0022
+LRR++PMM r = 1/512
0.9691 + 0.0015
0.9705 + 0.0015
99.28%
Table 4: Pixtral-12B VLM results on the AID satellite image classification benchmark with and without PAMM. We report Macro Fl (FI averaged across classes) and Weighted Fl (class- wrequency-weightted Fl average). PAMM at rank r 1/128 uses & 1/512 uses k 1 for both. We also report the
the vision encoder
and k in the language model, while memory saved (in MB and as percentage the Q, K, V activation memory required in standard LoRA finetuning). Each configuration is evaluated over runs compute standard deviations. PAMM effectively removes the Q, K, V activation memory cost during finetuning while maintaining competitive performance.
CONCLUSION
We presented PAMM, memory efficient technique, and demonstrated that it drastically reduces the memory footprint of the input activations to the Q, K, V projections in attention layers. PAMM is easy to implement, simple to use and easily composable with other efficient training techniques. Future work could attempt to extend PAMM approximate the activations of other layers in the transformer model, or to improve tradcoffs by choosing the generating set more strategically.
REPRODUCIBILITY STATEMENT
We provide anonymous downloadable source code in the supplementary materials. In the appen- dices, we include full pseudocode, detailed experimental settings and results, and complete theoret- ical proofs.
10

---

Published as a conference paper at ICLR 2026
REFERENCES
Menachem Adelman, Kfir Y. Levy, Ido Hakimi, and Mark Silberstein. Faster neural network training with approximate tensor operations, 2021. URL /ttps::/arxiv.oitt httg://b//180000 08079
Pravesh Agrawal, Szymon Antoniak, Emma Bou Hanna, Baptiste Bout, Devendra Chaplot, Jes- sica Chudnovsky, Diogo Costa, Baudouin De Monicault. Saurabh Garg, Theophile Gervet, So- ham Ghosh, Amélie Héliou, Paul Jacob, Albert Q. Jiang. Kartik Khandelwal, Timothée Lacroix, Guillaume Lample, Diego Las Casas, Thibaut Lavril, Teven Le Scao, Andy Lo, William Mar- shall, Louis Martin, Arthur Mensch, Pavankumar Muddireddy, Valera Nemychnikova, Marie Pel- lat, Patrick Von Platen, Nikhil Raghuraman, Baptiste Rozière, Alexandre Sablayrolles, Lucile Saulnier, Romain Sauvestre, Wendy Shang. Roman Soletskyi, Lawrence Stewart, Pierre Stock, Joachim Studnia, Sandeep Subramanian, Sagar Vaze, Thomas Wang, and Sophia Yang. Pixtral 12b, 2024. URL
arxiv. org/ http://aaaxxi.oooo/////////11
Joshua Ainslie, James Lee-Thorp, Michiel de Jong, Yury Zemlyanskiy, Federico Lebrón, and Sumit Sanghai. GQA: Training Generalized Multi Multi-Query Transformer Models from Multi-Head Check- points, December Utt.aa/pe///arat...oo/aaaa/////a [es].
13245 arXiv:2305.
Iz Beltagy, Matthew E. Peters, and Arman Cohan. Longformer: The Long-Document Transformer, December 2020. URL
arXiv:2004.05155 [cs].
Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhari- wal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler. Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Bemnrr, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners, 2020. URL https: //axxiv org/ abs 2005. 14165
Tianqi Chen, Bing Xu, Chiyuan Zhang. and Carlos Guestrin. Training deep nets with sublinear memory cost, 2016. URL httptttss///xiiii wrgg
Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, Afroz Mohiuddin. Lukasz Kaiser, David Belanger, Lucy Colwell, and Adrian Weller. Rethinking Attention with Performers, November 2022. URL http://arxiv. org/ abs abss2009..4.. 14794 arXiv:2009. 109.99999999944444 [es].
Zihang Dai, Guokun Lai, Yiming Yang. and Quoc Le. Funnel-transformer: Filter- ing out sequential redundancy for efficient language processing. In Hugo Larochelle, Mare' Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan-Tien Lin (eds.), Advances in Neural Information Processing Systems 33: Annual Conference on Neu- ral Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, vir- tual, 2020. URL bEpps://proceedings.ceem neurips copssccc/poe ct/ccc///p/aee hcc/paper/2220/hass/ 22221aa99660046542220065-------2 .httttttct.tttttll...
Tri Dao. Flashattention-2: Faster attention with better parallelism and work partitioning, 2023, URL
08691
Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory-efficient exact attention with io-awareness, 2022. URL htttps://exxvv.orrg
2205 14135
Jeffrey Dean, Greg Corrado, Rajat Monga, Kai Chen, Matthieu Devin, Mark Mao, Marc' au- relio Ranzato, Andrew Senior, Paul Tucker, Ke Yang, Quoc Le, and Andrew Ng. Large scale distributed deep networks. In F. Percira, C.J. Burges, L. Bottou, and K.Q. Weinberger (ods.), Advances in Neural Information Processing Systems, volume 25, Curran Associates, Inc., 2012. URL https://proceeddnggg. wtw.:/nnnnnniuiiin wwwpppeepppppppppppppppppeppeee. f files/paper/2011 11222
Tim Dettmers, Artidoro Pagnoni, Ani Holtzman, and Luke Zettlemoyer. Qlora: Efficient finetuning of quantized Ilms, 2023. URL wtt.itttpa//aaiii.orrrto http:///..g//aaaaaoooo//////////

---

Published as a conference paper at ICLR 2026
Aaron Grattafiori et al. The llama 3 herd of models, 2024a. URL URL.htoss//arxvv.orrr// 2407.21783
DeepSeek-AI al. DeepSeck-V22 A Strong, Economical, and Efficient Mixture-of- Experts Language Model, June 2024b. URL arXiv:2405.04434 [cs].
Borjan Geshkovski, Cyril Letrouit, Yury Polyanskiy, and Philippe Rigollet. The emergence of clusters self-attention dynamies, February 2024. URL arXiv:2305.054. [cs].
05465
Andi Han, Jiaxiang Li, Wei Huang, Mingyi Hong, Akiko Takeda, Pratik Jawanpuria, and Bamdev Mishra. SItraim: sparse plus low-rank approach for parameter and memory efficient pretraining, 2024.
org. abs/ 2406. 02214
Yongchang Hao, Yanshuai Cao, and Lili Mou. Flora: Low-rank adapters are secretly gradient compressors, 2024. URL httpt:tt:///ee http::://axxiv..ooog// http://wwww...rp///////........0 03293
Edward J. Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhi Yuanzhi Li, Shean Wang. Lu Wang, and Weizhu Chen. Lora: Low-rank adaptation of large language models, 2021. URL https:
09685
Yanping Huang, Youlong Cheng, Ankur Bapna, Orhan Firat, Dehao Chen, Mia Chen, Hy- oukJoong Lee, Jiquan Ngiam, Quoc Le, Yonghui Wu, and zhifeng Chen. Gpipe: Efficient training of giant neural networks using pipeline parallelism.
In H. Wallach, d'Alché-Buc, E. Fox, and R. Garnett (eds.), Ad-
H. Larochelle, A. Beygelzimer, vances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019. URL 7T11:993388888800333366677272222
Albert Q. Jiang, Alexandre Sablayrolles, Arthur Mensch, Chris Bamford, Devendra Singh Chap- lot, Diego de las Casas, Florian Bressand, Gianna Lengyel, Guillaume Lample, Lucile Saulnier, Lélio Renard Lavaud, Marie-Anne Lachaux, Pierre Stock, Teven Le Scao, Thibaut Lavril, Thomas Wang, Timothée Lacroix, and William El Sayed. Mistral 7B, October 2023. URL http: //tp:i/aaivvv../ org/ absaab//////// /2300 06825 arXiv:2310.066: [cs].
Xin Jin and Jiawei Han. K-Means Clustering, pp. 563-564. Springer US, Boston, MA, 2010. ISBN doi: 10.
978-0-387-30164-8.
425. URL httph:tt///eet /ttp://///o....//..../
424
Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The Efficient Transformer, Febru- ary 2020. RRLhttt::/// htttt:::/axxx...r////a.. 2000.200111110 04451 urXiv:2001.04455 [cs].
Mario Michael Krell, Matej Kosec, Sergio P. Perez, and Andrew Fitzgibbon. Efficient Se- quence Packing without Cross-contamination: Accelerating Large Language Models without Impacting Performance, October 2022. URL aattp://arxis wttp://aa//....02/203003 arXiv:2107.02027 [cs].
Mu Li, David G. Andersen, Jun Woo Park, Alexander J. Smola, Amr Ahmed, Vanja Josifovski, James Long, Eugene J. Shekita, and Bring Su. Scaling distributed machine learning with the parameter server. In Proceedings of the 11th USENIX Conference on Operating Systems Design and Implementation, OSDI' 14, pp. 583-598, USA, 2014. USENIX Association. ISBN 9781931971164.
Vladislav Lialin, Namrata Shivagunde, Sherin Muckatira, and Anna Rumshisky, Relora: High-rank training through low-rank updates, 2023. httpthtttts//iii/ee org/ abs/ 220./237...0 05695
Xiaoxuan Liu, Lianmin Zheng, Dequan Wang, Yukuo Cen, Weize Chen, Xu Han. Jianfei Chen, Zhiyuan Liu, Jie Tang, Joey Gonzalez, Michael Mahoney, and Alvin Cheung. Gact: Activation compressed training for generic network architectures, 2022. URL abs abs/2206. 22066 11357
12

---

Published as a conference paper at ICLR 2026
Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. Roberta: A robustly optimized bert pretraining approach, 2019. hRtpp://///axx
oww./rr///.....ooooo... 1907 11692
Zirui Liu, Guanchu Wang, Shaochen Zhong, Zhaozhuo Xu, Daochen Zha, Ruixiang Tang, Zhimeng Jiang. Kaixiong Zhou, Vipin Chaudhary, Shuai Xu, and Xia Hu. Wakner-take-all column row sampling for memory efficient adaptation of language model, 2024. URL httts:
bzg/abs/2333 20t.. 15265
Xingtai Lv, Ning Ding, Kaiyan Zhang. Ermo Hua, Ganqu Cui, and Bowen Zhou. Scalable Efficient Training of Large Language Models with ow-dimensional Projected Attention, November 2024.
URL http: wwwpsii...r...o...o. abs 2411 02063 arXiv:2411 104411002063 [cs].
Paulius Micikevicius, Sharan Narang, Jonah Alben, Gregory Diamos, Erich Elsen, David Garcia, Boris Ginsburg, Michael Houston, Oleksii Kuchaiev, Ganesh Venkatesh, and Hao Wu. Mixed precision training, 2017. (RLLattps///axxvvoooop
03740
Aashiq Muhamed, Oscar Li, David Woodruff, Mona Diab, and Virginia Smith. Grass: Com- pute efficient low-memory Ilm training with structured sparse gradients, 2024. URL https:
//ar //rrxii.orr//abs/226.111
Zizheng Pan, Peng Chen, Haoyu He, Jing Liu, Jianfei Cai, and Bohan Zhuang. Mesa: A memory- saving training framework for transformers, 2022.
11124
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. Exploring the limits of transfer learning with a unified text-to-text transformer, 2023. URL https: /ttp/v/rvvvv////////v////./a/av/ wttpv/orggabb/////// 10683)
Charbel Sakr and Brucek Khailany. Espace: Dimensionality reduction of activations for model compression, 2024. URL
http:////ab/241..0.533 05437
Yara Shamshoum, Nitzan Hodos, Yuval Sieradzki, and Assaf Schuster. CompAct: Compressed activations for memory-efficient L.M training. In Proceedings of the 2025 Conference of the Nations of the Americas Chapter of the Association for Computational Linguistics: Human Lan- guage Technologies (Volume 1: Long Papers), 2025. URL https::/aclattthloogg/ org/
2025 www.aalllloonn 717
Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catanzaro. Megatron-Im: Training multi- parameter language models using model par- allelism, 2020. URL httpttt/t//siiiii org/ abs / 009990000555.
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, Aurelien Rodriguez, Ar- mand Joulin, Edouard Grave, and Guillaume Lample. Llama: Open and efficient foundation language models, 2023. URL httptt:////aatt
org. /bsgg/bbs/// 222//333211111 13971
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is you need, 2017. URL httptst////axiii
org/abs/ 1706. 03762
Alex Wang. Amanpreet Singh. Julian Michael, Felix Hill, Omer Levy, and Samuel R. Bowman. Glue: A multi-task benchmark and analysis platform for natural language understanding, 2019.
URL http:httt:///aaetee
org/ abs/ 1804 07461
Sinong Wang, Belinda Z. Li, Madian Khabsa, Han Fang. and Hao Ma. Linformer: Self- Attention with Linear Complexity, June 2020. URL http:://axii..org/ab///0000
04768 arXiv:2006.04777 [cs).
Haocheng Xi, Han Cai, Ligeng Zhu, Yao Lu, Kurt Keutzer, Jianfei Chen, and Song Han. COAT: compressing optimizer states and activation for memory-efficient FP8 training. CoRR, abs/2410.19313. 2024. doi: 1044850/ARXIV.2410.1191. URL /tttps://dot /tts:///ooo...11....
48550/arXiv.222 2410 19313
13

---

Published as a conference paper at ICLR 2026
Gui-Song Xia, Jingwen Hu, Fan Hu, Baoguang Shi, Xiang Bai, Yanfei Zhong, Liangpei Zhang, and Xiaoqiang Lu. Aid: A benchmark data set for performance evaluation of aerial scene clas- sification. IEEE Transactions Geoscience and Remote Sensing, 55(7):3965- -99191,,6111,,, July 2017. http:tttp./dd......rr///..0....
ISSN 1558-0644. doi: 10. 109ht URL htt.htttpp//...... doi org wtrgg//000 10..........00.0.1....111....... TGRS.
Yuchen Yang, Yingdong Shi, Cheems Wang. Xiantong Zhen, Yuxuan Shi, and Jun Xu. Reducing fine-tuning memory overhead by approximate and memory-sharing backpropagation, 2024. URL
16282
Zhewei Yao, Xiaoxia Wu, Conglong Li, Connor Holmes, Minjia Zhang, Cheng Li, and Yuxiong He. Random-Itd: Random and layerwise token dropping brings efficient training for large-scale transformers, 2022. URL https: http::///axivv org/ abs /bs
Jingyang Yuan, Huazuo Gao, Damai Dai. Junyu Luo, Liang Zhao, Zhengyan Zhang, Zhenda Xie, Y. X. Wei, Lean Wang, Zhiping Xiao, Yuqing Wang. Chong Ruan, Ming Zhang, Wen- feng Liang, and Wangding Zeng. Native Sparse Attention: Hardware- and Natively Trainable Sparse Attention, February 2025. URL arXiv:2502.1 1089 [cs].
Hengyu Zhang. SinkLoRA: Enhanced Efficiency and Chat Capabilities for Long-Context Large Language Models, June 2024. URL arXiv.2406006678 [es].
Zhenyu Zhang, Ajay Jaiswal, Lu Yin, Shiwei Liu, Jiawei Zhao, Yuandong Tian, and Zhangyang Wang. Q-galore: Quantized galore with INT4 projection and layer-adaptive low-rank gradients.
CoRR, abs/2407.08299. 2024. doi: 10.45500/ARXIV.2400.002222 URL https::::/// org/
10. 18885/05aaii0 wwt.xivv.440... 2427
Jiawei Zhao, Zhenyu Zhang, Beidi Chen, Zhangyang Wang, Anima Anandkumar, and Yuandong Tian, Galore: Memory-efficient II training by gradient low-rank projection, 2024.
14

---

Published as a conference paper at ICLR 2026
A PAMM ALGORITHM
In this appendix, provide detailed pseudo-code for PAMM. Algorithm I describes PAMM as a two stage approximate matrix multiplication method given two input matrices E g E Ibxm; first, compressing A, and then approximating the product Ô. Algorithms and describe the forward and backward pass of linear layer with PAMM.
Algorithm PAMM: Point-Approximate Matrix Multiplication
1: procedure COMPRESS A RR k; 8)
2: Input: A E bx,, number of generators k, tolerance é 2 0.
3: Oufput: Compressed representation of
4: Sample indices I C [b] uniformly without replacement, [Z] = k.
5: C' A[Z,:]
generators, E R
6: Initialize 0: 4 09, f + 0
7: JA < ||A||] rows = norm A, dim= E 1x
A. Array row-norms of A
8:
nc <- IA E R
D Array of row-norms of C' (slice from array 114)
9:
csim( A, C) - A.C
g
> Division is element-wise
nAme
10:
f + arg max (csim( (A, C),, dim= [6] cA,m C)
11:
ne[s] RR
12:
= A - a a 2 =ll2 E R
>> Broadcast a along the column dimension n Neighborhood Condition via masked assignment
a/e 0
14:
return C, a,
15:
end procedure
I: procedure APPROXMM
2: Input: C, a, f. and B R
3: Output: Approximate multiplication Ô Raxm
4: Initialize + R
5: B' <- 3 B
Multiply each row of B with its value
6: B t index add( B,O, B*)
Efficient parallel computation of Σ [()=; 04 B
7:
0 E R Raxm
8: return
9: end procedure
Algorithm 2 Linear Layer: Forward Pass with PAMM
Input: An input X E BBLx, weight W E Raxm, and a rank r.
1: Z + XW E Bxxxm
2: C, J + (faptes(((atten()) (X),r) W)
3: save_for_ www.worrcccwwaaCC..omm
4: return Z
Algorithm 3 Linear Layer: Backward Pass with PAMM
Input: A gradient VZ E Bxmm A compressed representation of the activation C, a., f and weight W Raxm.
1:
2: VW -ApproxMMCC.OM f. (latttnn(())) VZ)) Raxm
3: VX t 72.W" BxL
4: return VX, VW
15

---

Published as a conference paper at ICLR 2026
B PROOF OF LEMMA 1
Proof.
arg arg ||A Ai.ll2
||
arg min ||A,, csim(A,, C₁)
|cyll2 Cylle
A
Cj
arg min
|Allle ||
csim( Cÿ)
||0/112 ll2
(8)
|Agll2
(A,, A)) 2csim Cj)
(Ai, Cj)
(C,, Cj)
arg min
||Aill2
+esim C2)²
|
|| 112 12/c
C
=osim| (A;,O)
arg min [114112(1 - cssm (A;, C,)3)}
(10)
arg max ccccccssss.. A,, Cÿ)]
(11)
j
For the last step, since the min operation is over j. ||A:]|2 is constant and can be removed. Also, the min becomes max because of the sign change.
C PROOF OF LEMMA
First, for x, y R", define h (. v) E IR" to be the point in Span(y) closest to x, i.e. h(x, y) (a.g) Tyl y. Note that in Section 3.2 we had ÅL, = A(A,, Cÿ)- Also define (i) = {A, ||A, (4),A))l2 S {|| to be the e-neighborhood of A,. In other words, it's the set of all points in capable of generating a representative for A; that will satisfy the neighborhood condition. Denote Mmin min|N, (i) the size of the smallest E-neighborhood in A. Note that nemin 2 1, A, can [b] j} is the
always generate itself with a We also remind the reader that I, {i index set of all rows whose representative is generated by Cj. and that I = U is the set of all indices which have representatives satisfying the neighborhood condition.
We now begin our proof.
Proof. Let be the set of generating points chosen uniformly at random from A, without replace- ment, and consider a row A, The probability that no C₁ generates representative for A, which satisfies the neighborhood condition is the probability of selecting all Cý outside of N=):):
Pr(i I PrivO, C. (i))
k-1 II ==0
(6- 1N-1N-00) R
(b
[Na(i)] 5)
II (b 8) ==0
k-1
N.(1)
4-1
[Ne(1)
==0
8=0
|Ne(1i)
IN)(1))
b
=
16

---

Published as a conference paper at ICLR 2026
Since 1 - e= for all real (1 [M((0] jk I exp( R|N.(8)) Thus:
Pr[i: i z S Pr(i $ I S cxp krmmin
Demanding the above probability sufficiently low:
k Romén
b mmin
b 5
b b.exp
< 6
7
k
In
b
Concluding our proof.
Proof of scaling of & In Section 3.2.1 we claim that is roughly constant with b, so k: scales like In(g). We will now prove this claim.
Time
First analyze the probability of some A₁ to be inside the e-neighborhood of another row As i.e. Aj E (i). The neighborhood condition ||A, - Aillz I e =||A;||2 defines an I2-ball around A; with radius A;|2: B(A,, A:lll))) Take the hypercone V(A,, 0) rooted at the origin which is tangent to this ball. The angle spanning the hypercone is simply 2 (r)sin(s), which isn't dependent on the specific choice of A, In order for any point æ R to also be in (i), there must be some constant e I such that is inside B(A (AA,, e alIAA (4:] therefore a must lie within the hypercone V (A,, 0), since it is the projective space of the ball B(A,, e||AA
2 arcsin [eIAA:l2)) [|A,
Next, denote the data distribution of rows A, with DA- The probability of sampling a row A, inside V(A,, 0) is Pr [A, E (AA, 0)[4, DA] = Assuming A, are i..d, when sampling b points from D the expectation of the number points inside V(A,, 0) becomes simply (
2-
Lastly, note that z; is a constant scalar value which doesn't depend on b; as b grows, even though more rows need to be covered, the probability of finding another row close the one we're consid- ering stays the same. This proves the claim.
To put it another way, assuming the data distribution given, and assuming that sampling longer sentences doesn't t inherently change that data distribution, we can expect that as the number of samples grows, so does the density of the data, and therefore the size of the e-neighborhoods all grows linearly with b.
Two small caveats to the above: 1. the A, aren't LLD along the sequence dimension they're context embeddings which are by definition dependent. However, even if the LLM can exploit this to adversarially "mess up" PAMM's approximation, e.g., by placing many embeddings with virtually neighbors, this doesn't 't happen in practice, evident by empirical success. 2. The claim above deals with E[N_(t) and not E Emmin, [i)]. However, the scaling of [E) the same for all even for the minimum
claim holds for Penin in particular.
17

---

Published as a conference paper at ICLR 2026
D PRE-TRAINING
As mentioned in Section
train LLaMA models TTourron et al., 2023 with 60M, 350M,
and IB parameters on the C4 dataset (Raffel et al., 2023 We follow the training setup outlined in previous works (Zhao al. 2024 Shamshoum et 2025). We train LLaMA-60M model for 10K steps, the LLaMA-350M for 60K steps and LLaMA- IB for 100K steps.
All runs maximum sequence length of 256 and global batch size of 512. We adopt the scaling factor = 0.25 and tune the base learning rate {le-3, 3, 3. -ie-3, 3, 8e - 3. 1e 2] separately for each model size. All weights are updated with 17, except those trained with PAMM, which use the reduced rate î = a for stability. Our learning rate scheduler employs linear warm-up over the first 10% of the training steps, followed by cosine decay that gradually reduces the learning rate to 10% of its maximum value over the remaining steps.
In Table we can see the perplexity for different ranks with PAMM.
Table 5: Pretraining perplexity and peak attention memory for different model sizes compared with the baseline. Reported memory corresponds to the memory consumed by the activations of all the Q, K, V projection layers. For PAMM measurements, this include the and f(-). PAMM reduces the memory by more than 97% while gaining 0.56 perplexity for the LLaMA IB model.
Model size
60M
350
1B
Perplexity
Memory
Perplexity
Memory
Perplexity
Memory
Full Rank
30.97
256 MB
18.80
1.5 GB
15.56
GB
PAMM r 1/128
31.94 1 0.008
MB
18.40 1 0.04
42 MB
15.01 1 0.04
78 MB
PAMM 1/256
32.18 1 0.03
5 MB
18.48 1 0.09
24 MB
15.06 1 0.02
42 MB
PAMM 1/512
32.53 + 0.05
3.5 MB
18.49 + 0.03
15 MB
15.36 + 0.04
24 MB
Training tokens
1.1B
6.4B
13.1B
D.1 PAMM AND FLASHATTENTION
All experiments use FlashAttention-v2 Dao 2023), a faster and more memory-efficient implemen- tation of the scaled-dot-product attention mechanism. FlashAttention avoids materialising the full attention-score matrix of size L L; instead, it stores only the operator's inputs and outputs to- gether with a few small work buffers. In particular, the output the scaled-do dot-product operation retained as part of Flash Attentonn 's activation set. However, this tensor is also the input and the activation tensor of the linear output-projection layer. Namely, it shared between Attennnt and the output projection, being saved in memory only once.
Therefore, When considering PAMM compression for the ouput-projection layer, several options arise:
1) Compress the shared activation itself. While memory efficceen,,,,,,,,, altering this shared tensor perturbs the gradient flowing into PlashAttention, introducing error accumulation through layers and to the residual stream effect we do not study in this work.
2) Compress separate copy kept by the output-projection layer. This increases memory usage, because the original (uncompressed) activation is still stored by Flash
3) Remove FlashA ttention completely so it does not retrain this tensor. This largely in- creases memory usage, as the naive implementation of attention is far more memory inten- sive, even if we apply PAMM for the output projection layer.
Thus, to avoid memory increase or performance degradation, we leave the untouched and defer joint treatment of FlashAttention and PAMM compression to future work.
layer
18

---

Published as a conference paper at ICLR 2026
E LLAMA-7B RESULTS
We pretrained a LLaMA-7B model with and without PAMM. using r = 1/256 and r r= -1/512, and report the models' perplexity in Table 6 The training setup follows previous works (Zhao et al.) 2024). We train the model for 150K steps, with global batch size of 512 and sequence length of 256, on GPUs. A scaling factor a 0.25 was used with learning rate ny E (1e 7e 4).
40K
80K
120K
150K
PAMM-255
17.73
14.93
13.91
13.81
PAMM-512
17.53
14.62
13.65
13.57
Baseline
18.09
15.47
14.83
14.61
Table 6: Pre-training LLaMA 7B on C4 dataset for 150K steps. without PAMM.
We report perplexity with and
F LLAMA-ID IB RUNTIME BREAKDOWN
We now report a breakdown of the runtime of the different operations in PAMM's forward and backward passes using LLaMA-IB model, well the runtime of the same model without PAMM. We use the same batch size used in Section 4.4 and report time in ms, as well as the operation's 's relative contribution the total iteration's 's runtime (forward backward + optimizer step). We also report the relative contribution of the appropriate phase (forward or backward). Results are Table and Table
Operation
Time (ms)
% iteration of forward
PAMM forward total
101
6.7%
19.1%
Forward pass matmul
53
3.5%
10.0%
Index selection
12
0.8%
2.3%
Normalization
22
1.5%
4.2%
Cosine matmul
8
0.5%
1.5%
Max/assign
3
0.2%
0.6%
Table 7: PAMM forward pass breakdown for LLaMA-7B. Times are averaged over 1000 update steps. Percentages are reported relative to the full training iteration and to the total forward time.
Operation
Time (ms)
of iteration
% backward
PAMM backward total
149
9.9%
15.8%
Input grad matmul
51
3.4%
5.4%
Index gathering
21
1,4%
2.2%
Alpha scaling
17
1. %
1.8%
Matmul
57
3.8%
6.0%
Table 8: PAMM backward pass breakdown for LLaMA-7B. Times are averaged over 1000 update steps. Percentages are reported relative to the full training iteration and to the total backward time.
We emphasize that the Argmax operation is implemented efficiently in parallel using tree reduction. The time of the random sampling (Index Selection) is negligible less than 1% of the total forward time shown in Table 7 We also experimented with reusing generators, but that produced no speed gain, and introduced small perplexity degradation. Therefore, we opted for the default choice of per-step sampling.
19

---

Published as a conference paper at ICLR 2026
G FINETUNING
Evaluation For the GLUE benchmark Waan et al. 2019), the metrics reported are F1-score for MRPC, Matthew's correlation for CoLA, Pearson's correlation for STS-B, and accuracy for the rest of the tasks, along with the peak memory consumed by the activations the three linear projections. Baseline scores are taken from CompAct (Shamshoum al. 2025 For every task, average results over three random seeds.
Training Setup We trained the various models for 30 epochs using batch size 16, except for CoLA where we used batch size 32. We perform hyperparameter search once on RTE for the optimal PAMM scale factor a for two ranks,
{1/128, 1/256]. The sweep a E {1, 2, 4, 8) yields
for 1/128 and a 2 for 1 1/256; these settings are fixed for all other tasks. For the learning
rate we consider {le 5, 2e - 5, 3e - 5} for every task, and consistently found that = le worked best.
We note that for some tasks, such as CoLA it happens that B < 128. Thus, following our method definition, such iterations used & = [r B L7. These provide empirical evidence for the effectiveness of PAMM even with an extreme = 1.
H PAMM COMPRESSION EDA
PC1
PCL
(a) Rows of Xx, colored by f(i)
(b) Rows of Xx₃ colored by f(i)
Figure 5: Two-dimensional PCA visualization of PAMM approximation. (a) rows of the input tensor of the K projection of layer LLaMA-60M. XK3- The rows are projected using PCA and colored by their assigned generating point, i.e. by f(i), using PAMM with = co. (b) The representative rows X shown in the same PCA space, colored by f(i). We can see that our method approximately clusters the input data, and transforms clusters into lines.
In this appendix we examine the behaviour of PAMM on real data taken from one of our LLaMA- 60M pretraining experiments, with co. We chose use the input activation the projection K at layer 3, at training step 3000. We write XK, E f in this section to denote this data tensor. For LLaMa-60M. 5 65536, n 512.
In Figure 5 we plot the first 2 principle components of the PCA decomposition of K3. We run the first stage of PAMMM selecting i generating points {C and color each point XK, according its chosen generator, i.e. according to f (t). Next, we plot the PCA of the representative points Kg, using the same PCA projection from the first step. We use the same coloring scheme. As can be seen, PAMM's compression approximately clusters the input data; close points share the same generator. We also see that the representatives of each cluster lie on the same line, which is expected from our construction of Á, = a,Cr Finally, notice that the variation in the first two principle components is mostly preserved for the whole set of points, while each cluster's 's variance is greatly reduced. Since our perplexity results were very good, this indicates that indeed most of the lost variance redundant.
20

---

Published as a conference paper at ICLR 2026
Next, measure the estimation error of PAMM for the same layer. In addition to the input tensor XK3- we also look at VK3, the gradient of the projection recorded during the backward step of the same training iteration. We compute the exact weight gradient VW K2 XK, V K3, and also
compute the PAMM approximation VW K3 = X Kg VK3. To assess how well PAMM approximates matrix multiplication, we measure the relative L2 error:
e( VWKx, VW TWKglFF
2-0.
C-0-205
0-125
KID.5
0.5
Relative Distance brror
0.928
3.0
Distnnee
Epolion
Rank
(a) Effect compression ratio on L2 error.
(b) Effect of E on L2 error.
Figure 6: Relative L2 error analysis. We compare the relative Le error E(r, E) for multiple com- pression ratios and with different epsilon values, including æ = oc which was used in our main experiments. The relative L2 errors calculated on the parameter gradients VW Kg of LLaMa- 60M model after 3K training steps. We can clearly see that the error decreases as é increases, indicating that oo the best choice this parameter. We also observe that the errors scale only logarithmically with which explains why we could reduce so drastically in our experiments.
Notice that E depends on our selection of compression ratio and neighborhood parameter We repeat this mcasurement for various values of and æ and present the results in Figure 6 Addi- tionally, we also measure the coverage of PAMM for each E, and present the results Figure
0.0625
0.125
0.8
Coverage
0.2
0.29
Figure 7: PAMM Coverage: Coverage the ratio between the number of points that have rep- resentative, i.e. their a; * 0, to the total number of points in the tensor. Coverage is reported for several compression ratios r and epsilon values s. As expected, coverage rises with larger e values, points satisfy the neighborhood criterion (Equation 2 Likewise, higher values increase coverage by introducing more generator points, increasing the number of rows for which there is a generator satisfying the neighborhood condition, and representing greater portion of the data.
We have several key insights from the results. First, as increases, and more points are covered by the approximation (Figure 7 the relative emor drops (Figure 6a). This further supports our results
21

---

Published as a conference paper at ICLR 2026
that as the relative error scales only logarithmically with r (Figure 6b). This shows that can indeed become very very small without significantly reducing performance, the enror scales up relatively slowly. Third, for & = 0, i.e. for Uniform-CRS, the error is much higher compared to g = 0.2. This further strengthens results from Section where show PAMM outperforms Uniform-CRS.
from 4.6 which indicate that choosing oc is the best option. Second,
Lastly, we note that the relative errors we measured, even for & = oo, are on the order of 1/2 to 1. These are large relative emror values; this means that for the data distribution present in real transformers, PAMM isn't It a very good approximation in terms of relative norm error. However, the perplexity we achieved practice was very close to the baseline, indicating again that the training process of LLMs is highly redundant. Even though PAMM doesn't 't approximate the full gradient very well, the model is still perfectly able to learn well as without any compression.
I LOss CURVES
Figure 8 compares the pretraining loss curves of LLaMa-IB with and without PAMM verify the stability of our method and its effect on training dynamics. We can see the resulting curves exhibit nearly identical behavior, and that PAMM maintains smooth and stable training, asserting that PAMM does not negatively affect optimization. This stability was consistent across all three random seeds.
(a) Baseline.
(b) PAMM with 1/512.
Figure 8: Pretraining Loss Curve Comparison. We validate PAMM's stability by evaluating on multiple seeds and comparing it to the baseline LLaMa-I pretraining in the same setting, The loss equivalent to model perplexity on the C4 dataset.
J COMPLEXITY ANALYSIS
We now compare the memory and time complexity PAMM with the exact matmul operation.
Memory Complexity The original matrix A requires storing bn scalars. PAMM only uses kn + 2b O(k)) scalars to represent Á, made up of the matrix C' E R*xn and the two vectors E [k] and E IR.. We show that in practice & < b. reducing the total memory required to represent Á to a very small fraction of the original memory.
Time Complexity The number of scalar multiplications in the full matrix multiplication O A B is bam. During compression, PAMM computes the scalar product of every A, and Cy. i.e. we compute ACT which has bkn multiplications, We also compute the squared norms of A, which require an additional bn multiplications. Finally we perform arg max operations which can be efficiently computed using parallel reduction algorithms in log(k) time. For decompression, in order to compute we have bm multiplications, and to compute Ô C ß we have knm multiplications. total, the number of multiplications used PAMM O(bkn + knm bn + bm bmbb-loggk)) = ((kn(b+m) Comparing this with the original brm we see that the speedup ratio bm
3(++mj determines if PAMM will be faster slower than exact multiplication: if 1, PAMM will be faster. This speedup highly dependent model size and parameter choice of ht and in
22

---

Published as a conference paper at ICLR 2026
experiments, had y values as high as 28 for LLaMa-IB during pretraining, using is b/256. However, the matmul operation PAMM is used to approximate in our work is only a small part of the entire runtime of the model, so the total training runtime is negligible for large models. Also, the theoretical speedup doesn't t account for implementation details, including memory movement the GPU between the GPU's RAM and the SMs doing the actual computation, as well adding several CUDA kernel launches. In practice, as model size increases, the total overhead incurred by PAMM becomes negligible. We give a detailed analysis of training throughput in Section 4.4
23


Reviews:
[{"review_id": "Tr7qUgv6pl", "reviewer": "Reviewer_LvZx", "summary": "This paper proposes PAMM (Projected Attention Memory Mapping), a memory-efficient training framework for Transformer-based architectures.\nThe method introduces learnable low-dimensional projection matrices for the Q/K/V representations, allowing the model to compress activations during training while maintaining recoverable gradient information through a projection-reconstruction scheme.\nThe approach achieves up to 97% memory reduction in attention activation storage with negligible performance loss, and generalizes well across multiple Transformer variants and tasks.", "questions": "See my Weakness", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper presents a clear motivation and strong practical relevance — training large Transformers is often constrained by activation memory, and this work directly addresses a crucial scalability challenge. The proposed method introduces learnable projection mappings that preserve gradient reconstructability, an elegant and theoretically grounded idea that distinguishes itself from prior techniques such as activation checkpointing or reversible layers. The empirical results are convincing, demonstrating consistent memory reduction with comparable performance across language modeling and vision benchmarks. Importantly, there is no additional computational overhead during inference, as PAMM operates without decompression or re-materialization, making it highly practical for deployment. Theoretical reasoning is also solid — the authors provide analytical evidence that the projection mapping retains full-rank subspace properties, explaining why gradient information is preserved.", "weaknesses": "While the proposed PAMM framework is conceptually sound and practically motivated, the experimental evaluation is somewhat limited. The current experiments are mainly conducted on mid-scale Transformer models, and it remains unclear how the method scales to very large architectures (e.g., >1B parameters) or longer sequence settings. Moreover, the ablation analysis on the projection dimension ratio and per-layer compression strategies is rather sparse—more systematic studies could strengthen the empirical conclusions.  Finally, while results across language and vision tasks are encouraging, broader validation on other modalities or architectures would further demonstrate the generality of the approach.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ceY6t5ktQj", "reviewer": "Reviewer_qFLe", "summary": "This paper tackles the less explored problem of reducing activation memory during LLM's training. The authors propose Point-Approximate Matrix Multiplication (PAMM), a simple technique that compresses the input activations of the Q, K, and V projections in attention layers using a small set of representative rows and scaling factors. This allows approximate gradient computation without storing full activations. Experiments on LLaMA models (60M–1B) and RoBERTa-base show that PAMM reduces activation memory by $512 \\times$ with little to no loss—and sometimes improvement—in perplexity and task accuracy. The method introduces only a small training throughput penalty and is compatible with existing efficiency techniques.", "questions": "* Do similar sequence-level redundancies exist in pretrained LLMs during inference. Does the degree of redundancy change over the course of training?\n\n* Could the authors report the total training-time memory consumption (including all activations, parameters, and optimizer states) for both the baseline and the PAMM-compressed models?\n\n* The current experiments use relatively short sequences (256 tokens) and moderate batch sizes. How does PAMM’s computational cost scale as sequence length or batch size increases?\n\n* How sensitive is PAMM to the shape of the training batch—for example, using fewer long sequences versus more short ones with the same total number of tokens? Does this affect model accuracy or compression effectiveness?\n\n* Can PAMM be applied to other activation tensors beyond the QKV projections, such as feed-forward or output-projection layers? How much redundancy do those tensors exhibit in comparison?", "limitations": "", "rating": 4, "confidence": 5, "soundness": 2, "presentation": 4, "contribution": 2, "strengths": "* The paper is well written, clearly structured, and easy to follow.\n* The proposed method is supported by solid theoretical analysis.\n* The approach effectively reduces training-time memory consumption while maintaining, or even slightly improving, model accuracy with minimal degradation.", "weaknesses": "* **Limited profiling on sequence redundancy:** The paper offers initial empirical evidence of sequence‑axis redundancy (Appendix F: PCA clustering; relative error and coverage), but the analysis is confined to a narrow slice (one layer/model/step). A broader study would better ground the motivation.\n\n* **Missing complexity/scaling analysis:** While the runtime breakdown is helpful (Table 2), the paper lacks an explicit complexity analysis of the compression and approximate multiply and a scaling study with respect to batch size, sequence length, and model dimension.\n\n* **Narrow evaluation scope:** Pretraining results are limited to a specific setup, with sequence length 256 and batch size 512. It is unclear how PAMM performs under different training regimes (e.g., smaller batch but longer sequences).\n\n* **End‑to‑end memory not reported:** Memory savings are reported for QKV activations only. End‑to‑end peak memory (including other activations, parameters, and optimizer states) and a memory timeline would better quantify the overall training benefit.\n\n* **Throughput comparability:** Throughput is reported for single‑GPU batches of 16K tokens (Section 4.4), whereas pretraining uses a global batch of 128K tokens (Appendix D). Aligning these settings or clarifying per‑GPU vs. global configurations would strengthen the runtime claims.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "JnJyx0u32c", "reviewer": "Reviewer_i9xN", "summary": "PAMM proposes to approximate GEMMs using a sequence of point-wise multiplications. These are obtained by projecting one of the GEMM operands into lines spanned by a small catalog of points randomly sampled from the rows of said operand.  Each point in the activation tensor is represented by its closet point from a generator, under a proposed \"neighborhood condition\".", "questions": "No questions, but please address the above.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "- The proposed method identifies the redundancy in the sequence dimension, which is typically enormous in modern LLM training.\n- Outperforms other methods like CompAct and Uniform-CRS in memory-performance tradeoffs.", "weaknesses": "- the experiments are limited to small to medium size LLMs. Scaling up models can introduce performance degradation of the approximation method. \n- Activations are dynamic during training. Random sampling is slow, particularly on GPUs. Note that random sampling requires a pseudo-random number generator, such as an LFSR or others. Most algorithms in this family are sequential and model generation from an irreducible in a Galois field. The authors should check the intrinsic details of cudaranddx. Can we just sample once and stick with the choice of catalog points Cjs, or perhaps periodically refresh? \n- The algorithm requires a matmul followed by an argmin on one of the dimensions. This means a reduction over one of the reduction, this is slow on a GPU.\n- The breakage of computation into a sequence of small GEMMs can harm tensor core utilization.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "rEg6yItJ0O", "reviewer": "Reviewer_16on", "summary": "This paper proposes Point-Approximate Matrix Multiplication (PAMM) to reduce memory consumption of the QKV projections, by a factor of up to 512x, thereby improving memory footprint while achieving similar PPL. PAMM can be composable with other efficient attention techniques such as FlashAttention and LoRA, making it practical and complementary for memory-efficient LLM training.", "questions": "My questions and suggestions are basically from \"Weaknesses\" as aforementioned.\n1. From Weakness 1: If my opinion is correct, PAMM can save up to (and at most) 20% of peak memory usage.\n2. From Weakness 2: Please discuss the effects of batch size and sequence length on the throughput degradation.\n3. From Weakness 3: Please use larger model(s) to demonstrate the consistency of PAMM's efficacy.", "limitations": "", "rating": 4, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. This paper shows that QKV projections are highly redundant along the sequence dimension and thus can be compressed by a factor of >100x for efficient training with almost no accuracy loss, marking a step toward large-size model training on limited hardware devices.\n2. The throughput degradation decreases when the model size gets larger, making PAMM more practical for large-size model training on limited hardware devices -- highly reduced memory consumption (hard constraint, must be satisfied) and low throughput degradation (soft constraint, the lower the better).", "weaknesses": "1. The memory consumption of QKV projections is significantly reduced, but that is just the memory consumption of QKV projections and accounts for 20% of peak usage -- there are many other tensors/activations consuming non-negligible memory. Therefore, the overall impact of this work may be limited (by 20%).\n2. The throughput degradation is proportional to batch_size * seq_len -- but only the effect of model size on the throughput degradation is discussed in the paper. The proposed method may be (much) less efficient with large batch size or long sequence length -- this should also be addressed.\n3. The largest model used in the experiments is LLaMA-1B -- larger model(s) should be considered.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)