Trion: FFT-based Dynamic Subspace Selection for Low-Rank Adaptive Optimization of LLMs
Authors: Ionut-Vlad Modoranu
Mher Safaryan
Erik Schultheis
Max Ryabinin
Artem Chumachenko
Dan Alistarh
Venue: ICLR 2026 Poster (2026)
Keywords: low-rank optimization
fast fourier transform
computational efficiency
memory efficiency
efficient optimization
large language models
Primary area: optimization
TL;DR: We use the Discrete Cosine Transform (DCT) via a FFT-based procedure called Makhoul's N-point algorithm to dynamically select columns from the DCT matrix to perform low-rank adaptive gradient optimization for LLMs to replace the SVD/QR methods.

Abstract: Low-rank optimization has emerged as a promising direction in training large language models (LLMs) to improve running time and reduce the memory usage of adaptive optimizers by constraining learning to a lower-dimensional space. Prior work typically projects gradients of linear layers using approaches based on Singular Value Decomposition (SVD) or QR-decomposition. Applying these techniques individually to each layer in large models is computationally expensive and incurs additional memory costs due to storing the projection matrices. In this work, we propose a computationally efficient and conceptually simple, two-step procedure to approximate SVD/QR-based gradient projections into lower-dimensional spaces by using a predefined orthogonal matrix of the Discrete Cosine Transform (DCT). We dynamically select columns from the DCT matrix based on their alignment with the gradient of each layer. The effective projection matrices are obtained via a simple \texttt{matmul} with the DCT matrix in $O(n^3)$ time, followed by a lightweight sorting step to identify the most relevant basis vectors. For large layers, DCT can be computed via \texttt{Makhoul}'s $N$-point algorithm based on Fast Fourier Transform (FFT) in $O(n^2 \log(n))$ time, yielding speed-ups for low-end GPUs. Due to the predefined nature of the orthogonal bases, they are computed once at the start of training. Our numerical experiments on both pre-training and fine-tuning tasks demonstrate the effectiveness of our dual strategy in approximating optimal low-rank projections, obtaining an approach with rank-independent running time that matches the performance of costly SVD/QR-based methods while achieving faster runtime and reduced memory usage by up to $25\%$ across different model sizes. Our code is available at \href{https://github.com/IST-DASLab/Trion}{\texttt{https://github.com/IST-DASLab/Trion}}.

TRION: FFT-BASED DYNAMIC SUBSPACE SELECTION FOR LOW-RANK ADAPTIVE OPTIMIZATION OF LLMs
Ionut- Vlad Modoranu" ISTA & Together Al
Mher Safaryan
Erik Schultheis ISTA
ISTA
Max Ryabinin Together AI
Artem Chumachenko Together
Dan Alistarh ISTA & Red Hat Al
ABSTRACT
Low-rank optimization has emerged as promising direction in training large language models (LLMs) to improve running time and reduce the memory usage of adaptive optimizers by constraining learning to lower-dimensional space. Prior work typically projects gradients of linear layers using approaches based on Singular Value Decomposition (SVD) or QR-dccomposition. Applying these techniques individually to each layer in large models is computationally expensive and incurs additional memory costs due to storing the projection matrices. In this work, we propose a computationally efficient and conceptually simple, two- step procedure to approximate SVD/QR-based gradient projections into lower- dimensional spaces by using a predefined orthogonal matrix of the Discrete Cosine Transform (DCT). We dynamically select columns from the DCT matrix based their alignment with the gradient of each layer. The effective projection matrices are obtained via a simple ma mul with the DCT matrix in ((")) time, followed by a lightweight sorting step to identify the most relevant basis vectors. For large layers, DCT can be computed via Makhoul'ss 's N N-point algorithm based on Fast Fourier Transform (FFT) in On log(n)) time, yielding speed-ups for low-end GPUs. Due to the predefined nature of the orthogonal bases, they are computed once at the start of training. Our numerical experiments on both pre-training and fine-tuning tasks demonstrate the effectiveness of our dual strategy in approximating optimal low- rank projections, obtaining approach with rank-indepeneer methods while achieving faster
running time that
matches the performance of costly SVD/QR- runtime and reduced memory usage by up 25% across different model sizes. Our code is available at https:/ /etp://gttttbbbt com/ MII/ISTT-- DASLab/Trhh
INTRODUCTION
The Adam optimizer (Kingma & Ba, 2014) and its regularized version AdamW (Loshchilov & Hutter, 2019) have become the standard approach for optimizing deep neural networks in various settings. With the recent increase in scale of LLMs up to billions and trillions of parameters, training with AdamW becomes more and more challenging, as its internal state requires two momentum buffers that scale with the model size. This practical problem paved the way for a line of research that focuses on reducing the memory usage of optimizer states in the context of adaptive gradient optimization. These approaches range from quantizing the states to 8 bits (Dettmers et al., 2021) to the recent GaLore optimizer (Zhao et al., 2024) inspired from the LoRA techniques (Hu et al., 2021: Lialin et al., 2023), that compresses the gradient matrix using low-rank decomposition based on SVD. Several improvements to GaLore have been proposed to enhance its performance, such as LDAdam (Robert al., 2025), FRUGAL (Zmushko et al., 2024), FIRA (Chen al., 2024). BAdam (Luo et al., 2024), Q-GaLore (Zhang et al., 2024). The key aspect of GaLore and its later improvements is the low-rank decomposition based on matrix factorization, such as SVD or QR decomposition.
"Corresoondnnce to Together AL.
www.wawww..o.o.deittttt...cmm ac. at. Part of work done during an internship at
I Institute of Science and Technology Austria
I Current affiliation: Lancaster University, UK

---

Recently, the Muon (Jordan al., 2024) optimizer sparked the community's attention due to the faster convergence in pretraining settings, achieved by orthogonalizing the momentum matrix using Newton-Schulz, iterative procedure difficult to parallelize for large-scale pretraining runs as the full-sized matrices must be materialized on GPU. The Dion optimizer (Ahn et al., 2025) aims to reduce this overhead by employing low-rank, orthogonal updates. However, requires storing a projection matrix for each layer and uses QR DR-decomposition to orthogonalize the low-rank components, making its running time dependent on the rank.
All these SVD/QR-based techniques are known to be computationally intensive as they have to be invoked for each linear layer, either at each step (for, e.g.. LDAdam, Muon) or once at a few steps (for, e.g.. GaLore). To mitigate the high computational and memory costs of these procedures, we ask: can we find an alternative low-rank projection approach to serve as an accurate replacement for the orthogonal matrices in SVD/QR for low-rank compression of optimizer states, that is much cheaper to compute, and portable across memory-effiiie optimizers?
Contributions. In this work, we address our research question by proposing cheaper alternative to the orthogonal matrices computed via SVD/QR. Concretely, we propose using the orthogonal matrices from the general class of Discrete Fourier Transforms, such as the Discrete Cosine Transform (DCT), which has been successfully used in image compression for the JPEG algorithm. To the best of our knowledge, are the first to use it in the context of low-rank adaptive gradient optimization. We summarize our contributions as follows:
We propose a dynamic column selection approach to adaptively choose columns from a fixed orthogonal matrix compute low-rank projection of the input matrix. The effective projection matrix obtained from the fixed orthogonal matrix by indexing the columns (Section 2.1);
We motivate that DCT is good candidate for our dynamic column selection approach due to the reduced time complexity to compute the alignments with the input matrix via the Makhoul's N -ponnn algorithm (Makhoul, 1980) to compute fast DCT with O log(n)) time complexity compared to (n²) incurred by basic matrix multiplication. This can yield speedups of up to 8 50x in computing the alignments for large layers (Appendix C) on low-end GPUs.
We theoretically justify our dynamic column selection approach to compute the most significant columns of any fixed orthogonal matrix (including the DCT-based) to obtain projection matrix tailored to each layer (Section 4);
We show the DCT-based projection a fast and accurate replacement for the orthogonal matrices computed via SVD/QR for the Muon- and Adam-lik, optimizers (such as Dion, FRUGAL, FIRA, LDAdamW) in the context of low-rank compression of optimizer states. We reduce the running time and improve the memory usage as store only one DCT matrix per GPU for the entire network, computed once at the beginning of the training. In addition, we store only (rank) integers, representing the indices of the most significant columns for each layer instead of storing one low-rank projection matrix.
We propose two standalone optimizers that use the DCT-based dynamic column selection approach:
1. Trion (Section 2.4), which improves Dion by replacing the Power-Iteration with our DCT- based dynamic column selection approach to compute low-rank representation of the momentum buffer, followed by orthogonalization via Newton-Schul iteration. We provide a DDP-compatible implementation that communicates orthogonal, low-rank momentum across devices and compute the final layer updates locally using the projection matrix obtained from DCT. To the best of our knowledge, we are the first ones to reduce the complexity of Newton-Schull using a low-rank approximation of momentum.
2. DCT-AdamW (Section 2.5), which replaces the SVD-based low-rank projections and op- tionally adds quantized error feedback for the projection error. Further, DCT-AdamW W rotates the momentum buffers such that new low-rank gradients are correctly incorporated, and thus allows changing the low-rank subspace every step.
2 METHOD
This section presents our dynamic column selection approach that works with any orthogonal matrix, briefly introduces the Diserete Cosine Transform (DCT) and the motivation of using it and finally
2

---

introduces the two algorithms propose: Trion an improvement Dion to replace Power- Iteration and DCT-AdamW as an improvement to low-rank AdamW variants to replace SVD/QR.
2.1 DYNAMIC COLUMN SELECTION
General View. Given an orthogonal matrix Q E R want to compute a projection matrix Q E Rexr by selecting r columns from Q to project the gradient G E R=x= to an r-dimensional space g GQ, E R First, we compute the similarities matrix S GQ containing the scalar products between rows of G and columns of Q. We rank the columns of S based on their f1- or 2--norm and then pick the indices of the largest r columns according to the chosen norm, which we use to index columns in Q to obtain Q.. The ith column in matrix S contains the scalar products between each row of G and 7th column of Q, as detailed in Appendix B. We view these scalar products as similarities (or alignments) of rows in G' with columns of Q and want to choose the columns of Q with largest similarity with rows of G. Using this approach, we ensure a dynamic mechanism to choose the most appropriate columns of Q for the current gradient matrix G' to minimize the projection error (see Section 4). Each layer will have its own set of r indices for the columns. The dynamic behavior comes from the n-choose-r possible sets of indices to select from.
The rule of thumb the low-rank factorization methods for optimization we are targeting in this work is to compress the smallest dimension of matrix dimensions (e.g. from - to gxxr for n 2 m). Usually, the smallest dimension of the gradient matrix is dmodel. the hidden (embedding) size of the model. As a result, the memory overhead of our dynamic column selection approach is the cost of storing only one orthogonal matrix Q E I Rooode Xdmodde (DCT in our case) per GPU for the entire model and integers for the corresponding column indices from Q for each layer to create Qr-
2.2 PROJECTION TYPES
We follow the standard projection approach introduced in GaLore, where for gradient G' E R Rxx (R rows and C columns) we choose to collapse the smallest dimension to rank r. This set of rules applies each layer by choosing the left right projection based on the layer dimensions.
For R 2 C, we apply a right projection and we choose the most appropriate r columns from the DCT matrix of shape (C, C), resulting projection matrix P RCxr and low-rank gradient E Rxr, For It < C, we apply a *left projection * with DCT matrix of shape (R, Rt,, resulting a projection matrix P E Rrx and low-rank gradient g Rexc
2.3 DISCRETE COSINE TRANSFORM
The Discrete Cosine Transform (DCT) is widely used in the signal processing and data compression literature (e.g., JPEG algorithm for image compression) and consists of n orthogonal basis vectors whose components are cosines (Strang, 1999). There are minor variations of DCT and in this work we will use DCT-II/III. We denote the DCT-III matrix of order n by Q E R=**, defined as Qi = /2/n //2/n-cos cos 1())+1)m with i,j E [n]. where the first row has to be divided by /2 in order for Q 2r
to be orthogonal, i.c. QTC Q Ln. The DCT-II matrix is the transpose of DCT-IIL. In Appendix A, we provide details on how we can efficiently materialize the DCT matrix on GPU. Once we compute Q, can it in the dynamic column selection approach to efficiently compute low-rank projections of any two-dimensional matrix.
In Appendix C, discuss why we choose DCT matrix. In short, the particular structure of DCT allows us to enable fast computation for the similarities matrix in O(n² log(n)) time using the Makhoul's N N-point algorithm (see Appendix D) instead of (mn²) time for basic matmul.
2.4 TRION: DCT-BASED IMPROVEMENT TO DION
In this section, we present how we can apply the DCT-based dynamic column selection approach to compute low-rank projection of momentum in the context of Dion (Ahn et al., 2025) to obtain faster and more accurate optimizer, which we call Trion, presented in Algorithm 1.
Dion uses Power-Iteration to compute low-rank projection and orthogonalizes it using QR- decomposition, whose rumning time depends on the rank r. Instead, we propose replacing these techniques with a rank independent approach to compute the low-rank projection based on DCT matrix, which requires computing the similarity matrix S, followed by a top-r ranking to determine the indices of columns that best align with the momentum matrix, denoted by the set it. The matrix S,
3

---

or S, = B, De compute similarities; E N" > select largest r columns by f//22-norm projection matrix containing most significant r columns > extract low-rank momentum from the similarity matrix S > the error b,Q, and replaces P,R, error from Dion b uudate momentum with error feedback represents the DCT of the momentum matrix and it can be computed using the Makhoul's algorithm or simply by only one matrix multiplication S = BeD where B₁ E R RxC is the momentum and De E RCxC is the DCT matrix.
Once we determine the indices i, of the most significant columns in De, we can extract the low-rank momentum b, from the similarity matrix St, as well the projection matrix Q: by indexing Dc. which will be used in (a) computing the projection error A and (b) projecting the orthogonal low-rank momentum back to the higher dimensional space update the model.
We would like to emphasize that we input the low-rank momentum b, E RR to Newton-Schut and not the original momentum buffer B, R RRx, which significantly reduces the computational overhead. Moreover, we can use the efficient triton kernels provided in the Muon implementation from the official Dion repository' to speed up the computations even further for large ranks r, as Newton-Schulz will operate with X r matrices.
Communication in Distributed Training. We develop the Trion optimizer on top of the published code for the Muon optimizer that leverages the ZeRO (Rajbhandari et al., 2020) approach in Distributed Data Parallel (DDP) settings. Specifically, the model update O, for a layer is computed only on one GPU, and the result is communicated to other GPUs using a 11 -gaa ther, drastically reducing the computation costs for large models (under the assumption that communication is cheaper than computation itself). Since we replicate the DCT matrix on each GPU, we would like to emphasize that we communicate only low-rank terms oe from the source GPU (the one computing the update) to other devices and perform the step O Q, locally each device. This way, do not communicate full size, orthogonal matrices O,, but only low-rank versions o,, thus reducing the overall iteration time.
This is particularly useful in the FSDP-compatible implementation, where we can materialize the full low-rank matrices o on each device using 1-lto-all (see this Muon implementation 9) to perform Newton-Schus on a low-rank input, followed by a resharding. After resharding, each GPU can compute the individual slice of O, and update its own shard. However, the FSDP implementation requires deciding how each layer should be sharded based on whether that particular layer requires a left- or projeetooettionn in order to avoid unwanted calls to communication primitives to materialize full tensors to compute the alignments S.
Algorithm Trion Optimizer
1: Define the DCT-II/III matrix Do RCC
2: for = (1.2..... T) do
G. VoL(O,) E Rxx
3:
4:
B. = M,-1 + G2 E RXC
5:
S. MAKHOUL B,) E RRCC
6:
à WWNAMCCOLUNNSELECIIONN
7:
Q Do(;,u] E Rexr
8:
D = S1l:, E R Rxr
9:
A, 1u-biQ1
10:
M. = ni; + (1 n) B - (1-m),1))
11:
OR (EWTONSCHUIZ(b)) E Rxr
> Newton-Schut applied on low-rank by b project orthogonalized low-rank momentum to original size an) -m /R/C) O₂ E RXC
12:
O, o E RxC
13:
9
14:
end for
2.5 DCT-ADAMW
We propose DCT-AdamW, standalone low-rank version of AdamW with DCT-based projection that has the option to use quantized eгror feedback (EF) (Seide et al., 2014; Karimireddy et al., 2019: Alistarh et al., 2018) and ensures momentum buffers integrate gradients from the same lower dimensional subspaces in similar way as in LDAdam (Robert et al., 2025). In contrast to LDAdamW, which has to store two consecutive projection matrices per layer, we only have to store
1118
4

---

two sets of r indices per layer. Prior work MicroAdam (Modoranu et al., 2024) quantized the error feedback down to 4-bits in the context of compressing the optimizer state using sparsity. In our setup, the lowest resolution we can quantize EF to is 8-bits without degrading the optimizer performance. For space constraints reasons, we present the pseudocode of our DCT-AdamW optimizer in Appendix G.
EXPERIMENTS
In this section we present our numerical results. Our main goal is to show that our DCT-based dynamic column selection approach least recovers the performance of the original algorithms which we integrate it in. Specifically, we evaluate Trion against Dion, where we directly compare the DCT-based projection followed by Newton-Schulz iteration with QR-based Power-Iteration. In addition, we replace the SVD- projection in FRUGAL and FIRA optimizers with our DCT approach. In the end, we compare LDAdamW with our standalone DCT-AdamW optimizer. In our pretraining experiments, we compare training/validation perplexity and for fine-tuning we compare the evaluation accuracy, as well as memory usage and running time for both scenarios. In the remaining part of the paper, we focus on presenting our pretraining (PT) results and the fine-tuning (FT) results are presented in Appendix K.
We train from scratch models from the Llama family with 350M, 800M and 1.3B parameters using Chinchilla-optimal token counts (20 tokens per parameter) from the C4 (Raffel et al., 2020) dataset and sequence length 512. All PT experiments 8x H100 Nvidia GPUs using global batch size 512 with local batch size 64 per GPU (unless otherwise specified explicitly).
in Distributed Data Parallel (DDP) settings on
PT with Trion. Our purpose is to show that our DCT-based dynamic column selection approach can successfully replace the QR-based performance for the best hyper-parametee reported by the original work. In Table 1, we present our results, which are obtained using the optimal learning rate ') 0.01 (as reported by the original Dion paper) and weight decay À= = 0.01. We report the average values across 3 seeds for train/ krin/valldation loss and perplexity, maximum allocated memory (read directly from the GPU) and the lowest running time across all runs. We experiment with ranks 128, 256 and 512, the Lunkto-ddimesion ratio is r/d E (1/16, 1/8, 1 1/4, 1 /2), equivalent 6. 6.25%, 12.5% 25% and 50% of the full rank d, where d is model's 's embedding dimension.
procedure in Dion and at least recovers the
Performance In Table we show Trion consistently achieves lower training and validation loss. which also translates to lower perplexities. At the top row of Figure I we show the training loss curve of Trion is lower than Dion across iterations. This is an indication that our DCT-based dynamic column selection algorithm, followed by orthogonalization via Newton-Schie Schulz can successfully replace the Power-Iteration procedure in Dion.
Memory Usage. In Table 1, we show Trion has a lower memory requirement across all experiments since it allocates only one DCT matrix per GPU of size d x d, from which we select r columns using the indices of the most significant columns. In contrast, Dion stores a projection matrix for each layer. This design difference translates to around 10% lower memory footprint for Trion compared to Dion.
Runtime. We would like to emphasize that the running time of Trion does not depend rank, as can be seen from the reported runtime across different ranks in Table 1: Trion achieves nearly constant runtime for each model size across all ranks, while the runtime of Dion clearly depends on the rank because of the QR-decomposition. This represents an advantage of Trion for much larger scales compared Dion. At the bottom row of Figure 1 we present the wall clock time for the two optimizers for rank 256 across all three models we tested on. Concretely, for a fixed running time budget, Trion consistently achieves lower training loss than Dion. Trion is faster than Dion by 2.5 4.5% for rank 128, 4.5 - 9% for rank 256 and 8 18% for rank 512 (the overhead of Dion increases with rank). is important to mention that the embedding size of the models we used in work is too small to see significant difference in the runtime between Makhoul's algorithm and matmul when computing the column similarities. However, our benchmark presented in Appendix D shows the advantage of using Makhoul's 's aachooullssllloooian algorithm at scale compared to matmul.
Additional Runtime Benchmarks. In Appendix F we provide a more comprehensive ablation for the optimizer step by comparing full-rank Muon with Dion and Trion to show the behavior of our Dynamic Column Selection algorithm when computed via matmul and via Makhoul's 's algorithm. It is important to mention that the benefits of Makhoul's 's algorithm can be seen older-generation
5

---

Rank r
Metric
350M
1024)
800M (d 2048)
1.3B'
2048)
Trion
Dion
Trion
Dion
Trion
Dion
r/d
1/8
1/16
1/16
Train Loss
2.726
2.764
2.532
2.555
2.475
2.503
Train PPL
15.29
15.89
12.59
12.89
11.90
12.24
128
Val Loss
2.736
2.773
2.505
2.527
2.422
2.448
Val PPL
15.43
16.00
12.25
12.52
11.27
11.56
Memory (GB)
42.42
45.56
67.45
71.64
63.62
68.57
Runtime
1h 53m
58m
8h
8h 24m
19h 13m
42m
r/d
1/4
1/8
1/8
Train Loss
2.715
2.741
2.533
2.546
2.476
2.492
Train PPL.
15.11
15.51
12.61
12.78
11.92
12.11
256
Val Loss
2.727
2.750
2.503
2.519
2.423
2,440
Val PPL
15.30
15.64
12.22
12.42
11.28
11.47
Memory (GB)
42.42
45.59
67.45
71.75
63.62
68.58
Runtime
1h 53m
2h 3m
8h 3m
8h 36m
19h 13m
20h 6m
r/d
1/2
1/4
1/4
Train Loss
2.708
2.718
2.528
2.535
2.470
2,481
Train PPL
15.01
15.16
12.54
12.63
11.84
11.97
512
Val Loss
2.718
2.727
2.499
2.509
2.420
2,430
PPL
15.15
15.29
12.18
12.30
11.25
11.36
Memory (GB)
42.42
45.99
67.45
71.81
63.62
68,73
Runtime
1h 54m
2h 14m
8h 3m
8h 48m
19h 13m
20h 44m
Train Loss
2.697
2.517
2,513
Train PPL
14.85
12.41
12.35
Val Loss
2.707
2.489
2.409
Muon
PPL
14.99
12.05
11.13
Memory (GB)
42.42 GB
67.45 GB
63. 64 GB
Runtime
1h 52m
8h 25m
19h 17m
Table Perplexity, Memory, and Running Time Comparison for Trion and Dion. Muon added for reference. d stands for the embedding dimensionality of the model. The 1.3B model was trained with local batch size 32.
GPUs, which have fewer/less powerful tensor-cores. However, on newer GPUs, the tensor-cores are much faster than the gains of FFT-based procedures. We benchmarked matmul and Makhoul's 's algorithm on Nvidia H100 and RTX 2080/3090 and provide details in Appendix E.
4.1
Trics
1h:m
Diss
314
2-4
23.2
1550
2000
1800
1007 vile
2330
2.2
2.8
11000
20000
4000
33000
53007
30703
10750
in
kx
4.5
4.5
4.5
Thmm
Tilm
3.8
7im
2.6
3.4
3.4
3.4
3-2
2.0
2.5
3:35
0:35
0.3
106
1.25
1.40
1.50
2:50
4.10
133
0.35
1.3
2.25
333
8.30
1136
B3.5
1800
1922
(a) 3500
(b) Llama-800M
(c) Llama- 1.3B
Figure Training Loss across iterations (top row) and wall clock time (bottom row) for three
Llama models trained with Dion and Trion optimizers C4 using 20 tokens per parameter and rank r = 256. For a clearer visualization, the training loss curve was smoothed using a moving average with window size of 200 and zoom-in the last 5000 training steps.

---

Factorization Accuracy. In order to understand why our low-rank projection in Trion is more accurate than Power-Iteration in Dion, we compute the & 22-norm of projection errors between the accumulators B, and the corresponding orthogonal updates that each optimizer uses to update the model: A dion = || Br P.Q. ||2 for Dion and Atrion = I| - Osll2 for Trion. In Figure 2 we show the projection errors for the linear layers of the first transformer block in Llama-30M model and provide more information in Appendix H.
PT with LDAdamW. W. We train Llama-800m 800M on 80B tokens (100 tokens per parameter) from C4 using LDAdam and DCT-Adam W described in Algorithm 2. In particular for our approach, we use EF quantized 8-bits and some further optimizations from the EedOdredunyyncy optimizer (Rajbhandari et al., 2020). where one layer is updated on a single GPU and then it is broadcasted to the other GPUs. This way, we obtain lower memory usage by replacing redundant operations in the optimizer with communication, since the GPUs receiving the updated layer parameters do not allocate any state. In Figure we present the training loss curves for full-rank Adam (for reference), L.Adam W and DCT-AdamW and we are interested in directly comparing the last two. We observe that DCT-AdamW has lower training loss than LDAdam W. which also translates to lower perplexities Table 2. Due to relatively high rank, the memory usage of LDAdamW is close to the memory usage of Adam because it stores two projection matrices to be able to rotate the momentum buffers. In contrast, DCT-AdamW W stores only two sets of indices instead of storing the actual projection matrices, which drastically reduces the memory usage, coupled with the ZeRO-redundancy trick. In terms of running time, DCT-AdamW is faster than LDAdamW by 10h 7m (= 25.75%) and slower than AdamW by Ih 55m ( 5%).
Other PT Ablations. We replace the SVD projection with our DCT approach and test it on Llama- 800M model using 16B tokens from the C4 against FRUGAL/FIRA. For space constraints reasons, we present the results in Appendix In addition, we provide some pre-training results for DCT-Adam and Sub Track++ Rajabi et al. (2025) in Appendix J.
11.25
monofrmer HAMEL zmoweige
Dian
11.25
Tarchmerhh rrb (et weight
rondheeeeh wright
Tikes
16.75
2..30
10.50
10.25
Dizi
10:00
1:.00
Donn
Touz
500
1000
1.50
2000
300
1600
1550
2000
30.20
300
1000
1000
2000
Treasing Step
Triising Sxg
Tribsing Skip
Figure 2: Projection errors for Dion and Trion for a few layers from the first transformer block on Llama-30M (d 640, = 128).
4.0
AdamW
LDAdamW DCT-AdamW W
3.8
LDAdlon
Adam W
Train PPL
12.88
15.10
14.95
g
3.6
DCT-AdumW
3.4
Val. PPL.
11.73
13.91
13.69
3.2
Mem. (GiB)
73.72
72.10
57.82
3.0
Time
1d 13h 22m
2d 1h 24m
1d 15h 17m
2.8
Wwwwena www. wwwpk//wwwwwwoe///////oooooo///.
2.6
Table Pre-training results on Alama-800M. for AdamW, LDAdamW and DCT-Adam with 100 tokens/parameter. AdamW is the full-rank optimizer and is added for reference.
2.4
0
50000 100000
150000
200000 250000 360000
iter
Figure 3: Pre-training Llama-800M with DCT-
AdamW and LDAdam on 80B tokens from C4.
4 THEORETICAL GUARANTEES
In this section, we provide a theoretical justification for our two-step procedure to approximate SVD-based gradient projections. First, we rigorously show that adaptively selecting columns of an orthogonal matrix based on their alignment with the gradient matrix is an optimal strategy for minimizing the reconstruction error. This approach leads to a contractive compression scheme, which is commonly exploited in the analysis of compressed adaptive optimization algorithms. Next,
7

---

to justify the specific use of the DCT matrix, we demonstrate that it naturally serves as linear approximation of the left or right eigenbases of the gradient matrices.
4.1 OPTIMALITY OF NORM-BASED RANKING PROCEDURE
Let G' E R= be the gradient matrix and Q E RR** is orthogonal, namely QQT = I, (without loss of generality, we consider left multiplication). For given rank S n, let Qr be sub-matrix of Q composed of r columns. We are interested in the reconstruction error between G and Q G, where Q, performs projection and Qr performs back-projection. While we may choose any matrix norm to measure the reconstruction error, the standard Frobenius norm makes the derivation cleaner. Specifically, using the definition of the Frobenius norm |G]| tr(G' G), linearity of trace and Q Qr
I, due orthogonality of Q, decompose the reconstruction error: || - Q,Q, G = tr(G" - (1-Q,),)) (,,) = tr(G (1 - Q,Q, GG)
= tr(G G' = tr G) tr(G' Q2Q,
|G]] 1Q, |
1 E G
where 's are the columns of Q.. From this identity, we conclude that to minimize the reconstruc- tion error, the optimal strategy is to maximize the alignments 11 G1 for all selected columns g. Furthermore, since 19_ G + و1 GI1B +...+ l1q GIB =IQ G tr NCTQQTT) G)
tr(G G) = |[ following the optimal strategy of selecting columns from Q, we have ||G Q, GIP || - Gr
||G| - + | (1 +) [G]] Thus, the proposed norm-based selection strategy for any orthogonal matrix Q induces low-rank compres- sion scheme that is contractive with factor 1 - 7//.. Contractivness of the compression scheme is the key property in the convergence analysis of compressed optimization (Stich et al., 2018; Richtárik al., 2021; Li et al., 2022; Modoranu al., 2024; Robert et al., 2025).
We can extend the above argument for any p-norm over vectorized matrices, for which the Frobenius norm is the special case of p = Then, we have:
||G Q.22 Gllp = | 9:9. G Ei= 99 GlI, | T.p.T..... GII, S Ei ll Gllp
(b) Ei 11: [ 2 max n Ei 11077 Glll:
where (a) follows from the triangle inequality of the p-norm, (b) follows from the definition of p-norms for matrices, namely lww 1|p = (Eng u/2r)// (E, ut E, uz)' ||w/|p Iellp for two column-vectors and of the same size, (c) follows from the relationship between Ep norms and that ||aal12 = i.,.e ||u||p S n² llll2 if p I 2 and ||и||p S ||l|a if 2 2.
4.2 DCT AS LINEAR APPROXIMATION OF THE GRADIENT RIGENBASTS
Given a real-valued gradient matrix G E T and its SVD decomposition G = UEVT our goal find fast approximation of (without loss of generality) its left eigenvectors stacked in U R As GG = U it is equivalent to approximate the cigenvectors of symmetric matrix GGT.
Our argument starts from a linear algebra decomposition result, originally motivated by the optical information processing literature to factorize linear transformations that can be implemented optically (Müller- Quade et al., 1998; Schmid al., 2000). The result states that any square matrix M of shape X over the complex numbers C can be decomposed into product of diagonal and circulant matrices, i.e., M D2D3 ... D2x-3C2k-pp D2x-1, where D's are diagonal matrices and C's are circulant matrices. Circulant matrices are a special class of Toeplitz matrices in which each row is cyclic right shift of the previous one as shown below:
Co
C1
C2
Cn-1
C=
Cp
CD
C1
Cn-2
F
1 vn
[w 1i=00
where w = e 2rd
Ci
C2
Cn-1
co
The number of factors 2k in this decomposition has been shown to be up to 2n for almost all matrices (in the sense of Lebesgue measure), and it is further conjectured that decomposition
8

---

with up to n factors is sufficient (Huhtanen & Perämäki, 2015). Since circulant matrices can be diagonalized using the discrete Fourier transform (DFT) matrix F (see equation 1), i... C' = F> D (Golub & Van Loan, 2013), we can decompose the matrix F MF with M GG and arrive at the following decomposition for GG
GG = (FDDF"")
2 (F F*).
Notice the analogy between this matrix decomposition and the Taylor expansion for functions. Since the space of matrices is finite-dimensional, the signal-matrix can be recovered using finitely many "variable" D's, in contrast to the infinite series required for function expansions. Keeping this analogy in mind, just as loss functions are linearly approximated at each iteration in first-order optimization algorithms, consider "linear" approximation of the decomposition equation 2 by including only a single variable D, i.e., GG 2 F.,
This directly implies that approximate the eigenvectors U by the DFT matrix F. However, since the matrix GG is real and symmetric by design, its eigenvalues are real, and the real part Re(F) also forms an approximate eigenbasis that better aligns with U. Finally, we observe that the real part Re( F) corresponds to the discrete cosine transform (DCT), up to minor variations.
5 RELATED WORK
Our approach aims to improve existing optimizers from prior work in the literature, which we group into two categories: optimizers that speed up convergence at the expense of increased FLOPs and optimizers focused on reducing the running time and memory usage by compressing the gradient via low-rank matrix factorization.
Fast Convergence Optimizers. Muon optimizer (Jordan al., 2024) has recently stood up in the literature for its fast convergence rate due to the orthogonalized momentum update. The purpose of the orthogonalization is push the singular values of the momentum matrix towards 1. It speeds up convergence by increasing the importance (singular values) of directions in the momentum matrix, which otherwise would have low impact over the optimization.
The most straightforward approach to computing an orthogonal update is to use the UVT from the SVD decomposition of the momentum matrix. Since SVD is expensive, it is replaced by an iterative procedure called Newton-Schulz which involves computing the odd powers of the momentum matrix up to 5th order and multiplying by some carefully chosen constants. While the Newton-Schule procedure delivers an accurate approximation of UVT, the odd powers in the polynomial involve full- size matrix multiplications. This is in particular difficult for large scale settings, where the full matrices have to be materialized on GPU before running Newton-Schule which increases communication and memory usage. Dion optimizer (Ahn et al., 2025) aims to reduce the communication overhead by using low-rank, orthogonal updates computed via Power-Iteration, that requires QR decomposition with running time depending on the rank.
Instead, propose to reduce the overhead of Nowton-Schula in Muon and QR- deoocoomnonition in Dion by factorizing the momentum to a low-rank matrix using our DCT-based dynamic column selection approach and orthogonalizing the low-rank momentum using Newton-Schulz, then project back to the original space to update the model parameters.
Low-Rank Compression of Optimizer States. Most approaches use individual projection matrices tailored to the gradient at each layer by invoking techniques based on quantization or matrix factor- ization, such as SVD, QR or PCA reduce the memory usage of the optimizer states stored in GPU memory. The most common approach is to factorize the gradient to low-rank matrices. One pi- oneering work in the context of low-rank adaptive optimization is the recent GaLore optimizer (Zhao et al., 2024), which performs SVD once at a few steps to project the gradient to a lower- divee--imeniional space. GaLore was followed by several other optimizers that improve certain aspects of it. LDAdam (Robert et al., 2025) aims to improve the computational runtime of Gal-ore by compressing the first order momentum in Adam W, replacing SVD with a Block Power-lteration (Bentbib Kanber, 2015) and performing smooth subspace transition by rotating the first and second momentum accordingly to incorporate gradients from the same subspace at each step.
By default, GaLore discards the projection error, which LDAdam stores and incorporates into the gradient at the next step to improve convergence. In contrast, FRUGAL (Zmushko et al., 2024)
" is the conjugate transpose of the complex matrix F.

---

makes the distinction between the low-rank gradient (called state-full) and the projection error (called state-free free). The state-full gradient used in AdamW, while the state-free gradient is fed to an optimizer without state, such SignSGD. The main idea is to leverage the remaining gradient information the projection error since is available at each step instead of discarding or storing it. concurrent work to FRUGAL is FIRA (Chen et al., 2024), which preserves the low-rank constraint for memory efficiency while achieving full-rank performance by properly scaling the projection error.
Another approach worth mentioning is Online Subspace Descent Liang al. (2024), which replaces SVD projection with Online PCA and involves computing the projection matrix as a solution of an optimization objective focused on: (a) minimizing the projection error and (b) foreing the projection matrix to be orthogonal. The authors indicate that performing one step with Adam to solve this additional optimization problem is enough to obtain a qualitatively good projection matrix P. In contrast, our method does not introduce such overheads during training since the DCT matrix is already computed the beginning of training.
Despite all aforementioned approaches being similar in using SVDIQR- -baed projections, they differ few aspects. The most important one our view is how they handle the projection error and how often they update the low-dimensional subspace, which we clarify in Table 3.
Our approach uses DCT projection coupled with a dynamic column selection determine a projection matrix tailored to the gradient/momentum for particular layer and can be integrated into any optimizer, regardless of the way it handles the projection error.
CONCLUSION AND LIMITATIONS
We introduced the Trion and DCT-Adam optimizers that use our DCT-based dynamic column selection to replace techniques used perform low-rank decomposition, that is, the inaccurate Nower-Iteration in Dion and the expensive SVD and QRdeccomposition in BRUGAL/FIRA/LDAACE We showed that our work improves running time and memory usage. Moreover, recovers the accuracy of the original methods and thus serves as a cheaper alternative to these expensive and inaccurate methods used to perform low-rank decomposition in adaptive gradient methods for both pretraining and finetuning.
Our experiments are limited to models with at most 1.3B parameters for pretraining and additional work is required to test our technique for larger models and beyond the Chinchilla-optimal token counts, which would require significantly more computational resources.
Low-rank Projection
Type
Frequency
Error
GaLore (Zhao al., 2024)
SVD
200
discard
FRUGAL (Zmushko et al., 2024)SVD, Random, RandPerm
200
feed to SignSGD
FIRA (Chen et al., 2024)
SVD
200
norm-based scaling
LDAdam (Robert et al., 2025)
Block Power-Iteration
error feedback
Dion (Ahn et al., 2025)
Power-Iteration
save momentum
Trion (this work)
DCT
same as Dion error feedback
DCT-AdamW (this work)
DCT
any
Table 3: Properties of prior low-rank adaptive optimizers. The update frequency* 200 is the default
in GaLore that made the approach computationally feasible.
7 ACKNOWLEDGEMENTS
We would like to thank the Scientific Computing Department ISTA for providing access to computational resources to develop this work.
REFERENCES
Kwangjun Ahn, Byron Xu, Natalie Abreu, and John Langford. Dion: Distributed orthonormalized updates, 2025. URL hwwph:tt/t//aaae org.
wrt//bs//////.oooo/o///...ssssss abs/ 2504. 05295.
10

---

Dan Alistarh, Torsten Hoefler, Mikael Johansson, Sarit Khirirat, Nikola Konstantinov, and Cédric Renggli. The convergence of sparsified gradient methods, 2018. URL htLpstt:://aaa /ttpav.////.x.oo//////
Abdeslem Bentbib and A. Kanber. Block Power Method for SVD Decomposition. Analele Stúntifice ale Universitatii Ovidius Constanta, Seria Matematica, 23:45-58, 06 2015. doi:
10.1515/auom-2015.5
Xi Chen, Kaituo Feng, Changsheng Li, Xunhao Lai, Xiangyu Yue, Ye Yuan, and Guoren Wang. Fira: Can we achieve full-rank training of Iims under low-rank constraint?, 2024. http::
//arxiv... wttgg//b/////eo.eee 2410. 623.
Tim Dettmers, Mike Lewis, Sam Shleifer, and Luke Zettlemoyer. 8-bit optimizers via block-wise quantization. arXiv preprint arXiv:2 10.02861, 2021.
Gene H. Golub and Charles F. Van Loan. Matrix Computations 4th Edition. Johns Hopkins University Press, Philadelphia, PA, 2013. doi: 10.1137/1.9781424407994 URL https:// epubs. siam. wttp/oor/////////e 110 1137/1.9781421111111
Edward J. Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen -hu, Yuanzhi Li, Shean Wang, Lu Wang. and Weizhu Chen. Lora: Low-rank adaptation of large language models, 2021. URL https:
///axxivoorg//ae absg///b//// /albss//10. 2206. 09685.
Marko Huhtanen and Allan Perämäki. Factoring matrices into the product of circulant and diagonal matrices. Journal of Fourier Analysis and Applications, 2015.
Keller Jordan, Yuchen Jin, Vlado Boza, Jiacheng You, Franz Cesista, Laker Newhouse, and Jeremy Bernstein. Muon: An optimizer for hidden layers in neural networks, 2024. URL https: ///e/llerreeee Jeroodnnn github. iae.ii///pa wttp/osss//moott
Sai Praneeth Karimireddy, Quentin Rebjock, Sebastian U. Stich, and Martin Jaggi. Error feedback fixes signsgd and other gradient compression schemes, 2019. URL httt:ps://arxvv.oo/
abs/ / 1 901.... 09847.
Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization, 2014.
Xiaoyun Li, Belhal Karimi, and Ping Li. On distributed adaptive optimization with gradient compres- sion. arXiv preprint arXiv:2205.. 00632,, 2022.
Vladislav Lialin, Namrata Shivagunde, Sherin Muckatira, and Anna Rumshisky. Relora: High-rank training through low-rank updates, 2023. URL htwps: /wwxivivvvvv. org/ hass//b///// 2207 05695.
Kaizhao Liang, Bo Liu, Lizhang Chen, and qiang liu. Memory-efficient LLM training with online subspace descent. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024. URL https: http://peeeeeeeeeeeeeooooraaaaaa
TTTCT 6g45.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In Proceedings of the Seventh International Conference on Learning Representations, 2019.
Qijun Luo, Hengxu Yu, and Xiao Li. BAdam: A memory efficient full parameter training method for large language models. arXiv preprint arXiv:2404.02827, 2024.
J. Makhoul, fast cosine transform in one and two dimensions. IEEE Transactions Acousties, Speech, and Signal Processing, 28(1):27-34, 1980. doi: 198.099AA.P.......1
Ionut-V Vlad Modoranu, Mher Safaryan, Grigory Malinovsky, Eldar Kurtic, Thomas Robert, Peter Richtarik, and Dan Alistarh. Microadam: Accurate adaptive optimization with low space overhead and provable convergence, 2024. URL https: http://awwww...gaa////////aaa
15593.
J. Müller-Quade, H. Aagedal, Th. Beth, and M. Schmid. Algorithmic design of diffractive optical systems for information processing. Physica D: Nonlinear Phenomena, 120(1):196-24) 1998. ISSN 0167-2789. doi: science/ http/iatiiiiiieeeee sciencedirech of the Fourth Workshop Physies and Consumption.
URL https: /wwww Proceedings

---

Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. Journal of machine learning research, 21((40):1-67, 2020.
2020.
Sahar Rajabi, Nayeema Nonta, and Sirisha Rambhatla. Subtrack++ Gradient subspace tracking for 01586.
scalable Ilm training, 2025. URL https:tt//// www.i.ooioo............ooo.....a
Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, and Yuxiong He. Zero: Memory optimiza- tions toward training trillion parameter models. In SC20: International Conference for High Performance Computing, Networking, Storage and Analysis, pp. 1-16. IEEE, 2020.
Peter Richtárik, Igor Sokolov, and Ilyas Fatkhullin. EF21: A New, Simpler, Theoretically Better, and Practically Faster Error Feedback. arXiv preprint arXiv:2106.05200, 2021.
Thomas Robert, Mher Safaryan, Ionut- Vlad Modoranu, and Dan Alistarh. Ldadam: Adaptive optimization from low-dimensional gradient statistics, 2025. URL https: /tttps:/aaxii..oo wwwiv.orggo..................... org/ abs/2410. 16133.
Michael Schmid, Rainer Steinwandt, Jörn Müller-Quade, Martin Rötteler, and Thomas Beth. Decomposing a matrix into circulant and diagonal factors. Linear Algebra
and its Applications. 306(1):131-143, 2000. ISSN 0024-3795. doi: http://doi.org/10
wtt.eeeeeeeeeeeeeeeeeeeeeeeeeeee com/sciencee/ 1016595022-399559900555555 URL https: http:://ww..ciicndir..ee... wtw.:/eeettcrccccooeeeeeeeeeeeee arricile/p11/50224775550000
Frank Seide, Hao Fu, Jasha Droppo, Gang Li, and Dong Yu. 1-bit stochastic gradient descent and its application to data- aata-parallel distributed training of speech DNNs. In Fifteenth Annual Conference of the International Speech Communication Association, 2014.
Sebastian U. Stich, Jean- Baptiste Cordonnier, and Martin Jaggi. Sparsified SGD with memory. arXiv preprint Xiv:
2018.
Gilbert Strang. The discrete cosine transform. SIAM review 41(1):135-147 1999.
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. arXiv preprint arXiv:2407.08296. 2024.
Jawei Zhao, Zhenyu Zhang, Beidi Chen, Zhangyang Wang, Anima Anandkumar, and Yuandong Tian. Galore: -emory-efficient Iim training by gradient low-rank projection. arXiv preprint arXiv:2403.03507, 2024.
Philip Zmushko, Alcksandr Beznosikov, Martin Takáč, and Samuel Horváth. Frugal: Memory- efficient optimization by reducing state overhead for scalable training, 2024. URL https:
//arxiv.ort www.voorrroooooo abs http//1s.2221...1 /ttp//211.00000 078377....
12

---

CONTENTS
Introduction
Method
2.1
Dynamic Column Selection
2.2
Projection Types
2.3
Discrete Cosine Transform
2.4 Trion: DCT-based improvement to Dion
2.5
DCT-AdamW
Experiments
4
Thcoretical Guarantees
4.1 Optimality of norm-based ranking procedure
4.2 DCT linear approximation of the gradient eigenbasis
5
Related Work
Conclusion and Limitations
10
7
Acknowledgements
10
A Efficient Computation of the DCT Matrix on GPU
14
Dynamic Column Selection Approach
14
Motivation of Discrete Cosine Transform Matrix
14
D Makhoul's Algorithm Explained
15
Benchmarking Makhoul's Algorithm on Different Harwdare
16
Benchmarking running time of optimizer step
16
G Pseudocode of DCT-AdamW Optimizer
18
H Projection Errors of Trion and Dion
19
I Pre-training with FRUGAL/FIRA
20
J Pre-training with DCT-Adam W and SubTrack+
20
K Fine- Tine-Tuning
21
13

---

A EFFICIENT COMPUTATION OF THE DCT MATRIX ON GPU
In this section we present how can efficiently compute the DCT-III matrix GPU using vectorized implementation that fast for large values of on the GPU, To create the DCT-IIII matrix 1]", which is used create the matrix
Rux=, first create vector L [0.-
E N by replicating on the columns of Z:
0
0
0
2 n
I (2I 1) π
L
I
Q
cOs
2n
nxn.
We restate that we need to divide the first row by v2 make Q orthogonal. The computational efficiency comes from the element-wise product I @ (2I + +) that actually computes the integer entries ((2 + 1). The DCT-II matrix can be obtained by transposing the DCT-III.
B DYNAMIC COLUMN SELECTION APPROACH
In this section provide further details about dynamic column selection approach. When compute the similarities S, we look at the columns of where on 7th column we have the scalar products between all rows of G and the hth column of Q. By choosing the columns with largest l or fo-norms we make sure we select the columns with largest overall alignment with rows in G.
G Q1
G
Q
G
G
Qn
S GQ
Q1
Q2
Q
=
(3)
Gn
G G
GI
C MOTIVATION OF DISCRETE COSINE TRANSFORM MATRIX
In general, any orthogonal matrix would yield similar quantitative results when used with our dynamic column selection technique. In this section are interested in exploring different types of orthogonal matrices and our goal is to find candidate matrix whose structure allows us to reduce the computational complexity for the operation S GQ, especially for large layers.
The first candidate is a random orthogonal matrix Q. which can be obtained by generating a random Gaussian matrix and then orthogonalizing it using Qecdecomnoiition (e.g. keeping only the Q- component from the decomposition), which should be done only once at the beginning of training. While this would work, has the drawback of always having Onn')) complexity when computing the similarities because the matrix does not have any structure that would allow us to use an algorithm with lower complexity.
The second candidate is the Hadamard matrix, a particular Fourier matrix containing only values +1, which is used in the model compression literature and is known for its fast multiplication routines tailored to GPUs. However, it has the drawback of being ill defined for certain values of model's s dimension dmodel and the current existing procedures provide matrix that not orthogonal in these cases, making it unusable for our context.
Our preferred candidate the DCT matrix, which also a Fourier matrix and has a particular structure that gives it the potential of computing S GQ (S is called the discrete cosine transform of G), making faster than the ((n²) complexity of matmul, as we present below.
When Q is the DCT-II matrix, we can reduce the complexity of S GQ from O(n²) to (og) by using the Makhoul's 's N N-point algorithm Makhoul (1980) to compute DCT faster. We summarize the algorithm in Appendix D and also provide a benchmark for different layer sizes and data types encountered in practice. In short, this is a FFT-based algorithm that benefits from the pre- and proc-proes steps of the input matrix G (such permutations and multiplications with complex values) to reduce the computational complexity. Interestingly, the link between Makhoul's algorithm compute the cosine transform and the same transformation computed via a basic matmul is that the
14

---

matmul version embeds all operations of the Makhoul's algorithm in the DCT-II matrix itself, but at higher computational cost.
Given the current data types implemented in PyTorch, Makhoul's algorithm be run only float32 inputs, which makes it infeasible for practical bíloatl6 inputs unless the imput size a power of 2. The limitation comes from the lack of complex- complex-bfloat16 6 type support in Py Torch to represent the real and imaginary parts complex number as bfloat16. See Appendix D for more details.
D MAKHOUL'S S ALGORIIHM EXPLAINED
Makhoul's -ppoint algorithm to compute a fast DCT can summarized as follows:
apply permutation the input signal X obtain XPERM: vector [a, b, c, e, /] becomes la, f. dd b), where odd indices are in increasing order, even indices are in decreasing order and they are interleaved (can be cached for the same input size);
2. compute FFT of the permuted signal XPERM and obtain XFFT;
3. compute Fourier coefficients W₁- exp( (ikk-ikkk/)) (can be cached for the same input size);
4. obtain DCT of X as X Real(XpFr" W) (multiplication by W accounts for permuting the input X).
Makhoul's algorithm performs one-dimensional, type-1l DCT for each row of matrix G. In essence, the Makhoul's algorithm S = MAKHOUI. is equívalent to the matmul version S = GQ, where the matrix Q embeds all the operations performed by the Makhoul's 's procedure.
We benchmark Makhoul's s algorithm against the matmul version on GPU for different layer sizes and show it is faster than matmul for large layers. For small layers (up embedding size 2048, the benefits of Makhoul's algorithm do not have practical benefits compared to matmul they have similar runtimes). Concretely, we run 10 iterations of warmup before starting to measure the time for 100 runs on the GPU for both matmul version and Makhoul's 's version, where G is initialized with random gaussian data float32 and stays fixed during the benchmark.
In Table we show our benchmarking results where gradient matrix G is in float32 format, compare against the DCT transform obtained via matmul by simply employing S = =Q and compare against the Makhoul's 's N-point algorithm. The column Ratio shows how much faster the Makhoul's implementation is compared to the standard matmul. can get up to 50x speedup for matrices with more columns than rows.
Input size
Source Model
Matmul time (s)
Makhoul time (s)
Ratio (@ /FFT)
(4096, 4096)
Llama-2-7B Qwen3-32B
0.00267808 0.02746414
0.00033124
8.09
(25600, 5120) (5120, 25600)
0.00259182
10.60
Qwen3-32B
0.13069152
0.00262571
49.77
Table 4: Comparison of Matmul vs. Makhoul runtimes across different input sizes and models for float32. Ratio measures how much faster the Makhoul's 's algorithm is compared to the standard matmul. Matmul is faster for ratio < and Makhoul is faster when ratio 1.
Input size
Source Model
Matmul time (s)
Makhoul time (s)
Ratio (@ / /FF)
(4096, 4096)
Llama-2-7B
0.00018977
0.00034137
0.56x
(25600, 5120)
Qwen3-32B
0.0018455 84539
0.00268 176
0.69x
(5120, 25600)
Qwen3-32B
0.00968907
0.00273731
3.54x
Table 5: Comparison of Matmul (bfloat16) Makhoul (float32) timings across different input sizes and models with different data types.
While get significant speedup for the case R < C for float32, the real practical settings use 6 instead of float32 and for this reason we also run our benchmark for bfloat16 as follows: we generate random matrix G in float32 which will be used for Makhoul's 's algorithm and the same
15

---

matrix G will be converted to 66loat16. Moreover, the DCT matrix Q will also be stored in 6 such that both operands in S = GQ are in bfloat16. On the other hand, since Makhoul's 's algorithm uses FFT from PyTorch we are currently restricted to using float32 imputs because at the moment of developing this work PyTorch does not have the complex-bfloat16 type, where both real and imaginary parts are stored in 6 (half precision is supported only for inputs with sizes power of two). As a result, we run Makhoul's 's algorithm only in float32 and compare with matmul in bfloat1
In Table 5 we show our benchmarking results where we compare matmul run with 6 inputs against Makhoul's s algorithm run on float32 input. For R 2 C. matmul-bfloat/ is consistently faster than Makhoul's 's algorithm. This is expected because 6 type has higher throughput on GPUs than float32. However, we see that Makhoul's 's algorithm is 3.5x faster than matmul-bfloat for R < C, which is a much lower speedup compared to the results Table for float32.
The FFT- algorithm shows reduction in the running time in our benchmark, the theory predicts (e.g. O(n9log(n)) compared to O(n²)) for float32, while the gains are limited when we compare against matmul-bffoatl6. In order to benefit from the faster computation of Makhoul's algorithm, we would need to convert 6 matrices to float32, run the faster procedure, then convert back to bfloat1 which is not feasible because of the additional memory and computational overhead. However, when running training in mixed precision, the gradients are computed in 66loat and accumulated in float32 buffer for precision reasons. Fortunately, we have access to this float32 buffer in the optimizer's 's step function in Py PyTorhh
E BENCHMARKING MAKHOUL'S ALGORITTM ON DIFFERENT HARWDARE
The primary speed up comes from our Dynamic Column Selection algorithm that replaces SVD/QR. not from the FFT-based algorithm, all results were obtained using matmul to compute the similari- ties S = GQ instead of Makhoul's algorithm.
Given two matrices A, B with shapes (m, n) and (n, the complexity of matrix multiplication A.B B is ((mp)). Modern GPUs have powerful tensor-cores specialized in matrix- (ali---mullllucuummmme (MMA) operations such as D A x B +C, which reduce the practical complexity by parallel processing. Of course, the number of multiplications is still of order O (mmp), but they are performed in parallel and therefore the effective running time is lower. The H100 and B200 GPUs from Nvidia have powerful tensor-cores that erase the theoretical benefit of Makhoul's 's algorithm.
Our benchmark in Appendix D shows the improvements of Makhoul's algorithm on an Nvidia H100 GPU with tensor-cores support for matmul turned off by setting a low_tf32=Fall se. When this flag is set to True, the
torch. brcchbaccendds matmul is faster than Makhoul's s algorithm, as we can see below for layer of shape (4096, 4096). However, on older GPUs such RTX 2080, we can see the benefit of Makhoul's 's algorithm as setting the flag allow awwawwtttl=aaa___aaaaaaaa wtwe=frreeeeeeeee Faasse does not have significant impact over running time for computing the similarities S via matmul or Makhoul.
Our Dynamic Column Selection algorithm can be used with any fixed orthogonal matrix because it minimizes the projection error regardless of the chosen matrix. In particular, the DCT matrix has the benefit of theoretical speedup, which can also be achieved in practice on GPUs with less powerful tensor-cores.
In Table 6 below we benchmark the operation of computing the similarities S = GQ computed via matmul and/or Makhoul's algorithm for some practical matrix sizes that can be found in a standard Llama block with embedding size 4096. We can see that Makhoul is always faster than matmul when the tensor-cores are turned off.
For RTX-2080/3090, using Makhoul's algorithm is indeed more beneficial.
Important mention. We would like to emphasize one observation in the following table. Looking at specific layer such as (4096, 4096), user owning Nvidia RTX 2080 GPU can compute the similarities S in 1.3 milliseconds using Makhoul (rows 11-12), which is equivalent to approximately 4 times slower than matmul on H100 with tensor-coree ON (row 7). This is a great achievement since using Makhoul on RTXX-080 (row 12) would be 8 times faster than applying matmul (row 11).
F BENCHMARKING RUNNING TIME OF OPTIMIZER STEP
We created a custom model with only one transformer block where the embedding dimensionality is 4096 and number of heads is 32, resulting in model with 464M parameters in total (embeddings,
16

---

Table 6: Matmul (@) vs. Makhoul (FFT) runtime comparison across GPUs and matrix sizes
Matrix Size
GPU
Vensor-Cores
Row #
Matmul (ms)
FFT (ms)
Ratio
Faster
1
1.31 ms
1.07 ms
1.21
FFT
H100
2
8.07 ms
1.06 ms
7.58
FFT
(12288, 4096)
V
3
11.07
2.62 ms
4.22
FFT
RTX 3090
4
16.46 ms
2.62 ms
6.28
FFT
/
5
30.51 ms
3.88 ms
7.85
FFT
RTX 2080
x
6
30.79 ms
3.89 ms
7.91 x
FFT
V
7
0.36 ms
0.39 ms
0.92 x
@
H100
X
8
2.67 ms
0.37 ms
7.08
FFT
(4096, 4096)
/
9
4.00 ms
0.89 ms
4.48
FFT
RTX 3090
X
10
5.84
0.91 ms
6.42
FFT
V
11
10.67 ms
1.33 ms
7.99
FFT
RTX 2080
12
10.81 ms
1.35 ms
7.95
FFT
13
1.19 ms
0.99 ms
1.21
FFT
H100
X
14
7.40 ms
0.95 ms
7.72
FFT
(11008, 4096)
15
9.93 ms
2.35 ms
4.22
FFT
RTX 3090
X
16
14.89 ms
2.37 ms
6.28
FFT
V
17
27.62 ms
3.48 ms
7.93 x
FFT
RTX 2080
X
18
27.96 ms
3.49 ms
8.01
FFT
V
19
1.20 ms
1.41 ms
0.85
@
H100
20
7.34 ms
1.41 ms
5.20
FFT
(4096, 11008)
V
21
10.03
3.57 ms
2.81
FFT
RTX 3090
22
14.86 ms
3.59 ms
4.13
FFT
23
27.83 ms
5.68 ms
4.89
FFT
RTX 2080
X
24
28.20 ms
5.69 ms
4.95
FFT
transformer block, Im-head), We then train this model on a single GPU with sequence length 256 and batch size to minimize the time spent in forward and backward, as well as the memory usage.
Second, used the model with 800M parameters from Table in paper and trained it one GPU with sequence length 512 and batch size to replicate the previous setup.
For both scenarios measure and report the running time of the optimizer step. This way, want to show the overhead of low-rank variants with respect to standard, full-rank Muon.
We test full-rank Muon and low-rank optimizers Dion and Trion (with matmul and Makhoul's algorithm) with rank 512 and 1024 to show the overhead of each setting. We perform experiments on GPU of type Nvidia H100 and Nvidia RTX 3090 because H100 has fast tensor that make matmul faster, while RTX 3090 has slower tensor cores that benefit Makhoul's algorithm.
Below we present the results for this synthetic benchmark that emphasizes the overhead of the optimizer step for Trion and Dion in comparison to full-rank Muon:
in Table 7 we show the benchmarking results for Llama-800M on H100. We observe that matmul is faster than Makhoul's algorithm for both ranks, which is expected because H100 has powerful tensor cores. Trion achieves lower rumning time than the full-rank Muon, while Dion is very expensive because of the Qe-ocoompointion and our Dynamic Column Selection successfully replaces it.
Table we show the benchmarking results for Llama-800M on RTX-3090. We observe that matmul is slower than Makhoul's algorithm for both ranks, which is explained by the slower tensor-cores for this older GPU. Trion achieves lower running time than full-rank Muon for both ranks.
Table we show the benchmarking results for our custom Llama-464M on H100. As in the previous case, matmul is faster H100 than Makhoul's 's algorithm. Since the layers are larger,
17

---

observe some increase in the overhead for low-rank methods. Still, Trion is much cheaper than Dion, whose running time increases significantly with the rank. in Table 10 we benchmark our custom Llama-464M
RTX-3090. Similarly, Makhoul is slightly faster than matmul and the low-rank methods have an additional overhead compared to Muon. Still, Dion is the most expensive also in this case, which is justified by the calls QR-decomposition.
Table 7: [H100) Llama-800M (d 2048, layers=16, batch-size=l, seq_len=112)
Method
Rank
Optimizer Step
Memory
Time Overhead
Memory Overhead
Muon
full
113 ms/it
19.07 GB
1x
1x
Trion+matmul
512
67 ms/it
18.61 GB
0.59
0.98
Trion+Makhoul
512
78 ms/it
19.02 GB
0.69
1
Dion
512
246 ms/it
20.47 GB
2.18x
1.07
maamnll
1024
90 ms/it
18.87 GB
0.80
0.99x
Trion+Makhoul
1024
101 ms/it
19.02 GB
0.89
1x
Dion
1024
570 ms/it
20.65 GB
5.04
1.08%
Table 8: [RTX-3090] Llama-800M (d 2048, layers=16, batch-size=l, seq len=51)
Method
Rank
Optimizer Step
Memory
Time Overhead
Memory Overhead
Muon
full
830 ms/it
18.46 GB
1x
1x
Trion+matmul
512
355 ms/it
19.06 GB
0.43 x
1.03
Trion+Makhoul
512
315 ms/it
18.66 GB
0.38
1.01
Dion
512
485 ms/it
19.59 GB
0.58x
1.06x
Trion+matmul
1024
551 ms/it
19.06 GB
0.66x
1.03x
Trion+Makhoul
1024
507 ms/it
18.66 GB
0.61
1.01
Dion
1024
1010 ms/it
19.84 GB
1.22x
1.07x
Table 9: [H100] Llama-464M (d 4096, layers=1, batch-size=1, seq len=256)
Method
Rank
Optimizer Step
Memory
Time Overhead
Memory Overhead
Muon
full
3.19 ms/it
13.01 GB
1x
Trion+matmul
512
4.07 ms/it
13.03 GB
1.28
1x
Trkhou
512
4.41 ms/it
13.48 GB
1.38x
1.04x
Dion
512
11.95 ms/it
15.36 GB
3.75x
1.18x
Trion+matmul
1024
4.02 ms/it
13.03 GB
1.26x
1x 1.04x
Trion+Makhoul
1024
4.45 ms/it
13.48 GB
1.39
Dion
1024
22.60 ms/it
15.63 GB
7.08 x
1.2x
G PSEUDOCODE OF DCT-ADAMW OPTIMIZER
In this section we present the pseudocode of the DCT-AdamW optimizer, which uses our DCT-based dynamic column selection approach create dynamic, low-rank projection matrix use to factorize the gradient to low-rank matrix.
Following LDAdamW (Robert et al., 2025), different subspaces and this requires rotating the momentum buffers using rotation matrix R. where the second momentum buffer is updated using simpler rule compared to LDAdamW. DCT-Adam also supports error feedback and it saves only the enror induced by the low-rank compression.
incorporates low-rank gradients from
18

---

Table 10: [RTX-3090] Llama-464M (d 4096, layers=1, batch-size=l, seq len=26))
Method
Rank
Optimizer Step
Memory
Time Overhead
Memory Overhead
Muon
full
6.70 ms/it
12.98 GB
1x
1x
Trion+matmul
512
8.01 ms/it
12.37 GB
1.2
0.95x
Trion+Makhoul
512
7.89 ms/it
12.95 GB
1.18x
1x
Dion
512
13.26 ms/it
14.58 GB
1.98x
1.12x
Trion+matmul
1024
7.01 ms/it
12.37 GB
1.05 x
0.95x
Trion+Makhoul
1024
5.78 ms/it
12.95 GB
0.86
1x
Dion
1024
57.29 ms/it
14.82 GB
8.55x
1.14x
The sets I /prew/ort hold the indices of the r columns for the previous/current projections and Q Qreen/ert contain the columns from Q specified by the these two sets. T, represents the subspace update interval (set to 200 for Grre and for L.Adam). The procedure RANKCOLS ranks the columns of S, which is computed as S MAKHOUL(G)) (see Appendix D) S G. Q explained in Section 2.1. In order to make sure the momentum buffers integrate gradients from the same subspaces, need to rotate m, and tt using rotation matrix R that first projects the momentum to the full-dimensional space using the previous projector and then projects it back to the current liver-dimensional subspace. We can perform this rotation directly in the r-dimensional space by multiplying the two projection matrices, resulting in R E Rrxr, Rotating v might introduce negative values and apply the absolute value function to force the non-negativity of ey. The low-rank gradient 9r is computed by multiplying the full-rank gradient G, with Qen and the full-rank update obtained multiplying ut by Qert!
Algorithm 2 DCT-AdamM (right projection)
Algorithm 3 Update Subspace procedure
1: Input: B1,
1: procedure UPDATESUBBSPACECC)
2: mg. Vu + Onxr: Onxr
+ I,xr
3: E1- + -Oxxm
error feedback (EF) buffer
3:
if > 1 then
4: QE Rmxm
o DCT matrix o column indices
Iprew Iert
5: Tort, Iprev <- 0,, 0,
5:
end if
6: for
T) do
if (1 1) (E mod T, 0) then S. MAKHOUL(G) & or S = GQ
7:
Ge + Vef(0,)
8:
R+ UPDATESUBSPACEE(H
Iert +- -RANCCOLS(GQ.r) R QpQep Qart
9:
9i G. Qert
projected gradient o update EF 10:
10:
E, t Gt - 9t
QT.
end if
m. Bi m R
(1-Bi)94
11:
return R
12:
Ut & 321v1-1 R + 32)9i
12:
end procedure
13:
0c+1 + O2 ctvo
th
Q
14: end for
H PROJECTION ERRORS OF TRION AND DION
We train two Llama models with 30M parameters (dmodel = 640) with rank r= 128 using Dion and Trion optimizers and save the gradients. For each set of gradients we simulate the operations of the other optimizer and record the metrics A dion and A Atrion Our experiments showed no significant differences in the patterns for the two simulations and we plot the projection errors for the simulation that used gradients from the model trained with Dion optimizer. In Figure 2 we show the projection errors for the linear layers in the first transformer block, where Trion yields lower projection error than Dion and the trend is the same across all transformer blocks in the model. These plots explain why Trion has lower training and validation perplexity, as can be seen in Table 1. Moreover, the projection error seems to be constant for Dion, while for Trion the projection error has decreasing trend for some layers, which once again shows the dynamic behavior of our column selection approach coupled with Newton-Schull iteration.
19

---

I PRE-TRAINING WITH FRUGAL/FIRA
Pre-training with FRUGAL. We integrate the DCT matrix into the FRUGAL optimizer and compare against the SVD, RandPerm and Random projections. RandPerm uses a random permutation as a projection matrix, while Random generates random, semi-orthogonal matrix. In Figure 4a show the training loss for all these runs. In the zoomed-in plot for the last 5000 training steps we see that SVD projection recovers the performance full-rank AdamW, while the DCT projection a good approximation of the SVD, which supports our claim. We present our numerical results in Table 11. We would like to emphasize the running time reduced by Ih 48m (= 22.6%) compared to the SVD projection, as well as the memory usage reduced by about 2.2GB ( 3.5%). while the increase in perplexity is less than one point. comparison RandPerm and Random projections, the runtime on par, while train and validation perplexities are lower by approximately one point in the favor of DCT, illustrating again the benefit of our approach.
Pre-training with FIRA. We integrate DCT into FIRA optimizer and compare against SVD. In Figure 4b we show the training loss, where our focus is on the comparison between SVD and DCT projections. We observe that DCT consistently yields lower training loss compared projection, which is also visible in the numerical results in Table 11, translated to lower perplexities in the favor of DCT. Moreover, the memory usage is smaller by 2GB (=s 3%) and the running time is lower by 1h 54m (== 23.8%).
the SVD
Table 11: Pre-training Results for AdamW, FRUGAL and FIRA with different projections using 20 tokens/parameter. DCT is good approximation SVD with lower runtime and memory. Adam is the full-rank optimizer and is added for reference.
FRUGAL
FIRA
AdamW
SVD
DCT
RandPerm
Random
SVD
DCT
Train PPL
15.55
15.35
15.63
16.52
17.02
19.37
18.88
Val. PPL
14.05
14.02
14.23
15.18
15.61
17.67
17.30
Mem. (GiB)
73.49
65.70
63.50
65.44
63.72
68.44
66.48
Time
7h 54m
9h 45m
57m
7h 56m
7h 56m
9h 53m
59m
5.0
5.0
FUUSALSV
F.k AdaofE
FRASET
4.5
4.5
Training Loss
2.79
Training Loss
4.0
4.0
3.5
55030
3.5
5900 35000 00307 RS0C
3.0
3.0
10000
20000
30000 iter
40000
50000
60000
2.5
10000
20000
30000 iter
40000
50000
60000
(a) FRUGAL
(b) FIRA
Figure 4: Pre-training Llama-800M using FRUGAL and FIRA on 16B tokens from C4.
J PRE-TRAINING WITH DCT-ADAMW AND SUBTRACK++
In this section we provide additional ablations.
DCT-AdamW vs aubTrack++ Rajabi et al. (2025). In Table 12 we provide some comparison of DCT-AdamW and SubTrack++ on a model with 350M parameters, rank 256 (25% of model's embedding size) and preconditioner update frequency 200. We run three experiments with different random seeds for each optimizer and report the average metrics. We observe that our DCT-AdamW achieves significantly lower values for train and validation metrics with slightly faster runtime and memory usage, which proves that our Dynamic Column Selection approach effective in tracking the low-rank subspace of the gradient.
20

---

Table 12: Comparison between DCT-AdamW and +ubTrack++
Metric
DCT-AdamW
SubTrack++
Train Loss
2.738
2.779
Train PPL
15.48
16.13
Val Loss
2.747
2.789
Val PPL
15.60
16.266
Runtime
6h 57m
7h 15m
Memory Usage (per GPU)
42.99 GB
46.95 GB
K FINE-TUNING
We fine-tune Llama-2 7B on GSM-8k dataset using SVD and DCT projections for FRUGAL and FIRA optimizers, as well as LD-AdamW and DCT-AdamW. We show our results for ranks r {32, 512) in Table 13. We also finetune Dwen-2.5-7B on GSM-8k using OCT-AdamMK GaLore and AdamW for reference and show the results in Table 14 for the same ranks.
Fine-tuning with FRUGAL. Both projections yield similar training loss for both rank values, which an indication that DCT indeed good approximation for SVD. Despite having different training loss, the SVD projection achieves the same accuracy for both ranks. It is surprising that DCT yields better accuracy for lower rank compared to higher rank. terms of memory usage, as expected for this large model, the DCT projection saves 8GB of memory (*2 23.4 23.4%) for r = 32, while for r = 512 the reduction in memory is only 6.3GB (= 18.2%). The running time is reduced by roughly 35m (== 75%) for both ranks.
Fine-tuning with FIRA. The DCT projection yields larger training loss compared to SVD. DCT recovers the accuracy, outperforms the SVD for large rank, and average reduces the memory usage by ~ 1.35 GB and the running by 2s 110m.
Fine-tuning with LDAdamW W & DCT-Adam W. In this setting compare the DCT projection with the block power-iteration used in LDAdamW as an approximation to SVD, both without EF. The training loss achieved is smaller for DCT on 32 and larger for 512 compared LDAdamW. LDAdam achieves more than 1% accuracy for r 32 and comparable accuracy for = 512. Regarding memory, DCT-AdamW W uses =1 GB less memory, while achieving speedup of about 20m for r = 512. We would like to emphasize that EF does not help DCT-Adamll W in comparison to LDAddm for r = 32, the accuracy of DCT-AdamW with EF being 29.11% vs 32.53% for L.Ddaa with EF for = 32, while for r = 512 the accuracy can be recovered, scoring 35.33% compared to 35.86% for LDAdam W with EF.
Fine-tuning Qwen-2.5 7B. In this setting compare the DCT projection with the original GaLore with our DCT-Adam W without EF, where both optimizers update the subspace once at 200 steps. In Table 14 shows that despite having slightly higher training loss, the DCT-AdamW has higher evaluation accuracy than GaLore.
Table 13: Fine-tuning results for Llama-2 7B on the GSM-8K dataset.
FRUGAL
FIRA
LD/DCT-Ada
rank 32
rank 512
rank 32
rank 512
rank 32
rank 512
SVD
DCT
SVD
DCT
SVD
DCT
SVD
DCT
LD
DCT
LD
DCT
Train Loss
0.046
0.059
0.051
0.071
0.123
0.209
0.094 0.192
0.261
0.176
0.101
0.208
Acc. (%)
33.81
35.93
33.81
34.26
32.15
32.45
34.27
35.25
31.38
30.09
35.61
35.17
Mem. (GB)
39.61
31.59
40.68
34.41
32.72
31.29
35.41
34.11
32.08
31.17
35.58
34.35
Running Time
1h 22m
46m
1h 19m
47m
1h 8m
1h
1h 9m
1h
55m
48m
1h 8m
48m
21

---

Table 14: Fine-tuning results for Qwen-2.5-7B on AdamW, DCT-AdamW and GaLore GSM-8K. DCT is good approximation to SVD with lower runtime and memory. AdamW is the full-rank optimizer and is added for reference.
AdamW
DCT-AdamW
GaLore
full rank
rank 32
rank 512
rank 32
rank 512
Train Loss
0.012313
0.053792
0.0545
0.0388
0.0184
Eval Acc (%)
67.70
64.59
65.35
63.23
64.41
Memory (GB)
64.5
40.4
44
40.9
45.15
Running Time
48m
49m
47m
58m
58m
22


Reviews:
[{"review_id": "6z3aUMR4lI", "reviewer": "Reviewer_ZMSS", "summary": "Low-rank optimization can speed LLM training and cut optimizer memory, but per-layer SVD/QR gradient projections are costly and require storing projection matrices. We propose a simple DCT-based alternative: multiply each layer’s gradient by a fixed orthogonal DCT basis, then rank-select the most aligned columns to form the projection. The DCT is efficiently computed (via FFT-based routines), precomputed once, and reused—so projections need only a matmul plus lightweight sorting. Across pre-training and fine-tuning, this yields rank-independent runtime, matches SVD/QR accuracy, and achieves faster training with lower memory use.", "questions": "1. **Error-feedback memory overhead**. Trion’s error-feedback buffer appears to store a full-size residual per parameter. Does this negate the memory savings from low-rank gradient projection?\n\n2. **Low-rank during forward–backward**. Beyond optimizer/state compression, can the low-rank structure be exploited to reduce the dominant forward–backward costs (FLOPs and memory)—e.g., via factored weight updates or structured bases that lower gradient-computation cost?\n\n3. I noticed that the work [R1] uses a similar idea to save SVD computations. Could the authors highlight the difference from [R1]\n\n[R1] Wavelet Meets Adam: Compressing Gradients for Memory-Efficient Training", "limitations": "", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- Provided a two-stage DCT-based projection method to avoide the computational cost of SVD\n- Developed the DCT variant Trion and DCT-AdamW\n- Demenstrated the contractivenss of the proposed compressor\n- Conduct extensive experiments", "weaknesses": "1. **Effectiveness in tracking gradients.** Constructing SVD-free, contractive projection matrices is straightforward (e.g., random projections). The real challenge is to remain SVD-free and contractive *while* faithfully capturing the gradient’s low-rank structure. The proposed two-stage method fixes a DCT basis $D_C$ for the entire training and, at each iteration, selects only a few of its columns. In effect, it tracks gradients using subsets of a preset basis—an approach not obviously aligned with evolving, layer-specific low-rank subspaces. The paper does not explain why this selection should match the true gradient subspaces or under what conditions it would; clearer insight or evidence (e.g., principal-angle analyses or SVD/QR approximation errors over training) is needed. \n\n2. **Necessity of compressing optimizer states.** ZeRO-style sharding already reduces optimizer-state memory by ~1/N per data-parallel replica without degrading quality. In Figure 2, DCTAdamW underperforms AdamW, suggesting that extra low-rank compression of states may be unnecessary—or even harmful—unless it delivers clear end-to-end gains. I personally think compressing optimizer states is unnecessary due to ZeRO optimizer. \n\n3.  **Necessity of saving computations in Newton–Schulz.** In Trion/Muon-style preconditioning, forward–backward passes dominate step time; Newton–Schulz iterations are typically a small fraction—especially with long context windows. The paper should provide profiler traces showing that low-rank $b_t$ meaningfully reduces *end-to-end* step time or time-to-target. Please also compare with vanilla Muon (no low-rank) to assess any convergence slowdown from low-rank preconditioning, and include both complexity estimates (per-token FLOPs) and wall-clock measurements to substantiate the claimed benefit.", "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": "KwXkmqNfXU", "reviewer": "Reviewer_akDN", "summary": "The paper proposes replacing per step SVD or QR low rank projections in adaptive optimizers with a fixed orthogonal basis using DCT and dynamic column selection. At each step the method scores basis columns by gradient to basis correlation and selects the top r to form projections, which avoids repeated factorizations and heavy state storage. Two instances, Trion and DCT AdamW, show lower memory use, faster training, and comparable or better perplexity on mid size LLM pretraining. The theory motivates norm based selection and provides simple error bounds. The implementation also explains how to integrate with distributed training to reduce communication.", "questions": "Please see weakness.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The paper uses a precomputed DCT with on the fly selection to replace repeated SVD or QR, and only column indices are stored. Column norm ranking aligns with minimizing reconstruction error, and there is clear intuition for why DCT approximates dominant directions.\n\nThe method has consistent memory reduction and wall clock speedups while matching or improving perplexity across several model sizes.\n\nThe notes on DDP or FSDP usage and local reconstruction make adoption straightforward.", "weaknesses": "Evidence for scaling to very large hidden sizes remains limited. The paper primarily reports results on mid size models where the benefit of fast transforms over plain matrix multiplication is muted. Without end to end wall clock measurements at dimensions around 8k to 16k, and without a breakdown of time in the similarity computation, basis selection, and reconstruction kernels, it is hard to assess whether the claimed speedups persist when layers become wide and deep. A thorough profiling study across model width, batch size, and rank would make the efficiency claims more convincing.\n\nIt is unclear whether all systems level optimizations such as ZeRO style redundancy removal, identical communication and precision settings, and error feedback or quantization choices are enabled symmetrically for both the proposed method and the baselines. Differences in these controls can easily dominate the observed speed or memory gains. The paper should report results under strictly matched configurations and, if desired, separately include best tuned variants for each method.\n\nThe paper mixes square and rectangular gradient matrices without clearly specifying when to apply left versus right projections, and how this choice is made across different layer types such as attention projections and output layers. The exact dimensional assumptions of the DCT matrices and the consistency of symbols drift across sections, which complicates reproduction and theoretical interpretation. A precise, layer wise rule set and a short ablation on these choices would improve clarity.\n\nThe set of comparative baselines is too narrow to establish superiority. Strong structured projections such as Hadamard or FWHT, CountSketch style mappings, and recent online or streaming low rank methods are not evaluated side by side under matched rank and refresh frequency. Because many of these alternatives also offer O(n log n) or near linear time with tiny memory footprints, omitting them leaves open whether DCT based selection is uniquely effective. Head to head comparisons on identical tasks and budgets are needed to justify the design choice.", "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": "wzyUN7e9aq", "reviewer": "Reviewer_tJpX", "summary": "The authors introduce a Discrete Cosine Transform (DCT)-based dynamic column selection technique that approximates optimal low-rank projections by selecting columns of a fixed orthogonal DCT matrix aligned with each layer’s gradients to replace expensive SVD/QR-based low-rank projections used in adaptive optimizers for large language models. This method reduces computation and memory overhead by avoiding per-layer decompositions and storing only column indices. They apply this idea in two new optimizers: Trion, which improves Dion, and DCT-AdamW, a low-rank variant of AdamW.", "questions": "Please refer to the weaknesses.", "limitations": "", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The work targets a bottleneck in large-model training: the cost of SVD/QR-based low-rank projections used in adaptive optimizers.\n2. The proposed DCT-based projection method is straightforward and easily integrable into existing optimizers.\n3. The approach achieves good efficiency gains.\n4. The paper provides mathematical rationale for why DCT approximates the gradient eigenbasis and the effectiveness of norm-based selection. \n5. The presentation is easy to follow and the paper is well-written.", "weaknesses": "1. All experiments stop at 1.3B-parameter models; there is no validation on pretraining larger LLMs (7B+), which undermines claims of scalability.\n2. The paper lacks certain ablations of critical design choices (e.g., norm type, rank sensitivity, DCT variant).\n3. Distributed and FSDP discussions are not empirically backed with wall-clock or communication cost benchmarks.\n4. While the paper targets efficiency, comparison with fast efficient baselines that are not using SVD/QR decomposition like APOLLO [1] and SubTrack++ [2] would help to strengthen the paper. \n---\n[1] Zhu et al., 2025. APOLLO: SGD-like Memory, AdamW-level Performance.\n[2] Rajabi et al., 2025. SubTrack++: Gradient Subspace Tracking for Scalable LLM Training", "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": "T0MTrWnUtR", "reviewer": "Reviewer_F7Ct", "summary": "This paper tackles the high compute and memory cost of SVD/QR decompositions in low-rank optimizers  The authors propose replacing these expensive, per layer projections with a \"dynamic column selection\" from a single, predefined orthogonal basis (the Discrete Cosine Transform, or DCT, matrix) computed once at the start . The method efficiently selects the top $r$ most aligned DCT basis vectors for each layer's gradient, replacing SVD/QR with a fast matrix multiplication (or FFT) and a sorting step. This saves memory, as each layer only stores $r$ indices instead of a full projection matrix. The authors integrate this technique into two new optimizers: Trion (improving Dion) and DCT AdamW (improving AdamW variants).", "questions": "Please see Weaknesses section", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The method's runtime is rank independent, whereas SVD/QR based methods like Dion get slower as the rank $r$ increases. This is a significant practical advantage for using larger and more expressive ranks.\n\n2. The core idea is simple yet elegant. The \"dynamic column selection\" is not just heuristic; the authors prove it is the optimal strategy for minimizing reconstruction error, given a fixed orthogonal basis $Q$. The motivation for using the DCT as that basis is also well justified.\n\n3. Strong Empirical Results: \n\ni. Trion vs. Dion (Table 1): Trion consistently achieves better validation perplexity, lower memory usage (~7-10%), and faster runtimes (up to 18% faster) than its direct baseline, Dion, across all tested model sizes and ranks.\n\nii. DCT-AdamW vs. LDAdamW: DCT-AdamW achieves better validation perplexity, drastically lower memory, and is significantly faster (~25%).", "weaknesses": "1. Theoretical speedup not realized. The paper heavily motivates using DCT by citing the fast $O(n^2 \\log n)$ FFT-based algorithm. However, the authors admit that for the model sizes tested (up to $d=2048$), this speedup was not significant, and a standard $O(n^3)$ matmul was used15. The primary speedup comes from replacing SVD/QR, not from the FFT.\n\n2. While DCT-AdamW beats its low-rank competitor LDAdamW, it is still significantly outperformed by full rank AdamW (Val. PPL 13.69 vs. 11.73). This shows the method is a better low rank compromise, but still a compromise.\n\n3. As the authors note, experiments are limited to 1.3B models. The method's true scalability and the benefit of the FFT based algorithm would only be evident on much larger models.", "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": "dNq0Z2uamU", "reviewer": "Reviewer_kABp", "summary": "In this paper, the authors have attempted to address the computational and memory overhead in optimizers such as AdamW, Muon, Dion, and GaLore in LLM training. Prior low rank method relies on using SVD or QR decomposition to project gradients into a lower dimensional subspace. However, they are computationally expensive. Therefore, in this paper, the authors attempt to replace these with an alternative low-rank projection approach, which is cheaper to compute.", "questions": "Please see the weaknesses.", "limitations": "", "rating": 8, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The strengths of this paper are summarized as follows:\n\n1. SVD is time-consuming in LLM training, and it is widely used in various methods, like GaLore, FIRA, and FRUGAL. Replacing SVD with the Fast Fourier Transform based algorithm can reduce the time complexity.\n\n2. Empirically, it has shown improvements on multiple optimizers and model sizes, like LLaMA 350M, 800M, and 1.3B. Also, on all of these model sizes, Trion has shown better performance and better running time compared with Dion. \n\n3. It has a strong theoretical guarantee for justifying the column selection approach may give the most significant column.", "weaknesses": "The weaknesses of this paper are summarized as follows:\n\n1. The largest model size in the experiment is 1.3B. It would be better if the authors may consider running an experiment on larger models since modern transformer architectures are getting much larger than this size. Also, only the C4 dataset and LLaMA are tested and there is no fine-tuning or downstream benchmarks.\n\n2. Galore has numerous follow-up works, such as Galore 2 [1], Golore [2], and Sara [3]. Could the authors give a comparison with these works?\n\n[1] DiJia Su, Andrew Gu, Jane Xu, Yuandong Tian, and Jiawei Zhao. \"Galore 2: Large-scale llm pre-training by gradient low-rank projection.\" arXiv preprint arXiv:2504.20437 (2025).\n\n[2] Yutong He, Pengrui Li, Yipeng Hu, Chuyan Chen, and Kun Yuan. \"Subspace optimization for large language models with convergence guarantees.\" ICML'25.\n\n[3] Haochen Zhang, Junze Yin, Guanchu Wang, Zirui Liu, Tianyi Zhang, Anshumali Shrivastava, Lin Yang, and Vladimir Braverman. \"Breaking the Frozen Subspace: Importance Sampling for Low-Rank Optimization in LLM Pretraining\". NeurIPS'25.", "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)