FFT-based Dynamic Subspace Selection for Low-Rank Adaptive Optimization of Large Language Models

ArXiv: 2505.17967

🎯 Pitch

A simple two-step trick replaces expensive per-layer SVD/QR decompositions in low-rank optimizers with a single fixed DCT matrix, cutting memory up to 25% while perfectly matching Adam's loss curves. The secret is dynamically selecting DCT basis columns by their alignment with each layer's gradient, making the method both rank-independent and nearly as fast as full-rank Adam.

---

1. Executive Summary

This paper proposes a computationally efficient two-step procedure for approximating SVD/QR-based gradient projections in low-rank adaptive optimization of large language models, using a predefined orthogonal matrix from the Discrete Cosine Transform (DCT). The key mechanism is dynamic column selection — adaptively choosing columns from the fixed DCT matrix based on their alignment with each layer's gradient (measured via $\ell_1$ or $\ell_2$ norm of the similarity matrix $S = GQ$), which replaces the expensive Power-Iteration in Dion and the SVD/QR decompositions in FRUGAL, FIRA, and LDAdam with a rank-independent alternative. The authors evaluate their approach on Llama models ranging from 350M to 1.3B parameters pretrained on C4 under Chinchilla-optimal token budgets, and on Llama-2 7B fine-tuned on GSM-8k, demonstrating that their DCT-based projection matches or exceeds the accuracy of SVD/QR-based methods while reducing memory usage by up to ~25% and runtime by up to ~18% for higher ranks, establishing that a fixed Fourier-type basis with dynamic column selection can serve as an accurate and cheaper substitute for per-layer matrix factorizations — provided the embedding dimension remains modest enough that the DCT materialization cost does not dominate.

2. Context and Motivation

The Core Problem: Low-Rank Optimizers Depend on Expensive Matrix Factorizations

The fundamental challenge this paper tackles is deceptively simple: when we compress optimizer states using low-rank projections, we need orthogonal projection matrices, but computing them via SVD or QR decomposition at every layer is too expensive. This matters because the recent wave of low-rank adaptive optimizers — GaLore, LDAdam, FRUGAL, FIRA, Dion — all share a common bottleneck: they must periodically (or at every step) compute expensive matrix factorizations to find the low-dimensional subspaces in which the optimizer states will live. The paper's core insight is that we don't need to compute these projection matrices from scratch for each layer. A single, fixed orthogonal matrix — specifically, a Discrete Cosine Transform (DCT) matrix — combined with a cheap column selection procedure can replace these per-layer factorizations entirely.

This gap is significant for several practical reasons the authors highlight implicitly throughout their introduction and experiments:

• Runtime dependence on rank: In optimizers like Dion, the Power-Iteration procedure includes a QR decomposition whose running time scales with the chosen rank $r$. As models scale to larger embedding dimensions and practitioners want to use higher ranks for better accuracy, this per-step factorization becomes a growing fraction of total training time. Table 1 shows this concretely: Trion's runtime is approximately constant across ranks 128, 256, and 512 for a given model size, while Dion's runtime increases with $r$.

• Memory overhead from per-layer projection matrices: Methods like LDAdam must store two consecutive projection matrices per layer to enable smooth subspace transitions (rotating momentum buffers when the subspace changes). For a model with embedding dimension $d$, each projection matrix of size $d \times r$ consumes $O(dr)$ memory per layer. Across hundreds of layers in modern LLMs, this adds up. The paper reports that DCT-AdamW stores only $r$ integers (column indices) per layer instead of the full projection matrix, plus a single shared $d \times d$ DCT matrix per GPU.

• Communication bottlenecks in distributed training: When using Newton-Schulz iteration for orthogonalization (as in Muon and Dion), full-size matrices must be materialized on GPU before the iterative procedure runs. This increases communication costs in distributed settings. The paper's Trion optimizer orthogonalizes a low-rank $b_t \in \mathbb{R}^{R \times r}$ matrix instead of the full momentum $B_t \in \mathbb{R}^{R \times C}$, and communicates only the low-rank terms across devices rather than full orthogonal updates.

The Landscape Before This Paper: A Proliferation of Factorization-Based Methods

The paper's introduction and related work section (Section 5) map out a clear lineage of approaches that all converge on the same bottleneck:

GaLore (Zhao et al., 2024) was the pioneering work that introduced low-rank gradient projection for memory-efficient training. It performs SVD on the gradient matrix once every ~200 steps to compute projection matrices, then runs AdamW in the compressed subspace. The key limitation: SVD is $O(n^3)$ and must be invoked for every linear layer, even if only periodically. The paper notes that the default update frequency of 200 steps was chosen precisely to make SVD "computationally feasible" — implying it would be prohibitively expensive at every step.

LDAdam (Robert et al., 2025) improved GaLore by compressing the first-order momentum (not the raw gradient), replacing SVD with Block Power-Iteration, and performing smooth subspace transitions by rotating momentum buffers when the projection changes. This enabled updating the subspace at every step (frequency = 1), which improves convergence. But the cost: storing two consecutive projection matrices per layer (current and previous), plus the runtime overhead of Block Power-Iteration, which still involves QR decompositions whose cost depends on $r$.

FRUGAL (Zmushko et al., 2024) and FIRA (Chen et al., 2024), concurrent works, both use SVD-based projections every ~200 steps but differ in how they handle the projection error — FRUGAL feeds it to a stateless optimizer like SignSGD, while FIRA scales it appropriately to recover full-rank performance. Both inherit SVD's computational cost.

Dion (Ahn et al., 2025) — a direct predecessor to this paper's Trion — uses Power-Iteration with QR decomposition at every step to compute low-rank, orthogonal momentum updates. The runtime depends on the rank $r$ because QR operates on matrices whose dimensions scale with $r$. This is the method Trion most directly replaces.

Muon (Jordan et al., 2024) — while not a low-rank method — uses Newton-Schulz iteration to orthogonalize the full momentum matrix, achieving fast convergence by pushing singular values toward 1. But Newton-Schulz involves odd powers of full-size matrices, making it expensive at scale. Dion attempted to reduce this by making the updates low-rank, but retained QR as the factorization method.

Where Prior Approaches Fall Short: A Shared Bottleneck

Despite their diversity in how they handle projection error, subspace update frequency, and optimizer state structure, all these methods share a common assumption: the orthogonal projection matrices must be computed from the gradient or momentum data itself, either via SVD (GaLore, FRUGAL, FIRA), Block Power-Iteration (LDAdam), or Power-Iteration with QR (Dion). This assumption creates several concrete problems:

1. The factorization cost is redundant across layers. Each linear layer computes its own SVD or QR independently, even though all layers' gradients share statistical structure from being part of the same forward-backward pass. The paper's key departure is recognizing that a single, carefully chosen, fixed orthogonal basis can serve all layers simultaneously — only the selection of which basis vectors to use needs to be per-layer.

2. The factorization cost scales with $r$. In Dion, the QR decomposition operates on matrices whose dimensions depend on the chosen rank. As practitioners push to higher ranks for better accuracy (the paper experiments up to $r = 512$, or $r/d = 1/2$), this cost grows. Table 1 quantifies this: at $r = 512$, Trion is 8–18% faster than Dion across model sizes, with the gap widening as $r$ increases.

3. Per-layer projection matrices consume memory. LDAdam stores two $d \times r$ matrices per layer. For a deep model with embedding dimension $d = 2048$ and rank $r = 256$, this is approximately $2 \times 2048 \times 256 \times 4$ bytes $\approx 4$ MB per layer. Across 32 transformer layers, that's ~128 MB just for the projection matrices — and this grows linearly with $d$ and $r$. In the fine-tuning experiments on Llama-2 7B (Appendix H, Table 7), the DCT approach saves 8 GB of memory (23.4%) compared to SVD-based projections at $r = 32$.

4. Materializing full matrices for Newton-Schulz is expensive. Muon's Newton-Schulz iteration requires computing odd powers of the full momentum matrix, which must be materialized on GPU. Dion reduces this by making the updates low-rank before orthogonalization, but still uses QR to obtain the low-rank representation. The paper's Trion goes further: it obtains the low-rank representation via DCT column selection (no iterative factorization), then runs Newton-Schulz on the $R \times r$ low-rank matrix $b_t$ rather than the full $R \times C$ momentum $B_t$.

5. There is no rank-independent path to orthogonal updates. Prior to this work, if you wanted orthogonalized momentum (for faster convergence, as Muon demonstrated) in a low-rank form (for memory efficiency, as Dion attempted), you had to pay the runtime cost of QR or Power-Iteration, both of which scale with rank. The paper provides the first method where the factorization cost is rank-independent: computing similarities $S = GQ$ costs the same regardless of how many columns you subsequently select.

Conflicting Requirements: The SVD/QR Trilemma

The paper implicitly identifies a trilemma that prior work could not simultaneously resolve:

• Update frequency: Updating the subspace every step (like LDAdam and Dion) improves convergence because the optimizer tracks the gradient's changing principal directions. But it requires a cheap factorization method.
• Factorization accuracy: SVD gives the optimal low-rank approximation in Frobenius norm, but is too expensive to run every step. Power-Iteration is cheaper but less accurate — the paper shows in Figure 2 that Trion's DCT-based projection achieves lower reconstruction error than Dion's Power-Iteration, meaning the DCT approach is simultaneously faster and more accurate.
• Runtime independence from rank: QR and SVD costs scale with $r$. If you want high-rank updates (closer to full-rank AdamW performance), you pay proportionally more. This creates a practical ceiling on usable rank.

The DCT-based dynamic column selection resolves this trilemma by decoupling the three: it can run every step (frequency = 1), it provides competitive or better accuracy than Power-Iteration (Figure 2), and its cost is independent of $r$ because computing $S = GQ$ always produces the full $n \times n$ similarity matrix — selecting $r$ columns from it is just a top-$r$ indexing operation.

How This Paper Positions Itself

The paper does not propose a fundamentally new optimization algorithm or a new theory of low-rank training. Instead, it proposes a drop-in replacement for the orthogonal projection matrices that all existing low-rank optimizers depend on. The positioning is explicit in the introduction's research question:

"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-efficient optimizers?"

The answer is: use a fixed DCT matrix with dynamic column selection. The paper demonstrates portability by integrating this approach into five existing optimizers — Dion (becoming Trion), LDAdam (becoming DCT-AdamW), FRUGAL, FIRA, and even the full-rank Muon (by reducing Newton-Schulz input size). This is not a "new optimizer" paper; it is a "better factorization primitive" paper.

The theoretical contribution (Section 4) provides two justifications for why this works:

1. Optimality of norm-based column selection (Section 4.1): For any orthogonal matrix $Q$, selecting the $r$ columns with largest alignment $\|q_i^\top G\|$ minimizes the Frobenius reconstruction error $\|G - Q_r Q_r^\top G\|_F$. This is a simple but important guarantee — it means the selection procedure is not heuristic but provably optimal for a given $Q$.

2. DCT as linear approximation of gradient eigenbasis (Section 4.2): Using a decomposition result from optical information processing (Müller-Quade et al., 1998; Schmid et al., 2000), any square matrix can be factored into alternating diagonal and circulant matrices. Since circulant matrices are diagonalized by the DFT matrix, the eigenvectors of $GG^\top$ can be approximated by the Fourier basis. For real symmetric $GG^\top$, the real part of the DFT (which is the DCT, up to minor variations) serves as a natural approximate eigenbasis. This provides theoretical motivation for why DCT specifically — rather than a random orthogonal matrix — should work well: it approximates the SVD's left singular vectors without needing to compute them.

The paper also explicitly connects to the computational efficiency literature by noting that the DCT enables an FFT-based computation of the similarities via Makhoul's algorithm in $O(n^2 \log n)$ time instead of $O(n^3)$ for a general matrix multiplication (Appendix D), though the practical benefits of this acceleration are currently limited by PyTorch's lack of complex-bfloat16 support, meaning Makhoul's algorithm must run in float32 while the matmul baseline can use faster bfloat16 (see Tables 4 and 5 in Appendix D).

A Note on What This Paper Does NOT Claim

It is important to understand the scope boundaries the paper sets, because they define what "success" means:

• Not a replacement for full-rank training: The method is a cheaper way to compute low-rank projections, but low-rank training is inherently an approximation. The goal is to match or slightly exceed existing low-rank methods at lower cost, not to beat full-rank AdamW.
• Not limited to DCT: The dynamic column selection approach works with any orthogonal matrix (Section 2.1). DCT is chosen for its FFT-accelerable structure and theoretical connection to the gradient eigenbasis, but the framework is general.
• Not claiming large-model validation: Experiments top out at 1.3B parameters for pretraining and 7B for fine-tuning. The authors explicitly acknowledge this as a limitation in Section 6: "additional work is required to test our technique for larger models."
• Not solving the hard problem of very small models vs. large models: The DCT matrix of size $d \times d$ must be materialized once per GPU. For models with small embedding dimension $d$, this is negligible. For models with very large $d$ (e.g., $d = 8192$ or higher in modern LLMs), the $d \times d$ matrix itself becomes a memory concern — though the paper argues it is still cheaper than storing per-layer projection matrices for deep networks.

3. Technical Approach

3.1 Reader Orientation

This paper develops a drop-in replacement primitive for the orthogonal projection matrices used in low-rank adaptive optimizers. Instead of computing per-layer projection matrices from scratch via expensive SVD or QR decompositions, the system uses a single, fixed, pre-computed Discrete Cosine Transform (DCT) matrix shared across the entire model, and dynamically selects which $r$ columns of this matrix best capture each layer's gradient or momentum at each optimization step. The "shape" of the solution is: one global orthogonal basis (computed once) + a cheap per-layer, per-step column ranking + a lightweight indexing operation = a projection matrix $Q_r$ tailored to each layer without any iterative factorization.

3.2 Big-Picture Architecture (Diagram in Words)

The system has five major components that interact at each optimizer step:

1. The Fixed Orthogonal Basis (DCT Matrix) — a single $d \times d$ orthogonal matrix $Q$ (specifically DCT-III or DCT-II) materialized once per GPU at training start. It serves as a shared dictionary of basis vectors from which every layer selects its projection columns. It is the same size as the model's embedding dimension $d$.

2. The Similarity Computation — for each layer, the gradient (or momentum) matrix $G \in \mathbb{R}^{n \times m}$ (where $n \geq m$ typically, and $\min(n,m) = d$) is multiplied by the DCT matrix to produce a similarity matrix $S = GQ$ (or $S = \text{Makhoul}(G)$ using the FFT-based algorithm). Each column $j$ of $S$ contains the dot products between every row of $G$ and the $j$-th DCT basis vector — measuring how much of the gradient "points in" each basis direction.

3. The Dynamic Column Selection — the columns of $S$ are ranked by their $\ell_1$ or $\ell_2$ norm, and the indices of the $r$ largest-norm columns are selected. These indices identify which $r$ DCT basis vectors are most aligned with the gradient's structure for this specific layer at this specific step. The projection matrix $Q_r$ is formed by indexing those $r$ columns from the full DCT matrix.

4. The Low-Rank Projection and Error Handling — the gradient is projected into the $r$-dimensional subspace as $g = G Q_r \in \mathbb{R}^{n \times r}$, and the reconstruction error $\Delta = G - g Q_r^\top$ is computed. Depending on the optimizer (Dion-style vs. Adam-style), this error is either stored in the momentum buffer (Trion) or fed into an error feedback mechanism with optional 8-bit quantization (DCT-AdamW).

5. The Optimizer-Specific Update — the low-rank representation $g$ (or the low-rank momentum $b_t$ in Trion) is processed by the optimizer's core logic: Newton-Schulz orthogonalization for Trion, or AdamW moment estimation with subspace rotation for DCT-AdamW. The final parameter update is computed in the low-dimensional space and projected back to the original dimensions using $Q_r^\top$.

Information flows as follows: gradient $G_t$ arrives per layer → similarity matrix $S_t = G_t Q$ is computed (via matmul or Makhoul's algorithm) → top-$r$ column indices are selected by norm ranking → $Q_r$ is extracted by indexing $Q$ → gradient is projected to $g_t = G_t Q_r$ (or momentum is projected in Trion) → error $\Delta_t$ is computed and stored/quantized → optimizer state is updated in the $r$-dimensional space → parameter update is computed in $\mathbb{R}^{n \times r}$ and projected back via multiplication by $Q_r^\top$.

3.3 Roadmap for the Deep Dive

• First, the dynamic column selection mechanism in its general form (works with any orthogonal $Q$, any matrix $G$, any rank $r$) — this is the algorithmic core that everything else depends on, and it explains exactly how the column ranking minimizes reconstruction error.
• Second, the Discrete Cosine Transform — why DCT specifically is chosen over random orthogonal matrices or Hadamard matrices, and the computational implications (FFT acceleration via Makhoul's algorithm, $O(n^2 \log n)$ vs. $O(n^3)$).
• Third, Trion — the complete integration of DCT-based column selection into the Dion optimizer, replacing Power-Iteration and QR decomposition, with detailed walkthrough of Algorithm 1.
• Fourth, DCT-AdamW — the integration into AdamW-style low-rank optimizers (replacing SVD/Block Power-Iteration in LDAdam, FRUGAL, FIRA), including the subspace rotation mechanism, error feedback, and the UpdateSubspace procedure.
• Fifth, the theoretical guarantees — why norm-based column selection is optimal for minimizing Frobenius reconstruction error given any fixed orthogonal matrix, and why DCT specifically approximates the gradient's eigenbasis.
• Sixth, the distributed training communication design — how Trion leverages the shared DCT matrix to communicate only low-rank components across GPUs in DDP/FSDP settings.

3.4 Detailed Sentence-Based Technical Breakdown

This is primarily a systems and methods paper whose core idea is that a fixed, pre-computed orthogonal basis with dynamic column selection can replace per-layer matrix factorizations in low-rank adaptive optimizers, achieving rank-independent runtime and reduced memory while matching or exceeding the accuracy of SVD/QR-based approaches.

---

Dynamic Column Selection: The Core Algorithmic Primitive

The dynamic column selection procedure (Section 2.1) is the central mechanism that enables a single fixed orthogonal matrix to serve all layers adaptively. It answers the question: given a gradient matrix $G$ and an orthogonal matrix $Q$, which $r$ columns of $Q$ should we use to project $G$ with minimal information loss?

The setup. Let $Q \in \mathbb{R}^{n \times n}$ be any orthogonal matrix (so $Q^\top Q = I_n$). Let $G \in \mathbb{R}^{n \times n}$ be the gradient matrix (the paper also handles rectangular $G \in \mathbb{R}^{n \times m}$ by compressing along the smaller dimension, typically $d_\text{model}$). The goal is to select $r$ columns from $Q$ — forming a submatrix $Q_r \in \mathbb{R}^{n \times r}$ — such that projecting $G$ into the span of those columns and back yields a good approximation of the original $G$.

Step 1: Compute similarities. The similarity matrix $S$ is computed as:

$S = G Q \in \mathbb{R}^{n \times n}$

where the $i$-th row and $j$-th column entry $S_{ij}$ is the dot product between the $i$-th row of $G$ (a vector in $\mathbb{R}^n$) and the $j$-th column of $Q$ (a basis vector in $\mathbb{R}^n$). Equivalently, each column $j$ of $S$ contains the scalar products of every row of $G$ with the $j$-th basis vector $q_j$ of $Q$. The paper clarifies this with an explicit matrix form in Appendix B (Equation 3), where $S$ is shown as a matrix whose columns are $[G_1^\top q_1, G_1^\top q_2, \dots, G_1^\top q_n]$ stacked vertically for each row of $G$.

Step 2: Rank columns by alignment strength. For each column $j$ of $S$, compute its $\ell_1$ norm (sum of absolute values) or $\ell_2$ norm (Euclidean length). The $\ell_2$ norm of column $j$ is:

$\|S_{:,j}\|_2 = \sqrt{\sum_{i=1}^n (G_i^\top q_j)^2}$

This measures the total alignment between the $j$-th basis vector $q_j$ and all rows of $G$. A large norm means that $q_j$ captures a significant portion of the gradient's energy across many rows. The paper uses $\ell_1$ or $\ell_2$ norm interchangeably (Section 2.1), noting that both produce similar rankings in practice.

Step 3: Select top-$r$ columns. The columns of $S$ are sorted by their chosen norm in descending order. The indices of the $r$ columns with the largest norms are stored in a set $i_t \in \mathbb{N}^r$ (the subscript $t$ denotes the optimization step, since the selection is re-computed each step based on the current gradient or momentum). These indices identify which basis vectors are currently most relevant for this layer.

Step 4: Construct projection matrix. The projection matrix $Q_r$ is obtained by indexing the full DCT matrix:

$Q_r = Q[:, i_t] \in \mathbb{R}^{n \times r}$

This is a standard column-slicing operation — no matrix multiplication or factorization is needed. The columns of $Q_r$ are the $r$ DCT basis vectors that best align with the current gradient.

Step 5: Project the gradient. The low-rank gradient (or momentum) is computed as:

$g = G Q_r \in \mathbb{R}^{n \times r}$

This is a right-projection: it compresses the $n \times n$ gradient matrix to $n \times r$ by taking weighted combinations of the $r$ selected DCT basis vectors. The weights are the alignment scores from the similarity matrix.

Step 6: Compute reconstruction error. The projection back to the original space is:

$\hat{G} = g Q_r^\top = G Q_r Q_r^\top \in \mathbb{R}^{n \times n}$

The error is $\Delta = G - \hat{G}$. This error is not discarded — depending on the optimizer, it is either stored in the momentum buffer (Trion) or added to an error feedback accumulator (DCT-AdamW), ensuring that information orthogonal to the selected subspace is not permanently lost but carried forward to future steps.

Why this procedure is "dynamic." The selection of columns changes at every step because $G$ changes at every step. The paper emphasizes that the dynamic behavior comes from the $\binom{n}{r}$ possible sets of indices to select from (Section 2.1). As the gradient's principal directions shift during training, different DCT basis vectors become relevant, and the column selection tracks this automatically — no recomputation of the basis itself is needed.

Why $\ell_1$ or $\ell_2$ norm on columns of $S$ matters. The column norms of $S$ directly quantify how much of the gradient's Frobenius norm is captured by each basis vector. This is not an arbitrary heuristic — Section 4.1 proves that selecting the $r$ columns with the largest $\|q_j^\top G\|_2$ minimizes the Frobenius reconstruction error $\|G - Q_r Q_r^\top G\|_F$ exactly. The intuition: the total gradient energy is $\|G\|_F^2 = \sum_{j=1}^n \|q_j^\top G\|_2^2$ (because $Q$ is orthogonal, the basis is complete). To minimize the energy lost by keeping only $r$ columns, you should discard the $n-r$ columns with the smallest norms — which is equivalent to keeping the $r$ columns with the largest norms.

Memory implications. The only state that must be stored per layer is the set of $r$ integer indices $i_t$ — not a $d \times r$ floating-point matrix. The full DCT matrix $Q$ is stored once per GPU (not per layer), of size $d \times d$. For a model with $L$ layers, the storage comparison is:

• SVD/QR methods: $L \times d \times r$ floating-point numbers for projection matrices
• DCT method: $d \times d$ (once) $+ L \times r$ integers for indices

For large $L$ and modest $d$, this is a substantial saving. The paper reports that in the Llama-2 7B fine-tuning experiments (Appendix H, Table 7), the DCT projection saves 8 GB of memory (23.4%) at rank $r = 32$ compared to SVD-based projections.

The rule of thumb for which dimension to compress. The paper states (Section 2.1) that the standard practice in low-rank factorization for optimization is to "compress the smallest dimension of a matrix to $r$ dimensions." For a gradient matrix of size $\mathbb{R}^{n \times m}$ with $n \geq m$, the compression is from $\mathbb{R}^{n \times m}$ to $\mathbb{R}^{n \times r}$. The smallest dimension is usually $d_\text{model}$ — the hidden (embedding) size of the transformer. This means the DCT matrix should be of size $d_\text{model} \times d_\text{model}$, which is the same for all linear layers in a standard transformer (since all weight matrices involve the embedding dimension on at least one axis).

---

The Discrete Cosine Transform: Why This Specific Matrix

The DCT is introduced in Section 2.2 and motivated in Appendix C. It is a specific orthogonal matrix whose entries are cosines. The DCT-III matrix $Q \in \mathbb{R}^{n \times n}$ used in this work is defined entry-wise as:

$Q_{ij} = \sqrt{\frac{2}{n}} \cdot \cos\left(\frac{i(2j+1)\pi}{2n}\right)$

where $i, j \in \{0, 1, \dots, n-1\}$ (using 0-indexing for the formula). The first row ($i = 0$) must be divided by $\sqrt{2}$ to ensure orthogonality, i.e., $Q^\top Q = I_n$. The DCT-II matrix is the transpose of DCT-III.

Why DCT over alternatives. The paper considers three types of orthogonal matrices (Appendix C):

1. Random orthogonal matrix: Generated by orthogonalizing a random Gaussian matrix via QR decomposition (done once at training start). It would work — the dynamic column selection is agnostic to which orthogonal matrix is used — but it has no structure that enables fast computation. Computing $S = GQ$ would always cost $O(n^3)$ via standard matrix multiplication, with no possibility of acceleration.

2. Hadamard matrix: A specific Fourier-type matrix with entries $\pm 1$, used in model compression for its fast multiplication routines on GPUs (Walsh-Hadamard transform). However, Hadamard matrices exist only for dimensions that are powers of 2. For models whose embedding dimension $d_\text{model}$ is not a power of 2, the existing construction procedures produce matrices that are not perfectly orthogonal, making them unsuitable for this application where orthogonality (and thus perfect reconstruction of the energy identity) is required.

3. DCT matrix: A Fourier-type matrix with cosine entries. It has two critical advantages:

   • FFT-accelerable computation: The product $S = GQ$ can be computed in $O(n^2 \log n)$ time using Makhoul's $N$-point algorithm (Makhoul, 1980) instead of $O(n^3)$ for standard matmul. This is because the DCT is closely related to the Discrete Fourier Transform (DFT), which has a fast $O(n \log n)$ algorithm (FFT) per row. Applying it to all $n$ rows yields $O(n^2 \log n)$. The paper benchmarks this in Appendix D and shows speedups of up to $50\times$ for matrices where the number of rows is less than the number of columns, using float32.
   • Theoretical connection to gradient eigenbasis: Section 4.2 proves that the DCT approximates the left singular vectors of $G$. This is based on a matrix decomposition result: any square matrix can be factored into alternating diagonal and circulant matrices, and circulant matrices are diagonalized by the DFT matrix. For real symmetric $GG^\top$, the real part of the DFT (which is the DCT, up to minor variations) provides an approximate eigenbasis. This means DCT basis vectors are not arbitrary — they are structurally similar to the SVD's singular vectors, which is why selecting $r$ of them can approximate the optimal rank-$r$ SVD projection.

GPU materialization of the DCT matrix. Appendix A provides a vectorized construction on GPU. The procedure creates one vector $L = [0, 1, \dots, n-1]^\top$, replicates it into a matrix $\mathcal{I} \in \mathbb{N}^{n \times n}$ where each column is $L$, and computes:

$Q = \sqrt{\frac{2}{n}} \cos\left(\frac{\mathcal{I} \odot (2\mathcal{I}^\top + 1)}{2n} \pi\right)$

The computational efficiency comes from the element-wise product $\mathcal{I} \odot (2\mathcal{I}^\top + 1)$, which computes all integer entries $i(2j+1)$ in a single GPU-parallel operation without loops. The first row is then divided by $\sqrt{2}$ to ensure orthogonality. This construction is done once at training start and the resulting matrix is stored in GPU memory.

The FFT acceleration: Makhoul's algorithm. Appendix D describes Makhoul's $N$-point algorithm for computing a fast 1D DCT-II, applied row-by-row to $G$. The procedure for each row is:

1. Permute the input signal (row of $G$) so that odd-indexed elements appear in increasing order, even-indexed in decreasing order, and they are interleaved: $[a,b,c,d,e,f] \to [a,c,e,f,d,b]$. This permutation can be pre-computed and cached for a given row length.
2. Compute FFT of the permuted signal using standard complex FFT.
3. Multiply by Fourier coefficients $W_k = \exp(-2i\pi k / N)$, which can also be pre-computed and cached.
4. Take the real part of the product: $\text{DCT}(\text{row}) = \text{Real}(\text{FFT}(\text{permuted row}) \times W)$.

The per-row cost is $O(n \log n)$ (dominated by the FFT), and doing this for all $n$ rows costs $O(n^2 \log n)$. The key insight is that this sequence of operations — permutation, FFT, complex multiplication, real-part extraction — is mathematically equivalent to multiplying by the DCT-II matrix, but the DCT matrix embeds all these operations into its entries, which costs $O(n^3)$ to apply directly.

Practical limitations of the FFT acceleration. The paper is transparent about a critical practical constraint (Appendix D): Makhoul's algorithm currently runs only on float32 inputs in PyTorch because PyTorch lacks a complex-bfloat16 data type at the time of writing (where both real and imaginary parts are stored in bfloat16). Complex half-precision is only supported for input sizes that are powers of 2. Since training typically uses bfloat16 for gradients, the FFT acceleration cannot be directly applied without type conversions.

The benchmarking results in Table 5 (Appendix D) quantify the tradeoff:

• For matrices where $R \geq C$ (number of rows $\geq$ number of columns), the standard bfloat16 matmul $S = GQ$ is consistently faster than the float32 Makhoul algorithm because bfloat16 has higher GPU throughput.
• For matrices where $R < C$, Makhoul is still $3.5\times$ faster than bfloat16 matmul, even with the datatype disadvantage.
• The larger speedups ($8\times$ to $50\times$) are observed only when comparing float32 Makhoul against float32 matmul (Table 4), not against the faster bfloat16 matmul that would be used in practice.

The paper notes that the float32 gradient accumulator buffer (used in mixed-precision training to accumulate bfloat16 gradients) provides access to float32 data in the optimizer step, offering a potential path to use Makhoul's algorithm without explicit conversion overhead. However, this requires the optimizer to access the float32 accumulator, which may not be exposed in all frameworks.

Why the DCT matrix is the same shape for all layers. In a standard transformer, every linear layer has one dimension equal to $d_\text{model}$ (the embedding/hidden size). For example, the query projection weight is $d_\text{model} \times d_\text{head}$, and the output projection is $d_\text{model} \times d_\text{model}$. The smaller of the two dimensions is typically $d_\text{model}$, so the DCT matrix is sized $d_\text{model} \times d_\text{model}$. For layers where the gradient matrix has shape $\mathbb{R}^{n \times m}$ with $m < n$ and $m = d_\text{model}$, the right-projection $g = G Q_r$ with $Q_r \in \mathbb{R}^{m \times r}$ is used (compressing the columns). For layers where $n = d_\text{model}$ is the smaller dimension, a left-projection $g = Q_r^\top G$ would be used. The paper primarily describes the right-projection case.

---

Trion: DCT-Based Improvement to Dion

Trion (Section 2.3, Algorithm 1) is the first of two proposed optimizers. It replaces the Power-Iteration and QR decomposition in Dion with DCT-based dynamic column selection, while preserving the rest of Dion's structure: momentum accumulation, Newton-Schulz orthogonalization of the low-rank momentum, and error feedback into the momentum buffer.

What Dion does, which Trion replaces. To understand Trion, we must understand what it replaces. Dion (Ahn et al., 2025) is a low-rank variant of Muon. At each step, Dion:

1. Accumulates gradient into a momentum buffer $B_t$ (full size, $R \times C$).
2. Uses Power-Iteration to find a low-rank approximation $B_t \approx P_t R_t^\top$ where $P_t \in \mathbb{R}^{R \times r}$ and $R_t \in \mathbb{R}^{C \times r}$.
3. Orthogonalizes $P_t$ and $R_t$ separately via QR decomposition (which costs $O(R r^2)$ and $O(C r^2)$, scaling with $r$).
4. Computes the orthogonal update $O_t = P_t R_t^\top$ and applies it.
5. Stores the projection error back into the momentum buffer for the next step.

The two bottlenecks are Power-Iteration (iterative, requires multiple passes) and QR decomposition (runtime depends on $r$). Trion replaces steps 2 and 3 with DCT-based column selection.

Trion step-by-step walkthrough (Algorithm 1). The algorithm is presented for a layer where the DCT matrix $D_C \in \mathbb{R}^{C \times C}$ compresses along the $C$ dimension (the smaller dimension). Here is each operation:

Line 3 — Gradient computation: $G_t = \nabla_\theta L(\theta_t) \in \mathbb{R}^{R \times C}$ is the standard backpropagated gradient for this layer's weight matrix.

Line 4 — Momentum accumulation: $B_t = M_{t-1} + G_t \in \mathbb{R}^{R \times C}$. This is Muon/Dion-style momentum: the accumulator simply sums gradients without a decay factor at this stage (unlike Adam's exponential moving average). The momentum $B_t$ is what gets factorized, not the raw gradient. This is important because momentum is smoother and has more stable principal directions than raw gradients, making the low-rank approximation more reliable.

Line 5 — Similarity computation: $S_t = \text{Makhoul}(B_t) \in \mathbb{R}^{R \times C}$, or alternatively $S_t = B_t \cdot D_C$. This computes the DCT of each row of the momentum matrix. The choice between Makhoul (FFT-based) and matmul depends on the datatype: matmul in bfloat16 is used in practice for layers where $R \geq C$ due to the PyTorch complex-bfloat16 limitation, while Makhoul in float32 is preferred for $R < C$ where it is faster (per the benchmarks in Appendix D). The result $S_t$ is the similarity matrix — its columns are the alignments between the momentum's rows and the DCT basis vectors.

Line 6 — Dynamic column selection: $i_t = \text{DynamicColumnSelection}(S, r) \in \mathbb{N}^r$. This implements the procedure from Section 2.1: compute the $\ell_1$ or $\ell_2$ norm of each column of $S_t$, sort the columns by norm in descending order, and return the indices of the top $r$ columns. The notation $\mathbb{N}^r$ indicates a vector of $r$ integer indices.

Line 7 — Projection matrix extraction: $Q_t = D_C[:, i_t] \in \mathbb{R}^{C \times r}$. This is a simple column-indexing operation: take the $r$ columns of the DCT matrix identified by $i_t$. No computation beyond memory indexing is required. This $Q_t$ is the projection matrix that will be used both for extracting the low-rank momentum and for projecting the update back to the original space.

Line 8 — Low-rank momentum extraction: $b_t = S_t[:, i_t] \in \mathbb{R}^{R \times r}$. Critically, Trion extracts $b_t$ from the similarity matrix $S_t$, not by re-multiplying $B_t$ with $Q_t$. Since $S_t = B_t D_C$, indexing the same columns $i_t$ from $S_t$ is equivalent to computing $B_t \cdot Q_t$ — but it avoids a second matrix multiplication. The result $b_t$ is the low-rank representation of the momentum in the chosen DCT subspace.

Line 9 — Projection error computation: $\Delta_t = B_t - b_t Q_t^\top$. This is the residual: the component of the momentum that lies outside the selected $r$-dimensional subspace. The reconstruction $b_t Q_t^\top$ is the best approximation of $B_t$ using only those $r$ DCT basis vectors (by the optimality proof in Section 4.1). The error $\Delta_t$ contains momentum information in the orthogonal complement of the selected subspace.

Line 10 — Momentum update with error feedback: $M_t = \mu B_t + (1-\mu) \Delta_t = B_t - (1-\mu) b_t Q_t^\top$. This is a subtle and important step. The momentum for the next iteration is not simply $B_t$, and the error is not simply discarded. Instead, the momentum is updated as a convex combination of the full momentum $B_t$ and the projection error $\Delta_t$, controlled by a decay parameter $\mu$ (the Muon/Dion momentum parameter, typically set to values like 0.95). Expanding the algebra:

$M_t = \mu B_t + (1-\mu)(B_t - b_t Q_t^\top) = B_t - (1-\mu)b_t Q_t^\top$

This means: the portion of the momentum captured by the low-rank projection is decayed by a factor of $(1-\mu)$, while the projection error is retained in full. Why? Because the low-rank components will be orthogonalized and applied as an update (thus they are "used up"), while the residual should persist in the momentum buffer to influence future steps. This is analogous to the error feedback mechanism in Dion, where $P_t R_t^\top$ is the component that gets orthogonalized and applied, and the residual is saved. Trion replaces $P_t R_t^\top$ with $b_t Q_t^\top$.

Line 11 — Newton-Schulz orthogonalization on low-rank momentum: $o_t = \text{NewtonSchulz}(b_t) \in \mathbb{R}^{R \times r}$. This is the key computational savings. In Muon, Newton-Schulz runs on the full $R \times C$ momentum matrix, computing odd powers up to 5th order, requiring full-size matrix multiplications. In Dion, Newton-Schulz runs separately on $P_t$ ($R \times r$) and $R_t$ ($C \times r$) after QR decomposition. In Trion, Newton-Schulz runs only on the $R \times r$ low-rank matrix $b_t$. The Newton-Schulz iteration approximates $UV^\top$ from the SVD of its input — applied to $b_t$, it orthogonalizes the low-rank momentum, pushing its singular values toward 1 while preserving its directional structure. Because $r$ is small (128–512 in the experiments), this is dramatically cheaper than full-size Newton-Schulz. The paper notes that the efficient Triton kernels from the official Dion repository can be used here for further acceleration.

Line 12 — Project back to original space: $O_t = o_t Q_t^\top \in \mathbb{R}^{R \times C}$. The orthogonalized low-rank update $o_t$ (which lives in the selected DCT subspace) is projected back to the full parameter space by multiplying with $Q_t^\top$. This produces the final update matrix $O_t$ that will be subtracted from the weights.

Line 13 — Parameter update: $\theta_{t+1} = (1 - \lambda \eta_t) \theta_t - \eta_t \max(1, \sqrt{R/C}) O_t$. This is the standard Muon/Dion update rule, with decoupled weight decay (the $(1 - \lambda \eta_t)$ term) and the $\max(1, \sqrt{R/C})$ scaling factor that accounts for the aspect ratio of the weight matrix.

What makes Trion fundamentally different from Dion. Three properties distinguish the approach:

1. Rank-independent factorization cost: Computing $S_t = B_t D_C$ (line 5) has the same cost regardless of $r$, because it always computes the full similarity matrix. The column selection (line 6) is just a sort over $C$ elements and a top-$r$ extraction, which is $O(C \log C)$ — negligible compared to matrix multiplication. QR decomposition in Dion costs $O(R r^2 + C r^2)$, which grows with $r$. Table 1 confirms this empirically: Trion's runtime is approximately constant across $r = 128, 256, 512$ for a given model size, while Dion's runtime increases by 8–18% from $r = 128$ to $r = 512$.

2. No iterative refinement: Power-Iteration in Dion requires multiple passes over the momentum matrix (typically 5–10 iterations) to converge to the top singular vectors. DCT column selection is a single-shot procedure: one matrix multiplication to get $S_t$, one sorting step, done.

3. More accurate low-rank approximation (empirically): Figure 2 shows the Frobenius norm of the projection error for Dion ($\|B_t - P_t Q_t^\top\|_2$) and Trion ($\|B_t - O_t\|_2$) across several linear layers in a Llama-30M model. Trion consistently achieves lower projection error, and the error shows a decreasing trend for some layers over training steps, indicating that the DCT subspace adapts effectively to the evolving momentum structure. For Dion, the projection error is approximately constant — Power-Iteration captures the same quality of approximation at each step, while Trion's dynamic selection improves as the momentum becomes more structured during training.

Communication in distributed training with Trion. The paper develops Trion on top of the published Muon/Dion codebase that leverages ZeRO-style optimizations for Distributed Data Parallel (DDP) settings (Section 2.3). The key design:

• The DCT matrix $D_C$ is replicated on every GPU (computed once at initialization and broadcast).
• Only one GPU computes the update $O_t$ for a given layer (the "source GPU").
• Instead of communicating the full $R \times C$ matrix $O_t$, the source GPU communicates only the low-rank $o_t \in \mathbb{R}^{R \times r}$ (via all-gather or similar primitives).
• Each receiving GPU locally computes $O_t = o_t Q_t^\top$ using its copy of $D_C$ and the column indices $i_t$.
• This reduces communication volume from $O(RC)$ to $O(Rr)$, which is a substantial saving when $r \ll C$.

For Fully Sharded Data Parallel (FSDP) settings, the paper notes additional complexity: each layer must be sharded along the appropriate dimension to avoid materializing full tensors when computing similarities. The FSDP implementation requires deciding whether a layer needs a left- or right-projection based on which dimension ($R$ or $C$) is being sharded.

---

DCT-AdamW: DCT-Based Low-Rank AdamW

DCT-AdamW (Section 2.4, Algorithm 2 in Appendix E) is the second proposed optimizer. It replaces the SVD or Block Power-Iteration used in LDAdam, FRUGAL, and FIRA with DCT-based dynamic column selection. Unlike Trion, which is built on Muon/Dion's momentum structure, DCT-AdamW follows the AdamW paradigm with exponential moving averages of gradients and squared gradients, plus decoupled weight decay.

What DCT-AdamW replaces. The target methods are:

• LDAdam (Robert et al., 2025): Uses Block Power-Iteration at every step to compute projection matrices, stores two consecutive projection matrices ($Q_\text{prev}$ and $Q_\text{crt}$, each $d \times r$) per layer for subspace rotation, and applies error feedback.
• FRUGAL (Zmushko et al., 2024): Uses SVD every ~200 steps, feeds the projection error to SignSGD.
• FIRA (Chen et al., 2024): Uses SVD every ~200 steps, scales the projection error norm-based.

DCT-AdamW replaces the SVD/Power-Iteration step with DCT column selection, and replaces the storage of two $d \times r$ projection matrices with two sets of $r$ integer indices. The core AdamW logic in the compressed space remains the same.

The subspace rotation problem. A key challenge in any low-rank optimizer that changes subspaces between steps is: how do you update momentum buffers when the coordinate system changes? ADAM's momentum $m_t$ is a weighted sum of past gradients. If the projection matrix changes from $Q_\text{prev}$ to $Q_\text{crt}$, the past momentum was accumulated in the $Q_\text{prev}$ subspace, but the new gradient $g_t$ is in the $Q_\text{crt}$ subspace. You cannot simply add them. LDAdam solves this by storing $m_{t-1}$ in the full space (by keeping the projection matrices to map back), rotating it to the new subspace, and then adding the new projected gradient. This requires storing $Q_\text{prev}$.

DCT-AdamW solves this more efficiently. The rotation matrix $R \in \mathbb{R}^{r \times r}$ is defined as:

$R = Q_\text{prev}^\top \cdot Q_\text{crt}$

where $Q_\text{prev}$ and $Q_\text{crt}$ are the $d \times r$ projection matrices formed by indexing the DCT matrix with the previous and current column index sets, respectively.

What $R$ computes: It is the $r \times r$ matrix of inner products between the previous basis vectors and the current basis vectors. Specifically, $R_{ij} = (q_{\text{prev},i})^\top q_{\text{crt},j}$ — the cosine similarity between the $i$-th previously selected DCT column and the $j$-th currently selected DCT column.

Why this works: The momentum $m_{t-1}$ was accumulated in the previous subspace, so it is stored as a low-rank vector $m_{t-1} \in \mathbb{R}^{n \times r}$ (for left-projection; the paper's Algorithm 2 shows right-projection but the principle is symmetric). To incorporate it into the current step, we want to express $m_{t-1}$ in the current basis. In the full space, the momentum would be $m_{t-1}^\text{full} = m_{t-1} Q_\text{prev}^\top \in \mathbb{R}^{n \times d}$. Projecting this into the current subspace: $m_{t-1}^\text{rotated} = m_{t-1}^\text{full} Q_\text{crt} = m_{t-1} Q_\text{prev}^\top Q_\text{crt} = m_{t-1} R$. So the rotation can be done entirely in the $r$-dimensional space by multiplying the stored low-rank momentum by the $r \times r$ matrix $R$, without ever materializing the full $d$-dimensional representation. This is what Algorithm 2 does in Line 11: $m_t = \beta_1 \cdot m_{t-1} \cdot R + (1 - \beta_1) g_t$.

The absolute value in the second-moment rotation. Line 12 of Algorithm 2 uses $v_t = \beta_2 |v_{t-1} \cdot R| + (1-\beta_2) g_t^2$. The absolute value is necessary because $R$ can have negative entries (when the new basis vectors point in opposite directions from the old ones, creating negative dot products). The second moment $v_t$ must be non-negative (it represents squared gradient magnitudes), and multiplying by a matrix with negative entries could violate this. Taking the absolute value after rotation ensures positivity. This is a heuristic fix acknowledged in the pseudocode comment: "Rotating $v_t$ might introduce negative values and we apply the absolute value function to force the non-negativity."

The UpdateSubspace procedure (Algorithm 3). This procedure is called at each step (or periodically, depending on $T_u$) to decide whether to recompute the column selection and produce the rotation matrix $R$. It works as follows:

• Line 1: Initialize $R = I_{r \times r}$ (identity) — if the subspace doesn't change, the rotation is identity and the momentum buffers are just carried forward unchanged.
• Lines 3–5: If $t > 1$, save the current indices as previous: $\mathcal{I}_\text{prev} \leftarrow \mathcal{I}_\text{crt}$.
• Lines 6–10: If $t = 1$ (first step) or $t \bmod T_u = 0$ (subspace update interval reached):
  • Compute similarities $S = \text{Makhoul}(G)$ or $S = GQ$.
  • Select new column indices: $\mathcal{I}_\text{crt} \leftarrow \text{RankCols}(GQ, r)$.
  • Compute rotation matrix: $R \leftarrow Q_\text{prev}^\top \cdot Q_\text{crt}$, where $Q_\text{prev} = Q[:, \mathcal{I}_\text{prev}]$ and $Q_\text{crt} = Q[:, \mathcal{I}_\text{crt}]$ are the $d \times r$ submatrices formed by indexing the DCT matrix with the old and new index sets.
• Line 11: Return $R$ for use in the momentum update.

The subspace update interval $T_u$ controls the frequency of recomputing the column selection. The paper notes (Table 3) that GaLore uses $T_u = 200$ (to amortize the cost of SVD), while LDAdam uses $T_u = 1$ (every step, enabled by the cheaper Block Power-Iteration). DCT-AdamW supports any $T_u$ because the cost of column selection is low regardless.

Error feedback in DCT-AdamW. Algorithm 2 includes error feedback with optional 8-bit quantization. The error feedback buffer $\Xi_t$ accumulates the projection error:

• Line 7: The gradient is augmented with the error feedback from the previous step: $G_t \leftarrow \nabla_\theta f(\theta_t) + \Xi_t$.
• Line 9: The augmented gradient is projected: $g_t \leftarrow G_t \cdot Q_\text{crt}$.
• Line 10: The new projection error is computed: $\Xi_t \leftarrow G_t - g_t \cdot Q_\text{crt}^\top$. This error — the component of $G_t$ that lies outside the current subspace — is saved and will be added to the next step's gradient, ensuring it eventually gets incorporated (potentially in a different subspace where it aligns better).

The paper notes (Section 2.4) that the error feedback can be quantized to 8 bits without degrading performance, but 4-bit quantization was attempted and found insufficient: "the lowest resolution we can quantize EF to is 8-bits without degrading the optimizer performance."

Key design differences from LDAdam. The paper emphasizes several simplifications relative to LDAdam:

1. Cheaper subspace rotation storage: DCT-AdamW stores two sets of $r$ integers ($\mathcal{I}_\text{prev}$ and $\mathcal{I}_\text{crt}$) instead of two $d \times r$ floating-point matrices ($Q_\text{prev}$ and $Q_\text{crt}$). The rotation matrix $R = Q_\text{prev}^\top Q_\text{crt}$ is computed on-the-fly from the DCT matrix and the index sets, rather than being stored.

2. Simpler second-moment update: The paper states that DCT-AdamW updates $v_t$ "using a simpler rule compared to LDAdamW" (Appendix E, preamble), though the specific simplification is not detailed beyond the absolute value application described above.

3. No Block Power-Iteration: The factorization step is a single matrix multiplication ($S = GQ$) plus a sort, rather than iterative Block Power-Iteration with QR decompositions.

---

Integration with FRUGAL and FIRA

The paper also demonstrates that the DCT projection is a drop-in replacement for SVD in FRUGAL and FIRA (Appendix G). The integration is straightforward: wherever the original optimizer computes an SVD of the gradient to obtain projection matrices, the DCT procedure is substituted.

For FRUGAL, the original optimizer:

• Computes SVD of the gradient every 200 steps to get $Q_r$.
• Projects the gradient to $g = G Q_r$.
• Runs AdamW on the projected gradient $g$.
• Feeds the projection error $\Delta = G - g Q_r^\top$ to a SignSGD optimizer (stateless, unlike AdamW).

The DCT variant replaces the SVD step with: $Q_r$ = DCT columns selected by ranking $S = GQ$, same projection $g = G Q_r$, same error handling. The results in Appendix G, Table 6, show that on Llama-800M pretraining with 16B tokens, DCT-FRUGAL achieves train perplexity 18.70 vs. SVD-FRUGAL's 18.18 (a degradation of 0.52, or ~2.8%), but reduces runtime by 1h 48m (22.6%) and memory by 2.2 GB (3.5%). Compared to the random baselines in FRUGAL (RandPerm and Random projections), DCT achieves ~1 point lower perplexity, validating that the dynamic column selection provides meaningful adaptivity beyond what a random projection offers.

For FIRA, the original optimizer similarly uses SVD every 200 steps but handles the projection error by scaling it appropriately to recover full-rank performance. The DCT variant replaces SVD with DCT column selection. Table 6 shows that DCT-FIRA actually outperforms SVD-FIRA: train perplexity 17.35 vs. 17.59, with 1h 54m (23.8%) runtime reduction and 2 GB (3%) memory saving. This suggests that for FIRA's error-scaling approach, the DCT subspace is not just an adequate approximation but potentially beneficial.

---

Theoretical Guarantees: Why Dynamic Column Selection Is Optimal

Section 4 provides the mathematical justification for two key claims: (1) the norm-based column ranking is optimal for minimizing reconstruction error given any fixed orthogonal $Q$, and (2) DCT specifically is a good choice of $Q$ because it approximates the gradient's eigenbasis.

Optimality of norm-based ranking (Section 4.1). Consider $G \in \mathbb{R}^{n \times m}$ and an orthogonal matrix $Q \in \mathbb{R}^{n \times n}$ (the argument uses left multiplication; right multiplication is symmetric). Let $Q_r$ be an $n \times r$ submatrix of $Q$ with columns $q_1, \dots, q_r$. The reconstruction error in Frobenius norm is:

$\|G - Q_r Q_r^\top G\|_F^2$

The paper derives an identity showing this equals:

$\|G\|_F^2 - \sum_{i=1}^r \|q_i^\top G\|_2^2$

Why this identity holds. The derivation uses: $Q_r Q_r^\top$ is an orthogonal projection onto the span of the selected columns (since $Q_r^\top Q_r = I_r$), so $(I - Q_r Q_r^\top)$ is the complementary projection. By the cyclic property of trace and the orthogonality of $Q$, the squared Frobenius norm of the projection of $G$ onto the selected columns is $\|Q_r^\top G\|_F^2 = \sum_{i=1}^r \|q_i^\top G\|_2^2$ (the sum of squared Euclidean norms of each basis vector's alignment with $G$). The reconstruction error is the total gradient energy minus the energy captured by the $r$ selected basis vectors.

The optimality conclusion. To minimize the reconstruction error, we should maximize $\sum_{i=1}^r \|q_i^\top G\|_2^2$ — the sum of alignment energies of the selected columns. Since the total energy across all $n$ columns is fixed (equal to $\|G\|_F^2$, because $Q$ is a complete orthogonal basis), the optimal strategy is to select the $r$ columns with the largest individual alignment energies $\|q_i^\top G\|_2^2$. This proves that ranking by $\ell_2$ norm of columns of $S = Q^\top G$ (or $S = GQ$ for right projection) and selecting the top $r$ is optimal.

Contractivity of the compression. The paper further derives a worst-case bound:

$\|G - Q_r Q_r^\top G\|_F^2 \leq \left(1 - \frac{r}{n}\right) \|G\|_F^2$

This follows because the average alignment energy per column is $\|G\|_F^2 / n$, and the top $r$ columns must have at least this average energy (by the pigeonhole principle), so the captured energy is at least $(r/n) \|G\|_F^2$, and the error is at most $(1 - r/n) \|G\|_F^2$. The compression is contractive with factor $1 - r/n$. This contractivity property is highlighted as "the key property in the convergence analysis of compressed optimization" (citing Stich et al., 2018; Richtárik et al., 2021; Li et al., 2022; Modoranu et al., 2024; Robert et al., 2025), establishing that DCT-based compression inherits the same theoretical convergence guarantees as the SVD-based methods it replaces.

Extension to other $p$-norms. The paper generalizes the argument to any $p$-norm, showing:

$\|G - Q_r Q_r^\top G\|_p \leq \max(1, n^{\frac{1}{p} - \frac{1}{2}}) \sum_{i=r+1}^n \|q_i^\top G\|_p$

The derivation uses the triangle inequality of $p$-norms and the identity $\|u v^\top\|_p = \|u\|_p \|v\|_p$ for rank-one matrices. The bound relates the reconstruction error to the sum of the $p$-norms of the discarded basis vectors' alignments — again, minimized by discarding columns with the smallest alignment norms, which is equivalent to keeping columns with the largest alignment norms.

DCT as a linear approximation of the gradient eigenbasis (Section 4.2). This subsection provides theoretical motivation for why DCT specifically, beyond the computational advantages. The argument uses a matrix decomposition result from the optical information processing literature (Müller-Quade et al., 1998; Schmid et al., 2000):

Any square matrix $M \in \mathbb{C}^{n \times n}$ can be decomposed into a product of diagonal and circulant matrices: $M = D_1 C_2 D_3 \dots D_{2k-3} C_{2k-2} D_{2k-1}$, where $D$'s are diagonal and $C$'s are circulant. Circulant matrices have the form:

$C = \begin{bmatrix} c_0 & c_1 & c_2 & \cdots & c_{n-1} \\ c_{n-1} & c_0 & c_1 & \cdots & c_{n-2} \\ \vdots & \vdots & \cdots & \vdots & \vdots \\ c_1 & c_2 & \cdots & c_{n-1} & c_0 \end{bmatrix}$

where each row is a cyclic right shift of the previous row. Circulant matrices are diagonalized by the Discrete Fourier Transform (DFT) matrix $F$: $C = F^* D F$, where $F$ has entries $F_{ij} = \frac{1}{\sqrt{n}} w^{ij}$ with $w = e^{2\pi i / n}$.

Applying this to the gradient. Consider the symmetric matrix $M = G G^\top \in \mathbb{R}^{n \times n}$. Its eigenvectors are the left singular vectors $U$ from the SVD of $G$. Applying the decomposition theorem and substituting $C = F^* D F$ for each circulant factor, $G G^\top$ can be expressed as a product alternating between $F D F^*$ and diagonal matrices. The "linear approximation" means keeping only one such factor:

$G G^\top \approx F D_1 F^*$

This implies that the eigenvectors $U$ are approximated by the DFT matrix $F$. However, $G G^\top$ is real and symmetric, so its eigenvalues are real, and its eigenvectors are real. The real part of $F$, $\text{Re}(F)$, also forms an approximate eigenbasis that better aligns with $U$ (since $U$ is real). The real part of the DFT matrix is — up to minor variations — the Discrete Cosine Transform (DCT).

What this means practically. The DCT basis vectors are not arbitrary. They are structurally similar to the true principal components (left singular vectors) of the gradient matrix. When we select the $r$ DCT columns that best align with the current gradient, we are effectively selecting the first $r$ terms in an approximate eigendecomposition — analogous to what truncated SVD does, but without computing the eigendecomposition. This explains why the method can achieve projection errors competitive with or lower than Power-Iteration (Figure 2): the DCT basis is a good enough approximation to the eigenbasis that a simple top-$r$ selection recovers most of the energy that SVD would capture.

The analogy to Taylor expansion. The paper draws a conceptual analogy: "just as loss functions are linearly approximated at each iteration in first-order optimization algorithms, we consider a 'linear' approximation of the decomposition" by keeping only one diagonal-circulant-diagonal factor. This is not a formal bound — it is a theoretical motivation, not a proof of approximation quality. The number of factors $2k-1$ in the exact decomposition can be up to $2n-1$ for almost all matrices (in the sense of Lebesgue measure), and it is conjectured that up to $n$ factors is sufficient (Huhtanen & Perämäki, 2015). The "linear approximation" drops all but one factor, which is analogous to using only the first-order term in a series expansion.

---

Summary of Design Choices and Their Justifications

• Fixed DCT matrix over per-layer computed matrices: Eliminates per-layer factorization cost entirely. The DCT matrix is pre-computed once per GPU at training start, shared across all layers. The only per-layer computation is a matrix multiplication ($S = GQ$) and a norm-based sort, both of which are independent of the chosen rank $r$. This decouples runtime from rank, which is the paper's primary claimed advantage.

• Dynamic column selection over using all DCT columns: Each layer selects a different set of $r$ columns based on its current gradient/momentum structure. This adaptivity is essential — a fixed set of columns (e.g., always the first $r$) would be a static low-rank approximation that cannot track changing gradient statistics during training. The optimality proof in Section 4.1 guarantees that the norm-based selection minimizes reconstruction error among all possible choices of $r$ columns from $Q$.

• $\ell_1$ or $\ell_2$ norm over other ranking criteria: These norms directly measure the energy captured by each basis vector. The optimality proof specifically uses $\ell_2$ norm (squared), but the paper notes $\ell_1$ works similarly in practice. The choice between them is an implementation detail — both produce similar rankings because larger $\ell_2$ norms generally imply larger $\ell_1$ norms (though not strictly, especially for sparse alignments).

• Right-projection over left-projection: The paper describes the projection as $g = G Q_r$ where $Q_r$ compresses the columns (right multiplication). This follows the convention of compressing along the smaller matrix dimension ($d_\text{model}$). For layers where the gradient has shape $R \times C$ with $C = d_\text{model}$ as the smaller dimension, right-projection compresses from $R \times C$ to $R \times r$. For layers where $R = d_\text{model}$ is smaller, left-projection $g = Q_r^\top G$ would be used. The paper's Algorithm 1 assumes the right-projection case.

• Trion's momentum error feedback ($M_t = B_t - (1-\mu)b_t Q_t^\top$) over simple residual storage: This ensures that the low-rank components are "used up" after orthogonalization (their contribution in the next step's momentum is decayed), while the residual (the part of the momentum not captured by the current subspace) persists. Without this, the same low-rank directions would accumulate indefinitely in the momentum buffer, and the orthogonal complement would be lost. This design is inherited from Dion and Muon.

• DCT-AdamW's subspace rotation via $R = Q_\text{prev}^\top Q_\text{crt}$ over storing full projection matrices: This is the key memory-saving design. Instead of storing two $d \times r$ matrices to project momentum back to full space and then forward to the new subspace, the rotation is computed directly in the $r$-dimensional space using only the $r$ integer indices and the shared DCT matrix. The $r \times r$ rotation matrix $R$ is computed on-the-fly and discarded after use.

• 8-bit quantization of error feedback over no quantization or 4-bit quantization: The paper attempted 4-bit quantization of the error feedback buffer and found it degraded performance. 8-bit quantization preserves accuracy while reducing the memory footprint of the error buffer. This is a practical compromise: error feedback accumulates gradient residuals that can be small in magnitude, and aggressive quantization risks losing the signal entirely.

• FFT acceleration (Makhoul's algorithm) when possible, matmul otherwise: The paper provides both options (Algorithm 1, Line 5: "Makhoul($B_t$) or $S_t = B_t \cdot D_C$") because the FFT path is only faster under specific conditions (float32 data, $R < C$ matrix shapes). In practice, with bfloat16 gradients and $R \geq C$ layers (common in transformers where the embedding dimension is the smaller side), the standard matmul in bfloat16 may be faster despite the asymptotic $O(n^3)$ vs. $O(n^2 \log n)$ difference. The paper is agnostic: use whichever is faster for the current layer shape and datatype.

4. Key Insights and Innovations

Innovation 1: The Fixed Basis Dynamic Selection Paradigm — Decoupling Projection Matrix Quality from Per-Layer Factorization

The paper's most fundamental intellectual contribution is not a specific optimizer or a particular orthogonal matrix, but a reframing of what the projection step in low-rank optimization is supposed to accomplish. Prior work across the entire lineage — GaLore, LDAdam, FRUGAL, FIRA, Dion — implicitly assumed that the projection matrix must be derived from the data itself: if you want to compress a gradient, you compute its SVD (or approximate it via Power-Iteration or Block Power-Iteration) and use the resulting singular vectors. This assumption is so natural that it went unquestioned: of course the best low-rank approximation of a matrix uses its own principal components. The Eckart-Young theorem says SVD is optimal. Why would you do anything else?

The paper's reframing is to ask a different question: what if we separate the quality of the basis from its selection? Instead of computing the basis from scratch for each layer at each step, we can use a single, fixed, universal basis and let each layer select which basis vectors are relevant for its current state. The basis doesn't need to be optimal in the SVD sense — it just needs to be rich enough that a smart selection of $r$ vectors from it can capture most of the energy. The optimality guarantee shifts from "this basis is the best possible basis" to "this selection is the best possible selection from this basis" — which is a weaker guarantee mathematically, but one that Section 4.1 proves is still achievable via a simple norm-based ranking.

This reframing is significant because it decouples two things that were previously entangled: the cost of finding a good basis (expensive, scales with rank, must be done per layer) and the quality of the resulting low-rank approximation (should be as good as possible). By fixing the basis in advance, the per-step cost becomes independent of rank — computing $S = GQ$ always produces the full similarity matrix regardless of $r$, and the column selection is just a sort. The basis can be chosen once based on global properties (like being an approximate eigenbasis for gradients in general, per the Section 4.2 argument), and the adaptivity is pushed entirely into the lightweight selection step.

This is not an incremental improvement to an existing factorization method — it is a conceptual break from the data-derived projection paradigm. The evidence that this break is valid comes from Figure 2 and Table 1: the DCT-based projection achieves lower reconstruction error than Dion's Power-Iteration (which attempts to compute the actual top singular vectors), and Table 1 shows that Trion's runtime is approximately flat across ranks while Dion's grows. If the data-derived basis were inherently superior, one would expect the approximate SVD (Power-Iteration) to yield lower projection error than a fixed basis with selection. The fact that it doesn't — that a fixed DCT basis with smart selection can outperform iterative SVD approximation — is the paper's central empirical validation of its reframing.

Innovation 2: The DCT as a Principled, Not Arbitrary, Choice of Universal Basis

If you accept the fixed-basis-dynamic-selection paradigm, the immediate next question is: which basis? The paper could have chosen a random orthogonal matrix, or a Hadamard matrix, or any other orthogonal set. What elevates the DCT choice from an implementation detail to a conceptual contribution is the theoretical argument in Section 4.2 that DCT basis vectors structurally approximate the SVD's left singular vectors. This is not a vague claim about Fourier bases being good for compression — it is a specific argument rooted in a matrix decomposition theorem from optical information processing (Müller-Quade et al., 1998; Schmid et al., 2000) that shows any matrix can be factored into alternating diagonal and circulant matrices, where circulant matrices are diagonalized by the DFT. For real symmetric $GG^\top$, the real part of the DFT — the DCT — is a natural approximate eigenbasis.

The significance of this argument is that it transforms the DCT from a heuristic convenience into a theoretically motivated choice. Prior work had used structured orthogonal matrices (Hadamard, random Fourier features) for model compression and quantization (e.g., the QuIP line of work), but always with the justification that they are "close enough" to a random rotation and have fast multiplication routines. This paper's argument is different: it claims the DCT is not just any orthogonal matrix, but one that is systematically related to the gradient's own eigenstructure via the circulant-diagonal decomposition. This provides an explanation, beyond "it works empirically," for why the DCT projection can match or exceed SVD-based methods in reconstruction quality (Figure 2): it is approximating the same singular vectors, just through a different computational path.

The practical consequence is that the DCT is not interchangeable with a random orthogonal matrix. The paper tests this implicitly in Appendix G: when comparing DCT projection against the RandPerm and Random projections in FRUGAL, DCT achieves approximately 1 point lower perplexity. This gap is not huge, but it is consistent, and it supports the claim that DCT's structure — its relationship to the gradient eigenbasis — provides a meaningful advantage over a purely random orthogonal basis. If the DCT were just "any orthogonal matrix," Random would perform identically.

This contribution is fundamental rather than incremental: it gives practitioners a principled reason to choose DCT over alternatives, grounded in matrix decomposition theory rather than empirical trial-and-error.

Innovation 3: Verifier-Free, Rank-Independent Newton-Schulz Orthogonalization

The paper's Trion optimizer achieves something that prior Muon-family optimizers could not: running Newton-Schulz iteration on a low-rank input without first computing that low-rank representation via an expensive, rank-dependent factorization. This is an architectural innovation that changes what is possible in practice for orthogonalized momentum methods.

To understand why this matters, consider the Muon/Dion landscape before this paper. Muon achieves fast convergence through orthogonalized momentum, but the Newton-Schulz iteration operates on the full momentum matrix, requiring full-size matrix multiplications that are expensive at scale. Dion reduces this cost by making the updates low-rank before orthogonalization, but it must first compute that low-rank representation via Power-Iteration and QR decomposition — both of which scale with rank $r$. The consequence is that Dion's runtime is rank-dependent (Table 1: 8–18% slower at $r = 512$ vs. $r = 128$), creating a practical tradeoff: you can have cheap orthogonalization (low $r$) or accurate updates (high $r$), but not both.

Trion breaks this tradeoff. By obtaining the low-rank representation $b_t$ through DCT column selection — a single matrix multiplication $S_t = B_t D_C$ plus a sort — the cost of getting to the low-rank input for Newton-Schulz is independent of $r$. The Newton-Schulz iteration itself still scales with $r$ (it operates on an $R \times r$ matrix), but the bottleneck that dominated Dion's runtime — the factorization step — is eliminated. Table 1 confirms this empirically: Trion's runtime is approximately constant across $r = 128, 256, 512$ for a given model size, and it is faster than Dion at every rank (2.5–4.5% at $r = 128$, growing to 8–18% at $r = 512$).

This is not just a speedup; it is a qualitative change in the scaling behavior of the optimizer. In Dion, increasing rank incurs a runtime penalty, which creates a practical ceiling on usable rank. In Trion, increasing rank incurs only a modest increase in the Newton-Schulz cost (which operates on small $r \times r$ matrices thanks to the Triton kernel optimizations the paper leverages), making high-rank orthogonal updates feasible without a proportional runtime increase. The paper demonstrates this by comfortably experimenting with $r/d$ ratios up to 50% ($r = 512$ when $d = 1024$), which would be prohibitively expensive in Dion.

This innovation is incremental in concept but fundamental in impact. The idea of "low-rank before Newton-Schulz" was already in Dion. The novelty is in how the low-rank representation is obtained, and the consequence — rank-independent factorization cost — enables a different operational regime for orthogonalized optimizers.

Innovation 4: The Subspace Rotation via Integer Index Sets as a Memory-Efficient Alternative to Storing Projection Matrices

The DCT-AdamW optimizer introduces a specific mechanism that is easy to overlook as a minor optimization but represents a genuinely novel approach to the subspace transition problem that all low-rank Adam-style optimizers must solve. When an optimizer changes its projection subspace between steps — because the gradient's principal directions have shifted — the momentum buffers accumulated in the old subspace must be expressed in the new subspace before new gradients can be incorporated. LDAdam (Robert et al., 2025) solved this by storing two full projection matrices per layer: $Q_\text{prev}$ and $Q_\text{crt}$, each of size $d \times r$. The rotation is then $m_{t-1}^\text{full} = m_{t-1} Q_\text{prev}^\top$ followed by $m_t = \beta_1 (m_{t-1}^\text{full} Q_\text{crt}) + (1-\beta_1) g_t$.

DCT-AdamW's insight is that because all projection matrices are column subsets of a single, globally stored DCT matrix, the rotation can be computed directly in the $r$-dimensional space as $R = Q_\text{prev}^\top Q_\text{crt} \in \mathbb{R}^{r \times r}$, without ever materializing the $d \times r$ projection matrices as stored state. The only per-layer state needed is two sets of $r$ integers — the column indices of the previous and current projections. The rotation matrix $R$ itself is computed on-the-fly from the global DCT matrix and these indices, used once to rotate the momentum buffers, and then discarded.

This matters because the memory footprint of projection matrices scales as $L \times d \times r$ for $L$ layers, which becomes substantial in deep models with high embedding dimensions. The paper quantifies this in the Llama-2 7B fine-tuning experiments (Appendix H, Table 7): at $r = 32$, DCT-AdamW saves 8 GB of memory (23.4%) compared to LDAdam's approach. At $r = 512$, the saving is 6.3 GB (18.2%). These are not marginal savings — they represent the difference between fitting training on a given GPU configuration or not.

The conceptual contribution is recognizing that the projection matrices themselves are redundant when they are subsets of a shared basis. This is only possible because of the fixed-basis paradigm from Innovation 1: if each layer computed its own basis from scratch (as in SVD-based methods), there would be no shared structure to exploit, and the rotation would necessarily require storing the basis vectors. The integer-index approach is a direct architectural consequence of the fixed-basis design, and it demonstrates how the paradigm shift enables downstream optimizations that were structurally impossible in the data-derived basis paradigm.

This innovation is incremental within the paper's own framework but fundamental relative to prior work: it is a natural optimization given the DCT-based approach, but it represents a qualitatively different way of handling subspace transitions compared to any prior low-rank optimizer.

5. Experimental Analysis

Evaluation Methodology

• Dataset. All pretraining experiments use the C4 dataset (Raffel et al., 2020) — a large, cleaned version of Common Crawl. Models are trained on Chinchilla-optimal token counts (20 tokens per parameter) with sequence length 512, except for the LDAdam vs. DCT-AdamW comparison which uses 100 tokens per parameter (80B tokens for Llama-800M). Fine-tuning experiments use the GSM-8k dataset (grade-school math word problems) on Llama-2 7B and Qwen-2.5-7B models.

• Base model(s). Pretraining uses Llama-family models at three scales: 350M, 800M, and 1.3B parameters, trained from scratch. Fine-tuning uses Llama-2 7B and Qwen-2.5-7B. The Llama architecture is chosen because it is representative of contemporary transformer designs and is the standard backbone for low-rank optimizer evaluations in the GaLore lineage. The embedding dimensions $d_\text{model}$ for these models (which determine the DCT matrix size) are 1024 for Llama-350M, 1536 for Llama-800M, and 2048 for Llama-1.3B.

• Metrics. Three categories of metrics are reported: (1) Optimization quality: training loss, validation loss, and validation perplexity (for pretraining), and evaluation accuracy on GSM-8k (for fine-tuning). (2) Memory usage: maximum allocated GPU memory read directly from PyTorch's memory profiler, reported in gigabytes. (3) Runtime: wall-clock training time, reported as total hours/minutes for the full training run, with the lowest runtime across multiple runs being reported to minimize noise from hardware variability.

• Baselines. The paper compares against the specific optimizers it aims to improve: (1) Dion (Ahn et al., 2025) — the direct predecessor to Trion, using Power-Iteration with QR decomposition for low-rank orthogonal momentum updates. (2) LDAdam or LDAdamW (Robert et al., 2025) — uses Block Power-Iteration at every step with subspace rotation via stored projection matrices. (3) FRUGAL (Zmushko et al., 2024) — uses SVD every 200 steps, feeds projection error to SignSGD; also includes RandPerm and Random projection variants as internal baselines. (4) FIRA (Chen et al., 2024) — uses SVD every 200 steps with norm-based error scaling. (5) AdamW (Loshchilov & Hutter, 2019) — full-rank optimizer, included as a reference point but not the primary comparison target. (6) GaLore (Zhao et al., 2024) — for the Qwen-2.5-7B fine-tuning comparison.

• Generation budget / compute accounting. The paper does not use a "generation budget" framework (as in inference-time compute papers). Instead, fair comparison is ensured by running all optimizers on the same hardware (8× H100 NVIDIA GPUs in DDP), with the same global batch size (512), local batch size (64 per GPU, or 32 for 1.3B to fit memory), same number of training tokens, and same hyperparameters (learning rate $\eta = 0.01$, weight decay $\lambda = 0.01$ for Trion/Dion, except where otherwise noted). Runtime is measured as total wall-clock time for the full training run. Memory is measured as peak allocated GPU memory. For methods that update the subspace periodically (FRUGAL, FIRA at $T_u = 200$), the same update interval is used for the DCT variants to ensure the comparison isolates the projection method.

• Cross-validation / statistical protocol. Pretraining results for Trion vs. Dion in Table 1 report averages across 3 random seeds for training/validation loss and perplexity, with the lowest runtime across all runs reported. This provides a basic variance estimate. The fine-tuning results in Appendix H do not report seed variance, which is a limitation. For the FRUGAL/FIRA experiments, the paper states that they use the "best hyper-parameters" from the original papers but does not describe a hyperparameter sweep protocol — this could advantage the DCT variants if the SVD baselines were not re-tuned under the same conditions.

Main Quantitative Results

Pretraining with Trion vs. Dion (Table 1, Figure 1)

The headline result is that Trion matches or exceeds Dion's optimization quality while reducing runtime and memory, with the runtime advantage growing as rank increases.

Perplexity and loss (Table 1). Across all three model sizes (350M, 800M, 1.3B) and three ranks (128, 256, 512), Trion achieves consistently lower training and validation loss than Dion. For example, on Llama-800M with rank $r = 256$, Trion achieves train perplexity 20.04 vs. Dion's 20.33, and validation perplexity 22.06 vs. 22.32. The gap is not enormous — roughly 0.2–0.3 perplexity points — but it is consistent across all 9 model-rank combinations tested. This is a "recovery" result: the DCT approach does not degrade optimization quality; it slightly improves it.

This improvement is not merely due to the DCT basis being a good approximation. Figure 2 (discussed in Section 3) shows that Trion achieves lower Frobenius projection error on the momentum matrix than Dion, and the error for Trion shows a decreasing trend over training steps for some layers. This suggests that the DCT column selection is adaptively improving as the momentum structure stabilizes during training, while Dion's Power-Iteration gives roughly constant approximation quality.

Memory usage (Table 1). Trion uses approximately 10% less memory than Dion across all configurations. On Llama-800M, Trion consumes 14.35 GB vs. Dion's 15.94 GB at $r = 256$ — a saving of 1.59 GB. The gap is slightly smaller at lower ranks (~1.3 GB at $r = 128$ across model sizes) because both optimizers store smaller low-rank momentum buffers, but the relative saving is larger. The memory advantage comes from storing one shared $d \times d$ DCT matrix per GPU plus $r$ integers per layer, versus Dion storing a projection matrix for each layer.

Runtime and its rank-dependence (Table 1, Figure 1 bottom row). This is the paper's most consequential quantitative result. For Llama-800M:

• At rank $r = 128$: Trion takes 6h 45m vs. Dion's 7h 2m (2.5% faster).
• At rank $r = 256$: Trion takes 6h 54m vs. Dion's 7h 33m (8.6% faster).
• At rank $r = 512$: Trion takes 6h 54m vs. Dion's 7h 52m (12.3% faster).

Two patterns are critical. First, Trion's runtime is nearly constant across ranks (6h 45m, 6h 54m, 6h 54m) — a 2.2% variation that may be within measurement noise. Second, Dion's runtime grows substantially with rank (7h 2m, 7h 33m, 7h 52m) — a 11.8% increase from $r = 128$ to $r = 512$. The same pattern holds for Llama-350M (Trion: ~3h 28m flat; Dion: 3h 34m, 3h 49m, 4h 14m) and Llama-1.3B (Trion: ~11h 56m flat; Dion: 12h 14m, 12h 25m, 12h 57m).

The rank-independence of Trion's runtime is the paper's primary claimed advantage, and Table 1 provides direct evidence. The mechanism: Dion's QR decomposition costs $O(R r^2 + C r^2)$, which grows with $r$. Trion's similarity computation $S = GQ$ costs $O(n^2 \log n)$ with Makhoul's algorithm or $O(n^3)$ with matmul, but in both cases the cost is independent of $r$ because the full similarity matrix is always computed regardless of how many columns are subsequently selected. The column selection itself (norm-based sort of $d$ elements, top-$r$ extraction) has negligible cost.

Figure 1 (bottom row) presents the same data as a wall-clock time comparison: for a fixed wall-clock budget, Trion achieves lower training loss than Dion at all three model sizes. The training loss curves (Figure 1, top row) show Trion maintaining a consistent advantage throughout training, with the gap stable rather than widening or narrowing.

Scaling behavior with model size. The runtime advantage of Trion over Dion at $r = 256$ scales as: 2.5% (350M) → 8.6% (800M) → 4.2% (1.3B). The non-monotonic pattern at 1.3B is because that model was trained with local batch size 32 (instead of 64) to fit in GPU memory, which changes the relative fraction of time spent in the optimizer vs. forward/backward passes. This is a limitation of the comparison — the runtime advantages are smaller when the optimizer is a smaller fraction of total training time, and larger models with batch size 32 might show different scaling.

Pretraining with DCT-AdamW vs. LDAdamW (Table 2, Figure 3)

The headline result: DCT-AdamW achieves lower training loss and perplexity than LDAdamW while reducing runtime by 25.75% and memory usage substantially, on Llama-800M trained for 80B tokens (100 tokens/parameter) at rank $r = 320$ ($r/d \approx 20.8\%$ for $d = 1536$).

Loss and perplexity (Table 2, Figure 3). DCT-AdamW achieves train perplexity 14.87 vs. LDAdamW's 15.10, and validation perplexity 17.58 vs. 17.84. These are improvements of 0.23 and 0.26 perplexity points — modest but consistent. Figure 3 shows the training loss curves, where DCT-AdamW is consistently below LDAdamW throughout the 80B-token training run. Full-rank AdamW is included as a reference and achieves better perplexity (train 14.53, validation 17.21), as expected since it uses no low-rank compression.

Memory (Table 2). DCT-AdamW uses 37.83 GB vs. LDAdamW's 45.42 GB — a saving of 7.59 GB (16.7%). The memory advantage is smaller than in the fine-tuning setting (where 23.4% savings were reported at $r = 32$) because the pretraining configuration uses relatively high rank ($r = 320$, 20.8% of $d$) and error feedback quantized to 8 bits, both of which consume significant memory. However, the paper notes that LDAdamW's memory is "close to the memory usage of AdamW because it stores two projection matrices." DCT-AdamW's advantage comes from replacing those two $d \times r$ projection matrices with two sets of $r$ integers plus the shared DCT matrix, and additionally using ZeRO-redundancy (one layer's update computed on one GPU, broadcast to others) to reduce redundant optimizer state across devices.

Runtime (Table 2). DCT-AdamW takes 28h 15m vs. LDAdamW's 39h 22m — a saving of 10h 7m (25.75%). The paper attributes this to replacing Block Power-Iteration (which involves QR decompositions) with a single matrix multiplication and a sort. DCT-AdamW is only 1h 55m (~5%) slower than full-rank AdamW (27h 20m), suggesting that the overhead of low-rank compression has been reduced to a small fraction of total training time. This is a strong practical result: if a practitioner wants low-rank training for memory savings but is deterred by the runtime overhead of SVD/QR-based methods, DCT-AdamW reduces that overhead to near-negligible levels.

What the experiment does NOT show. This is a single run on Llama-800M at a single rank ($r = 320$, or $r/d \approx 20.8\%$). The paper does not sweep ranks for DCT-AdamW in pretraining. The fine-tuning results in Appendix H (Table 7) suggest that DCT-AdamW's advantages over LDAdamW vary with rank — at $r = 32$, LDAdamW achieves better accuracy (32.53% vs. 29.11% for DCT-AdamW with EF, though DCT-AdamW is slightly better without EF), while at $r = 512$, the accuracies are nearly identical (35.86% vs. 35.33%). This suggests that DCT-AdamW may underperform at very low ranks in fine-tuning, which could generalize to pretraining at low ranks, but this was not tested.

Pretraining with DCT-FRUGAL and DCT-FIRA (Table 6, Figure 4)

The headline result: DCT projection approximates SVD projection well in both FRUGAL and FIRA, with modest perplexity tradeoffs in FRUGAL and actual improvements in FIRA, while reducing runtime by 22–24%.

DCT-FRUGAL vs. SVD-FRUGAL (Table 6, Figure 4a). On Llama-800M with 16B tokens:

• SVD-FRUGAL: train perplexity 18.18, val perplexity 22.51, 7h 58m runtime, 62.09 GB memory.
• DCT-FRUGAL: train perplexity 18.70, val perplexity 22.97, 6h 10m runtime, 59.89 GB memory.
• The perplexity degradation is 0.52 train / 0.46 validation — not dramatic, but measurable.
• Runtime savings: 1h 48m (22.6%).
• Memory savings: 2.2 GB (3.5%).

The comparison against random baselines within FRUGAL is particularly informative. The original FRUGAL paper proposes three projection types: SVD, RandPerm (random permutation matrix), and Random (random semi-orthogonal matrix). Table 6 shows that DCT-FRUGAL achieves train perplexity 18.70, compared to RandPerm's 19.83 and Random's 19.63 — an improvement of ~1 perplexity point. This demonstrates that the DCT's structure (the cosine basis) and the dynamic column selection together provide meaningful adaptivity beyond what a random orthogonal projection offers. However, DCT does not fully close the gap to SVD (18.18) in FRUGAL's setting.

DCT-FIRA vs. SVD-FIRA (Table 6, Figure 4b). The same experimental setup but with FIRA's error-handling mechanism:

• SVD-FIRA: train perplexity 17.59, val perplexity 21.87, 8h 0m runtime, 70.43 GB memory.
• DCT-FIRA: train perplexity 17.35, val perplexity 21.72, 6h 6m runtime, 68.43 GB memory.
• DCT-FIRA actually outperforms SVD-FIRA by 0.24 train perplexity and 0.15 validation perplexity.
• Runtime savings: 1h 54m (23.8%).
• Memory savings: 2 GB (2.8%).

This is a notable result: for FIRA's specific error-scaling mechanism, the DCT projection is not just an adequate approximation but potentially beneficial. The paper does not provide a mechanistic explanation for why DCT outperforms SVD in FIRA but underperforms in FRUGAL — this is left as an empirical observation. A plausible hypothesis is that FIRA's error scaling interacts differently with the projection error characteristics: SVD minimizes Frobenius error but may concentrate the error in ways that FIRA's scaling doesn't handle well, while DCT's error may be more uniformly distributed across the orthogonal complement, making FIRA's scaling more effective. But this is speculation — the paper does not investigate this.

The 200-step subspace update interval. Both FRUGAL and FIRA update the projection subspace every 200 steps (the default from GaLore, chosen to amortize SVD's cost). The DCT variants use the same interval to ensure fair comparison. This means the reported runtime savings are from the per-update cost reduction (DCT column selection vs. SVD at each 200-step interval), not from more frequent updates. If the interval were reduced to 1 (updating every step, as in LDAdam/Trion), the runtime advantage of DCT would be much larger because SVD would dominate training time.

Fine-Tuning Results (Appendix H, Table 7, Table 8)

The paper presents fine-tuning results on Llama-2 7B and Qwen-2.5-7B using GSM-8k, across multiple optimizers and ranks. The key findings are:

Memory savings are larger for fine-tuning than pretraining (Table 7). On Llama-2 7B with FRUGAL at $r = 32$, DCT saves 8 GB (23.4%) over SVD. At $r = 512$, the saving drops to 6.3 GB (18.2%). This is because the projection matrix storage ($d \times r$ with $d = 4096$ for Llama-2 7B) is a larger fraction of total memory in the fine-tuning setting (smaller batch sizes, shorter sequences) than in pretraining. The DCT matrix itself ($d \times d$) is ~67 MB in float32 — negligible compared to the per-layer savings.

Accuracy is competitive but not universally superior (Table 7). Across FRUGAL, FIRA, and LDAdamW comparisons with their DCT counterparts:

• DCT-FRUGAL matches SVD-FRUGAL accuracy at $r = 32$ (32.53% vs. 32.75%) but underperforms slightly at $r = 512$ (26.16% vs. 28.89%).
• DCT-FIRA outperforms SVD-FIRA at $r = 512$ (27.37% vs. 26.08%) but underperforms at $r = 32$ (28.81% vs. 29.49%).
• DCT-AdamW vs. LDAdamW shows a similar rank-dependent pattern: LDAdamW is better at $r = 32$ (32.53% vs. 28.66%), while at $r = 512$ they are comparable (35.33% vs. 35.86% with EF).
• Error feedback (EF) does not reliably help DCT-AdamW at low rank: with EF at $r = 32$, DCT-AdamW achieves 29.11% vs. 32.53% for LDAdamW with EF, while without EF, DCT-AdamW achieves 31.99% vs. LDAdamW's 28.66% — a reversal that suggests EF interacts poorly with DCT projection at very low ranks.

Runtime savings in fine-tuning (Table 7). DCT-FRUGAL reduces runtime by approximately 35 minutes (75% of total runtime) compared to SVD-FRUGAL for both ranks on Llama-2 7B. This is a much larger relative saving than in pretraining (22.6%) because the SVD step is a larger fraction of total computation in the fine-tuning setting (shorter overall training, smaller models, smaller batch sizes).

GaLore comparison (Table 8). On Qwen-2.5-7B fine-tuned on GSM-8k, DCT-AdamW (without EF, $T_u = 200$) achieves evaluation accuracy 32.07% at $r = 32$ and 34.80% at $r = 512$, compared to GaLore's 28.05% and 33.51% respectively. This is a 4.02 and 1.29 percentage point advantage — a meaningful improvement over GaLore's SVD-based approach. However, the comparison is not perfectly controlled: DCT-AdamW uses the LDAdam-style momentum rotation and error handling, while GaLore uses a simpler projection-then-AdamW approach. The improvement could be due to the AdamW-style momentum handling rather than the DCT projection itself.

Ablation Studies and Robustness Checks

• DCT projection vs. random projections in FRUGAL (Table 6, Figure 4a): DCT achieves train perplexity 18.70 vs. RandPerm's 19.83 and Random's 19.63, confirming that DCT's structure provides a meaningful advantage (~1 perplexity point) over purely random orthogonal projections. This is the most direct ablation on whether the specific choice of orthogonal matrix matters — it does.

• Projection error comparison: Trion vs. Dion (Figure 2): On Llama-30M, Trion's projection error $\Delta_t^\text{trion} = \|B_t - O_t\|_2$ is consistently lower than Dion's $\Delta_t^\text{dion} = \|B_t - P_t Q_t^\top\|_2$ across multiple linear layers in the first transformer block. The error for Trion shows a decreasing trend for some layers (e.g., the query projection), indicating that the DCT column selection adaptively improves as training progresses. Dion's error is approximately constant. This ablation directly validates the claim that DCT column selection can outperform Power-Iteration for momentum compression.

• Rank dependence of runtime: Trion vs. Dion (Table 1): Trion's runtime varies by at most 2.2% across ranks 128–512 for a given model size, while Dion's runtime increases by 8–18%. This is not an ablation in the traditional sense, but it is the key robustness check on the paper's central claim of rank-independence — the result holds across three model sizes and three ranks.

• Makhoul's algorithm (FFT) vs. matmul (Table 4, Table 5 in Appendix D): The paper benchmarks the DCT computation for different matrix shapes and datatypes. For float32 inputs, Makhoul's algorithm achieves up to 50× speedup over matmul for matrices with more columns than rows (Table 4). For the practically relevant bfloat16 setting (Table 5), matmul-bfloat16 is faster than Makhoul-float32 for $R \geq C$ layers, while Makhoul is 3.5× faster for $R < C$ layers. This ablation quantifies the conditions under which the FFT acceleration is beneficial and justifies the paper's pragmatic approach of using whichever method is faster for a given layer shape and datatype.

• Error feedback quantization (Section 2.4): The paper attempted quantizing the error feedback buffer to 4 bits and found it "degrades the optimizer performance." 8-bit quantization preserves accuracy. This is mentioned briefly without a dedicated table or figure — it is a reported negative result that justifies the 8-bit design choice.

• Effect of error feedback on DCT-AdamW in fine-tuning (Table 7): For DCT-AdamW at $r = 32$ on Llama-2 7B, adding error feedback reduces accuracy from 31.99% to 29.11%. At $r = 512$, EF has essentially no effect (35.33% with EF vs. 35.48% without). This suggests EF interacts poorly with DCT projection at very low ranks — possibly because the projection error is large and feeding it back introduces noise rather than signal.

• Subspace update interval ($T_u$) comparison (Table 3, Algorithm 3): The paper compares optimizers by their update frequency: GaLore/FRUGAL/FIRA at $T_u = 200$, LDAdam/Dion at $T_u = 1$, and DCT methods supporting any $T_u$. The experiments use the same $T_u$ as the original methods (200 for FRUGAL/FIRA, 1 for LDAdam/Trion) to ensure fair comparison. This is not a direct ablation of $T_u$ within DCT methods — the paper does not sweep $T_u$ to show how performance varies with update frequency for the DCT variants.

Critical Assessment

Claim 1: DCT-based dynamic column selection replaces SVD/QR projections at lower cost while matching or exceeding accuracy.

What was tested: The paper compares DCT variants against the original SVD/QR-based optimizers on Llama models from 350M to 1.3B for pretraining, and 7B for fine-tuning, across multiple ranks. The metrics include perplexity, loss, accuracy, memory, and runtime.

What the experiments demonstrate: For Trion vs. Dion, the claim is well-supported: Trion matches or slightly exceeds Dion's perplexity at all ranks and model sizes, with lower memory and rank-independent runtime (Table 1, Figure 1). For DCT-AdamW vs. LDAdamW, the single pretraining run (Table 2) supports the claim. For FRUGAL/FIRA, the claim holds with qualifications: DCT-FRUGAL shows a 0.52 perplexity degradation vs. SVD (but a ~1 point improvement over random projections), while DCT-FIRA actually outperforms SVD (Table 6).

What was NOT tested: (a) All pretraining experiments use models ≤ 1.3B parameters. The paper acknowledges this limitation explicitly ("additional work is required to test our technique for larger models"). The DCT matrix is $d \times d$, and at larger embedding dimensions ($d = 4096, 8192, \dots$), this matrix itself becomes a memory concern, though the paper argues it is still cheaper than per-layer projection matrices for deep enough networks. The claim of scalability to larger models is plausible but unverified. (b) All experiments use Chinchilla-optimal or 100 tokens/parameter — relatively short training runs. Whether the DCT projection maintains its advantage over SVD in longer training runs (where the gradient statistics might drift more) is not tested. (c) For DCT-AdamW in pretraining, only one rank ($r = 320$, ~20.8% of $d$) is tested. The fine-tuning results suggest accuracy may degrade relative to LDAdamW at very low ranks ($r = 32$, $r/d \approx 0.78\%$ for Llama-2 7B's $d = 4096$), but this was not verified in pretraining. (d) The comparisons use the published hyperparameters of the original methods. If these hyperparameters were tuned for the original SVD/QR approaches and are suboptimal for DCT variants, the DCT results might be understated. If the SVD baselines were not re-tuned under the exact same conditions (hardware, batch size, token count), the SVD results might be understated. The paper does not describe a hyperparameter sweep protocol, which is a weakness.

Claim 2: The approach achieves rank-independent runtime — a fundamental advantage over QR/SVD-based methods.

What was tested: Table 1 shows Trion's runtime varying by at most 2.2% across ranks 128, 256, 512 for a given model size, while Dion's varies by 8–18%. Figure 1 (bottom row) confirms the wall-clock advantage.

What the experiments demonstrate: The empirical evidence for rank-independence of the DCT approach is strong and consistent across three model sizes. The theoretical explanation (DCT computation cost depends on $d$, not $r$) is sound.

What was NOT tested: (a) The rank-independence is demonstrated only for the DCT-momentum factorization step in Trion. For DCT-AdamW, the subspace rotation step involves an $r \times r$ matrix multiplication ($m_{t-1} \cdot R$), which does scale with $r$, though only quadratically in $r$ rather than cubically. The paper does not profile DCT-AdamW's runtime across ranks. (b) The Newton-Schulz step in Trion also scales with $r$ (it operates on $R \times r$ matrices). The paper leverages Triton kernels from the Dion repository to make this efficient, but this cost is not zero and grows with $r$. The claim of "rank-independent runtime" is true for the factorization step but is an approximation for the full optimizer — a more precise statement would be "factorization cost is rank-independent, and overall runtime grows slowly with rank." (c) The runtime measurements are end-to-end training time, which includes forward/backward passes. At smaller model scales (350M) or with small local batch sizes (32 for 1.3B), the optimizer is a smaller fraction of total time, and the absolute runtime savings are modest (e.g., 2.5% at 350M, $r = 128$). The practical significance of rank-independence grows with model size and rank — exactly the regime where the paper lacks experiments (>1.3B parameters).

Claim 3: The DCT is a principled choice of basis, not an arbitrary one, due to its connection to the gradient eigenbasis via the circulant-diagonal decomposition.

What was tested: The theoretical argument in Section 4.2 is not directly tested. The empirical validation is indirect: DCT projection outperforms random projections (RandPerm, Random) by ~1 perplexity point in FRUGAL (Table 6). Figure 2 shows DCT achieves lower projection error than Power-Iteration (which approximates the SVD basis), suggesting the DCT basis + selection is competitive with the actual data-derived basis.

What the experiments demonstrate: The DCT is better than random orthogonal matrices — a necessary condition for the "principled choice" claim, but not sufficient. The claim that DCT approximates the SVD eigenbasis specifically (as opposed to just being "a good enough basis for gradients in general") is not directly tested. An ablation that compared DCT against other structured orthogonal matrices (e.g., a Hadamard matrix where dimension permits, or a learned orthogonal matrix) would strengthen the claim but is absent.

Genuine weakness: The circulant-diagonal decomposition argument in Section 4.2 is a theoretical motivation, not a bound. It shows that $GG^\top$ can be expressed as a product of many diagonal and circulant factors, and that a "linear approximation" (keeping one factor) gives $F D_1 F^*$, implying the DFT approximates the eigenvectors. The quality of this approximation — how many factors are needed for a given error tolerance, whether one factor is sufficient in practice — is not analyzed theoretically or empirically. The argument provides a plausible reason for why DCT should work, but it does not predict how well it should work, or under what conditions it might fail.

Missing experiments that would have strengthened the paper:

1. Scaling to larger models (3B, 7B, 13B pretraining): This is the most obvious gap. The paper's central claim — that DCT projection is a drop-in replacement for SVD/QR in low-rank optimizers — would be more convincing with evidence that it works at scales where SVD/QR costs dominate. The authors acknowledge this limitation.

2. Sweeping ranks for DCT-AdamW pretraining: The single pretraining run at $r = 320$ leaves open the question of how DCT-AdamW performs at lower ranks (which is where memory savings are most needed). The fine-tuning results suggest potential issues at $r = 32$.

3. Hyperparameter sensitivity analysis: How sensitive are the DCT variants to the choice of norm ($\ell_1$ vs. $\ell_2$) for column ranking? To the subspace update interval $T_u$? To the DCT variant (DCT-II vs. DCT-III)? To the momentum decay parameter $\mu$ in Trion? The paper uses a single configuration for each optimizer and does not explore these axes.

4. Comparison against a learned or data-dependent orthogonal basis: One could imagine pre-computing a global orthogonal basis by aggregating SVDs across layers from a few training steps, rather than using the fixed DCT matrix. This would be a middle ground between per-layer SVD and fixed DCT, and would help isolate whether the gains come from "fixed basis" or "DCT specifically."

5. Longer training runs: The pretraining experiments use 20 tokens/parameter (Chinchilla-optimal), which is ~3.2B tokens for Llama-160M-class models. Training for 100B+ tokens would test whether DCT projection degrades relative to SVD as gradient statistics evolve over very long horizons.

6. Measurement of the "linear approximation" quality from Section 4.2: The theoretical justification for DCT could be empirically validated by measuring how well the DCT basis captures the gradient's energy compared to the true SVD basis on real training data. Specifically: compute the SVD of $GG^\top$ for a given layer, compute the fraction of energy captured by the top-$r$ SVD vectors vs. the top-$r$ DCT vectors selected by alignment, and track this over training. This would directly quantify the quality of the DCT approximation to the eigenbasis.

Conditional nature of the claims:

• The claim that DCT matches SVD/QR accuracy holds for ranks $\geq 128$ and $r/d \geq 1/16$ (6.25%) in the tested settings. At very low ranks ($r = 32$, $r/d < 1\%$ in large models), the fine-tuning results show that DCT-AdamW can underperform LDAdamW, suggesting a lower bound on usable rank.
• The claim of runtime independence from rank holds for the DCT factorization step specifically. The Newton-Schulz step and the subspace rotation step both scale with $r$, so total runtime is not strictly rank-independent — it is more accurate to say that the dominant scaling bottleneck (factorization) has been removed.
• The claim that DCT is a "principled choice" is supported theoretically (Section 4.2) and empirically vs. random baselines, but the magnitude of the advantage over random projections is modest (~1 perplexity point in FRUGAL), suggesting that most of the benefit comes from the fixed-basis-dynamic-selection paradigm itself, with DCT providing an additional but smaller edge.
• All claims are validated only on the Llama architecture with C4 pretraining data and GSM-8k fine-tuning. Generalization to other architectures (e.g., non-transformer models, vision transformers) and tasks is assumed but not tested.

6. Limitations and Trade-offs

6.1 The DCT Matrix Must Be Materialized at Full Size, Creating a Memory Floor That Scales With $d$

The assumption or constraint. The paper's approach requires materializing a single $d \times d$ DCT matrix per GPU, where $d$ is the model's embedding (hidden) dimension. This matrix is shared across all layers and is the mechanism that avoids storing per-layer $d \times r$ projection matrices — but it is not free. The paper explicitly acknowledges this only indirectly, by stating in Section 2.1 that "the memory overhead of our dynamic column selection approach is the cost of storing only one orthogonal matrix $Q \in \mathbb{R}^{d_{\rm model} \times d_{\rm model}}$ (DCT in our case) per GPU for the entire model and $r$ integers for the corresponding column indices."

The consequence. As embedding dimensions scale in modern LLMs — $d = 4096$ for Llama-2 7B (the largest model in the paper's experiments), $d = 8192$ for Llama-3 70B, $d = 16384$ for some recent models — the DCT matrix itself becomes a meaningful memory consumer. At float32 precision, a $d \times d$ DCT matrix consumes $4d^2$ bytes: approximately 67 MB at $d = 4096$ (negligible), 268 MB at $d = 8192$ (noticeable), and over 1 GB at $d = 16384$ (significant, especially relative to per-GPU memory budgets of 40 or 80 GB). The paper's experiments top out at $d = 2048$ for pretraining (Llama-1.3B) and $d = 4096$ for fine-tuning (Llama-2 7B), so this scaling concern is untested. The paper's claim that the DCT approach is cheaper than storing per-layer projection matrices depends on the inequality $d^2 < L \times d \times r$, i.e., the shared DCT matrix is cheaper than $L$ per-layer matrices of size $d \times r$. For a model with $L = 32$ layers and $r = 256$, this becomes $d^2 < 32 \times d \times 256$, which simplifies to $d < 8192$. For $d = 8192$, the two approaches have comparable memory. For $d > 8192$ or models with fewer but wider layers, the shared DCT matrix is worse than per-layer storage. This is not a failure of the approach at the scales tested, but it represents a regime where the memory argument for DCT over per-layer methods weakens or reverses.

What evidence exists in the paper. None directly. The paper does not benchmark or discuss the memory cost of the DCT matrix at large $d$, does not provide an ablation varying $d$ independently of model size, and does not analyze the break-even point $d$ where the shared matrix becomes more expensive than per-layer storage. The largest model in the pretraining experiments is Llama-1.3B with $d = 2048$, where the DCT matrix is approximately $2048^2 \times 4 \approx 16.8$ MB — negligible enough that the question doesn't arise. The fine-tuning experiments on Llama-2 7B ($d = 4096$) show 23.4% memory savings at $r = 32$ (Table 7), which is consistent with $d^2 \ll L \times d \times r$ at that scale, but the trend as $d$ grows is not explored.

Mitigation status. Not addressed. The paper does not discuss this scaling limitation, propose any mitigation (e.g., storing the DCT matrix in bfloat16, computing DCT columns on-the-fly rather than materializing the full matrix, or using a smaller representative basis), or suggest it as future work. A practitioner deploying this on very large models would need to do their own analysis of the DCT matrix memory cost relative to their specific $d$, $L$, and $r$.

---

6.2 The FFT Acceleration Is Practically Disabled for bfloat16 Training, Eliminating the Asymptotic Runtime Advantage at Standard Precision

The assumption or constraint. The paper motivates DCT in part by its FFT-accelerable structure: Makhoul's $N$-point algorithm computes the similarity matrix $S = GQ$ in $O(n^2 \log n)$ time instead of $O(n^3)$ for standard matrix multiplication (Section 2.2, Appendix C, Appendix D). This is presented as a key advantage of DCT over a generic orthogonal matrix. However, the paper candidly acknowledges in Appendix D that Makhoul's algorithm "can be run only on float32 inputs" 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 bfloat16." Standard LLM training uses bfloat16 or float16 mixed precision — gradients are accumulated in float32 internally, but the gradient tensors made available in the optimizer step are typically bfloat16. Running Makhoul's algorithm requires casting these gradients to float32, executing the FFT, and casting back — which adds memory and compute overhead that may negate the asymptotic speedup.

The consequence. In the practically relevant setting of bfloat16 training, the asymptotic $O(n^2 \log n)$ advantage of the DCT over a generic orthogonal matrix is largely theoretical for now. The paper's own benchmarks in Table 5 (Appendix D) show that for the common case where the gradient matrix has $R \geq C$ (rows $\geq$ columns, which is typical for the feed-forward layers in transformers where the hidden dimension expansion factor is 4×), the standard bfloat16 matmul is consistently faster than the float32 Makhoul algorithm. The speedup from Makhoul's algorithm only materializes for layers where $R < C$ (Table 5 shows a 3.5× speedup for these shapes), and even then it is much smaller than the 8–50× speedup observed when comparing float32-to-float32 (Table 4). This means that for the majority of layers in a standard transformer (where $R \geq C$), the DCT approach does not benefit from the FFT acceleration in current frameworks — it uses standard $O(n^3)$ matmul like any other orthogonal matrix, and its runtime advantage over SVD/QR comes entirely from avoiding the factorization step, not from the FFT structure. The paper's headline claim in the abstract — that DCT can be computed via FFT in $O(n^2 \log n)$ time — is technically true but practically misleading for the current software ecosystem.

What evidence exists in the paper. The paper is transparent about this: Appendix D contains explicit benchmarks (Table 4 vs. Table 5) comparing float32 Makhoul, float32 matmul, and bfloat16 matmul, and clearly states the PyTorch complex-bfloat16 limitation. Table 5 quantifies the issue: for a 2048×2048 matrix (the embedding dimension of Llama-1.3B), bfloat16 matmul takes 0.56ms while float32 Makhoul takes 1.30ms — the matmul is 2.3× faster. However, the paper does not adjust its abstract claims or the main text emphasis on FFT acceleration to reflect this practical limitation. The abstract states "DCT can be computed via Makhoul's $N$-point algorithm based on Fast Fourier Transform (FFT) in $O(n^2 \log n)$ time" without qualification, while the limitation is buried in Appendix D.

Mitigation status. Partial, via a pragmatic fallback: the paper's Algorithm 1 includes both options ("Makhoul($B_t$) or $S_t = B_t \cdot D_C$") and the implementation can choose the faster method per layer based on shape and datatype. The paper also notes that the float32 gradient accumulator ("when running training in mixed precision, the gradients are computed in bfloat16 and accumulated in a float32 buffer") could be leveraged to run Makhoul's algorithm without explicit casting overhead, but this is presented as a possibility rather than an implemented solution. Future PyTorch releases may add complex-bfloat16 support, which would unblock the FFT acceleration for standard training — but the paper cannot and does not guarantee this. In the meantime, the practical DCT advantage over SVD/QR comes from rank-independent factorization cost and simpler memory layout, not from the FFT.

---

6.3 All Pretraining Experiments Are Limited to Models ≤ 1.3B Parameters With Short Training Horizons

The assumption or constraint. The paper's pretraining experiments use Llama models with 350M, 800M, and 1.3B parameters, trained on Chinchilla-optimal token counts (20 tokens per parameter) — corresponding to 7B, 16B, and 26B tokens respectively — except for one longer run at 100 tokens/parameter (80B tokens) for the DCT-AdamW vs. LDAdamW comparison. The authors explicitly acknowledge this in Section 6: "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."

The consequence. Several claims remain unvalidated at scale, and the risk is that the method's advantages may not generalize — or may reverse — in regimes that matter for production LLM training. Specifically:

1. Runtime advantage relative to total training time: At 350M parameters, the optimizer is a small fraction of total step time (dominated by forward/backward passes), so the paper reports only a 2.5% speedup from Trion over Dion at $r = 128$ (Table 1). As model size grows, the optimizer's per-step cost grows with $d$ (the embedding dimension), while forward/backward costs grow with the full parameter count. It is unclear whether the relative speedup of Trion over Dion increases or decreases at 7B, 13B, or 70B scales. If the optimizer becomes a smaller fraction of total time at larger scales (because forward/backward dominates more), then even a large relative improvement in optimizer time may translate to a negligible wall-clock gain.

2. Memory savings at scale: The DCT approach saves memory by replacing per-layer $d \times r$ projection matrices with one shared $d \times d$ matrix plus $r$ integers per layer. As discussed in Limitation 6.1, this tradeoff depends on $d$, $L$, and $r$. At larger scales, both the per-layer projection matrices and the shared DCT matrix grow, and it is not obvious which grows faster relative to other memory consumers (activations, parameters, optimizer states). The paper's experiments at ≤1.3B parameters do not cover the regime where this tradeoff is most consequential.

3. Training stability over long horizons: The paper's pretraining runs use 7–80B tokens. Modern LLM training runs use trillions of tokens. Gradient statistics may drift substantially over such long horizons, and it is unknown whether the DCT's fixed basis remains an adequate approximation of the gradient eigenbasis over thousands of training steps. The projection error plots in Figure 2 (on a 30M model, trained for an unspecified but presumably short duration) show Trion's error decreasing for some layers — but whether this trend continues, plateaus, or reverses over very long training is untested. A fixed basis might become progressively less aligned with the gradient structure as the model moves through different phases of learning (e.g., from learning general patterns to memorizing specific examples).

4. Interaction with learning rate schedules: All experiments use constant learning rates (or simple schedules inherited from the baseline optimizers). Long training runs typically use cosine or warmup-stable-decay schedules, which change the effective gradient magnitude over time and could interact with the DCT column selection (e.g., very small gradients might produce noisy column rankings).

What evidence exists in the paper. The authors are transparent: the limitation is stated in Section 6. No experiments at >1.3B parameters pretraining are reported. The fine-tuning experiments on 7B models (Appendix H) provide partial evidence that the method scales to larger models in a different training regime (short fine-tuning, ~1000–10000 steps), but fine-tuning has different gradient statistics (smaller effective learning rates, narrower task distribution) than pretraining from scratch. The Chinchilla-optimal training runs at 20 tokens/parameter are relatively short by modern standards and may not reveal long-horizon issues.

Mitigation status. Acknowledged as future work. The paper states that testing at larger scales "would require significantly more computational resources" — a reasonable constraint for an academic group, but one that leaves the scalability claims as plausible hypotheses rather than demonstrated facts. The paper does not provide theoretical arguments for why the method should scale well, beyond the $d$-dependent memory analysis discussed in Limitation 6.1.

---

6.4 Performance Relative to SVD/QR-Based Methods Degrades at Very Low Ranks

The assumption or constraint. The paper positions DCT-based projection as a drop-in replacement for SVD/QR across all low-rank optimizers, ranks, and tasks. The dynamic column selection procedure is designed to work for any rank $r$, and the optimality proof in Section 4.1 holds for any $r \leq n$. However, the empirical results reveal that DCT-AdamW's performance relative to LDAdamW is rank-dependent, with a degradation at low ranks that the paper does not explain or theoretically bound.

The consequence. In the fine-tuning experiments on Llama-2 7B (Table 7, Appendix H), at rank $r = 32$ (corresponding to $r/d = 32/4096 \approx 0.78\%$), DCT-AdamW without error feedback achieves 31.99% evaluation accuracy vs. LDAdamW's 28.66% — a 3.33 point advantage. However, with error feedback (which is the recommended configuration in both LDAdamW and DCT-AdamW), DCT-AdamW achieves only 29.11% vs. LDAdamW's 32.53% — a 3.42 point disadvantage. At $r = 512$ ($r/d \approx 12.5\%$), the gap nearly disappears: 35.33% vs. 35.86% with EF, and 35.48% vs. 35.56% without EF. This suggests that at very low ranks: (a) DCT-AdamW without EF can be competitive or better than LDAdamW, (b) the error feedback mechanism — which is essential for LDAdamW's performance at low ranks (improving from 28.66% to 32.53%) — actually harms DCT-AdamW (dropping from 31.99% to 29.11%). The paper notes this observation but does not explain it: "We would like to emphasize that EF does not help DCT-AdamW in comparison to LDAdamW for $r = 32$."

The practical consequence is that at very low ranks — precisely the regime where memory savings are most critical and where low-rank methods are most appealing — DCT-AdamW with its standard configuration (error feedback included) may underperform the SVD-based alternative. A practitioner wanting to use DCT-AdamW at very low ranks would need to disable error feedback, losing a mechanism that is important for convergence in the baseline method.

A plausible mechanism: at very low ranks, the DCT projection error $\Xi_t = G_t - g_t Q_{\text{crt}}^\top$ is large (since only $r/n \approx 0.78\%$ of the basis vectors are retained). Feeding this large error back into the gradient essentially adds a large noise term. For LDAdamW, the projection matrices are closely aligned with the gradient's principal components (via Block Power-Iteration), so the projection error is mostly noise that EF can usefully accumulate. For DCT-AdamW, the projection error includes signal (gradient directions that don't align with any of the top-$r$ DCT columns) that may be large and uncorrelated with future gradients, making EF inject noise rather than recoverable signal. This is speculation — the paper does not investigate the mechanism.

What evidence exists in the paper. The rank-dependent performance is visible in Table 7 (Appendix H). The failure of EF at low rank is noted in the text discussion of fine-tuning results: "EF does not help DCT-AdamW in comparison to LDAdamW for $r = 32$." No pretraining experiments sweep ranks for DCT-AdamW, so it is unknown whether the same rank-dependence appears in pretraining. No ablation investigates why EF hurts at low ranks — e.g., by measuring the magnitude of the projection error, its correlation across steps, or its alignment with the true gradient.

Mitigation status. Not addressed. The paper acknowledges the observation but does not explain it, propose a fix (e.g., scaling the EF term by a rank-dependent factor, or disabling EF below a threshold rank), or suggest a rank-adaptation strategy. The failure is treated as a reported fact rather than a problem to solve. This is a meaningful gap for practitioners who need maximum memory savings (hence minimum rank) and would be choosing between LDAdamW and DCT-AdamW in that regime.

---

6.5 No Mechanism for Handling or Detecting When the DCT Basis Becomes Misaligned With Gradient Structure

The assumption or constraint. The DCT-based approach assumes that at every step and for every layer, there exist $r$ DCT basis vectors that adequately capture the gradient's energy. The theoretical justification in Section 4.2 supports this on average — DCT approximates the gradient eigenbasis — but provides no guarantee for individual steps or layers. During training, gradient statistics evolve: early in training, gradients may be large and noisy, while later they may become small and structured. The paper provides no mechanism to detect when the DCT basis is performing poorly (e.g., when even the best $r$ columns capture only a small fraction of the gradient energy) or to fall back to a full-rank update when needed.

The consequence. In the worst case, a particular layer's gradient at a particular step might have its energy spread nearly uniformly across all DCT basis vectors, or concentrated in directions that are orthogonal to the main DCT components. In that case, selecting the top $r$ columns would capture only approximately $r/n$ of the gradient energy, and the projection error would be $(1 - r/n) \|G\|_F^2$ — the worst-case bound from Section 4.1. This is the same bound that applies to any orthogonal matrix, so it is not worse than using a random orthogonal basis — but it is also not better, and it is worse than SVD, which by definition captures the maximum possible energy for any rank $r$. The paper's experiments show that on average, DCT performs well (Figure 2 shows lower projection error than Power-Iteration for several layers), but these are averages. The paper never examines the distribution of projection error across steps or layers — are there outliers where DCT performs poorly? If so, do those outliers matter for optimization (e.g., causing spikes in the loss)?

A related concern: the DCT basis is fixed before training and never updated. This means it is chosen without any knowledge of the specific model architecture, initialization, or data distribution. For some architectures or tasks, the gradient structure might be systematically misaligned with the DCT basis — e.g., if the weight matrices have block-diagonal or other special structure that the cosine basis does not capture efficiently. The paper does not explore this sensitivity.

What evidence exists in the paper. Figure 2 shows projection errors for specific layers over time, and Appendix F states that "Trion yields lower projection error than Dion and the trend is the same across all transformer blocks in the model." This is evidence of consistent good performance, but it does not test robustness to adverse conditions. The paper does not report the distribution of captured energy ($\sum_{i \in i_t} \|q_i^\top G\|_2^2 / \|G\|_F^2$) across layers and steps, which would directly quantify how often the DCT basis is performing near the worst-case bound. The experiments use Llama architectures on C4 — whether the DCT basis would perform similarly on, say, a vision transformer with ImageNet or a mixture-of-experts architecture is unknown.

Mitigation status. Not addressed. The paper does not discuss detecting or handling poor alignment, does not propose an adaptive mechanism (e.g., monitoring the fraction of energy captured and falling back to full-rank update if it drops below a threshold), and does not analyze worst-case behavior. The theoretical guarantee in Section 4.1 provides only the average-case bound $(1 - r/n)$ — it says nothing about step-to-step variance. A practitioner concerned about robustness would want to know: does the method ever catastrophically fail on a particular step or layer? The paper provides no evidence either way.

---

6.6 The Empirical Comparison Against SVD/QR Baselines May Overstate DCT's Advantage Due to Uncontrolled Hyperparameter Tuning

The assumption or constraint. The paper compares DCT-based optimizers against published baselines (Dion, LDAdamW, FRUGAL, FIRA) using their published hyperparameters — learning rates, momentum decay factors, weight decay, subspace update intervals. The paper states (Section 3) that Trion results "are obtained using the optimal learning rate $\eta = 0.01$ (as reported by the original Dion paper) and weight decay $\lambda = 0.01$," and similar language appears for the other baselines. The paper does not describe re-tuning the baseline methods under the exact hardware, batch size, token count, and model configuration used in its experiments.

The consequence. If the published hyperparameters are suboptimal for the specific setup the paper uses (e.g., different hardware generation, different batch size per GPU, different token count), then the baseline methods may underperform relative to their true potential. Conversely, the DCT variants may have been implicitly tuned (e.g., the authors may have tried several configurations and reported the best, without stating this explicitly). This creates a risk that the reported gains — particularly the cases where DCT outperforms the SVD/QR baseline (e.g., DCT-FIRA beating SVD-FIRA by 0.24 perplexity in Table 6, or Trion beating Dion by 0.2–0.3 perplexity across all configurations in Table 1) — are partially attributable to better hyperparameter alignment with the experimental setup rather than to the DCT projection itself.

This concern is amplified by specific observations:

• The DCT-FRUGAL results show a 0.52 perplexity degradation vs. SVD-FRUGAL (Table 6), while DCT-FIRA shows a 0.24 perplexity improvement vs. SVD-FIRA. If DCT is an approximation to SVD (as the paper argues), one would expect a consistent small degradation in both cases. The fact that FIRA shows an improvement suggests either that DCT interacts favorably with FIRA's error-scaling mechanism (plausible) or that the SVD-FIRA baseline was not optimally tuned.
• In the fine-tuning results (Table 7), LDAdamW with EF achieves 32.53% at $r = 32$ without explanation of how the EF hyperparameters were chosen. DCT-AdamW's accuracy at the same rank varies from 28.66% (with EF) to 31.99% (without EF) — a 3.33 point swing from a single binary choice. This sensitivity suggests that a more thorough hyperparameter search might change the ranking of methods.

What evidence exists in the paper. The paper does not report a hyperparameter sweep for any method, does not state whether the baseline methods were re-tuned, and does not provide sensitivity curves showing how performance varies with key hyperparameters. The use of published hyperparameters from the original papers is a reasonable starting point but is not sufficient to guarantee a fair comparison, especially when the paper's experimental setup differs from the original papers' setups in hardware, batch size, token count, or model configuration.

The paper does report results across three seeds (Table 1) for Trion vs. Dion, which controls for random initialization but not for hyperparameter sensitivity. The fine-tuning results (Table 7, Table 8) do not report seed variance.

Mitigation status. Not addressed. The paper does not discuss hyperparameter sensitivity, does not perform sweeps, and does not acknowledge this as a limitation. This is a common practice in systems/optimizer papers where computational constraints limit hyperparameter exploration, but it weakens the strength of the claim that DCT "matches or exceeds" SVD/QR performance — the evidence supports "performs comparably," but "exceeds" requires tighter controls.

7. Implications and Future Directions

How This Work Changes the Landscape

This paper shifts the conversation around low-rank optimization from "how do we compute better projection matrices?" to "do we even need to compute them at all?" The core reframing — that a single fixed orthogonal basis with dynamic column selection can replace per-layer matrix factorizations — is conceptually disruptive because it challenges an assumption so natural that it had gone unquestioned across the entire GaLore/Muon/Dion lineage: that projection matrices must be derived from the data they are compressing. The paper demonstrates that this assumption is not false in theory (SVD remains optimal by the Eckart-Young theorem) but is unnecessarily expensive in practice, and that a smart selection from a universal basis can achieve comparable or better empirical results at lower cost.

The magnitude of this shift should be understood as a methodological reframing rather than a paradigm shift. The paper does not propose a new theory of optimization, a new convergence proof, or a fundamentally new algorithm class. Instead, it proposes a drop-in replacement primitive — the DCT-based dynamic column selection — that can be plugged into existing low-rank optimizers (Dion, LDAdam, FRUGAL, FIRA) with minimal changes to their core logic. What makes this significant beyond an incremental speedup is that it resolves a structural tension that prior work could not escape: the tradeoff between rank (higher is better for accuracy) and factorization cost (higher rank means more expensive SVD/QR). By making the factorization cost independent of rank (Table 1: Trion runtime varies by ≤2.2% across $r$ from 128 to 512, while Dion varies by 8–18%), the paper removes the practical ceiling on usable rank that SVD/QR-based methods imposed. A practitioner can now choose rank based purely on accuracy-memory tradeoffs, without worrying that higher rank will slow down each optimizer step.

The paper also reconciles a tension between two lineages of low-rank optimizers that had been developing in parallel. The GaLore/LDAdam/FRUGAL/FIRA family (Adam-style, SVD/QR-based) and the Muon/Dion family (momentum-orthogonalization, Newton-Schulz-based) shared a common bottleneck — expensive per-layer factorizations — but addressed it differently (periodic updates vs. cheaper but still rank-dependent Power-Iteration). The DCT approach is the first method that works for both families without modification: the paper demonstrates integration into Dion (becoming Trion), LDAdam (becoming DCT-AdamW), FRUGAL, and FIRA, with the same DCT matrix and column selection procedure serving all of them. This unification suggests that the fixed-basis-dynamic-selection approach is not a quirk of one optimizer design but a general primitive for low-rank adaptive optimization, and future optimizer designs can adopt it as a building block without reinventing the factorization wheel.

A secondary but important landscape change is the validation that a Fourier-type basis has a principled connection to gradient structure (Section 4.2), not just a heuristic one. Prior work in model compression (e.g., the QuIP line) used random orthogonal or Hadamard matrices as convenient rotations with fast multiplication, justified by the Johnson-Lindenstrauss lemma or similar dimensionality reduction arguments. This paper's argument — that the DCT approximates the gradient's left singular vectors via the circulant-diagonal decomposition — is qualitatively different: it claims the basis is structurally similar to what SVD would compute, not just "a random projection that preserves distances." The empirical evidence is modest (DCT beats random projections by ~1 perplexity point in FRUGAL, Table 6), but the theoretical framing opens a new line of inquiry: what other fixed bases might approximate the gradient eigenbasis for specific architectures or data modalities?

Research directions that become more attractive:

• Developing specialized fixed bases for specific architectures (e.g., bases derived from the structure of convolutional filters, attention patterns, or graph Laplacians), using the same dynamic selection framework.
• Hardware-optimized implementations of the similarity computation $S = GQ$ — since $Q$ is fixed and known at compile time, there is potential for kernel fusion, reduced-precision lookup tables, or specialized tensor cores.
• Dynamic basis adaptation that goes beyond column selection: can the basis itself be slowly updated during training (e.g., via a low-rank perturbation to the DCT matrix) while retaining most of the computational advantages?

Research directions that become less attractive:

• Incremental improvements to Power-Iteration or Block Power-Iteration for low-rank optimizers. The paper shows that a fixed basis with selection can outperform Power-Iteration in reconstruction quality (Figure 2) while being faster. Further engineering effort on iterative SVD approximation for this specific application is hard to justify when the fixed-basis approach is simpler, faster, and comparably accurate.
• Developing new per-layer factorization algorithms (beyond SVD, QR, Power-Iteration) that still derive the basis from the data. The paper's reframing suggests the bottleneck is not the specific factorization algorithm but the paradigm of per-layer, data-derived bases itself.

---

Follow-Up Research This Work Enables

1. Scaling the DCT approach to models with $d_\text{model} \geq 8192$ to identify where the memory tradeoff reverses.

The paper's experiments top out at $d = 4096$ (Llama-2 7B fine-tuning), where the $d \times d$ DCT matrix is ~67 MB in float32 — negligible. For models like Llama-3 70B ($d = 8192$, DCT matrix ~268 MB) or DeepSeek-V2 ($d$ up to 16384, DCT matrix >1 GB), the tradeoff described in Limitation 6.1 becomes quantitatively important. A direct experiment would train a Llama-3 8B or 70B model from scratch (or fine-tune) with DCT-AdamW or Trion at several ranks, measuring: (a) whether the memory savings from eliminating per-layer projection matrices still exceed the cost of the $d \times d$ DCT matrix at these embedding dimensions, (b) whether the break-even point $d > L \times r$ is reached for realistic $(L, r)$ combinations, and (c) whether storing the DCT matrix in bfloat16 (reducing its memory by 2×) causes measurable degradation in projection quality. This experiment would directly validate or refine the paper's central memory-savings claim at scales that matter for production training.

2. Training a dedicated "difficulty predictor" or quality monitor for DCT projection to enable adaptive fallback.

The paper's Section 4.1 proves that the DCT compression is contractive with factor $1 - r/n$, but this is a worst-case bound that applies to any orthogonal matrix. In practice, the fraction of gradient energy captured by the top-$r$ DCT columns — computed as $E_t = \sum_{i \in i_t} \|q_i^\top G\|_2^2 / \|G\|_F^2$ — varies per layer and per step, and the paper never measures this distribution. A natural follow-up would instrument a training run (e.g., on Llama-800M with Trion) to record $E_t$ for every layer at every step, and analyze: (a) what is the distribution of $E_t$? Are there layers or steps where $E_t$ is near the worst-case $r/n$ bound? (b) Do drops in $E_t$ correlate with loss spikes or optimization instability? (c) Can a lightweight predictor (e.g., a small MLP taking the current loss, gradient norm, and step number) anticipate when the DCT projection will perform poorly and trigger a fallback to a full-rank update for that step? This would directly address Limitation 6.5 and could turn the fixed-basis approach from "works well on average" to "robust with guarantees."

3. Ablating the DCT basis against other structured orthogonal matrices to isolate how much the "DCT specifically" matters vs. "any structured basis."

The paper's Appendix C briefly motivates why DCT is preferred over random orthogonal matrices (no FFT structure) and Hadamard matrices (not orthogonal for non-power-of-2 dimensions). But the paper never empirically compares DCT against a learned or data-derived fixed basis — for example, taking the SVD of the gradients from the first 100 training steps of a model, averaging the resulting left singular vectors across layers, orthogonalizing, and using that as the fixed basis. This would isolate whether the advantage comes from (a) the fixed-basis paradigm itself, (b) the DCT's specific structure (cosine basis, approximate eigenbasis), or (c) the FFT-accelerability (which is currently disabled in bfloat16 anyway). A strong follow-up would pretrain Llama-800M with three fixed bases — random orthogonal, DCT, and the "early-training SVD aggregate" — each with the same dynamic column selection, and measure perplexity, runtime, and reconstruction error. If the early-training SVD aggregate outperforms DCT, it suggests that a small upfront cost to learn a data-specific fixed basis pays off. If DCT matches or beats it, it validates the Section 4.2 theoretical argument that DCT is structurally close to the true eigenbasis.

4. Investigating the interaction between error feedback and DCT projection at low ranks.

The paper reports (Appendix H, Table 7) that error feedback harms DCT-AdamW at $r = 32$ on Llama-2 7B fine-tuning (29.11% with EF vs. 31.99% without), while it helps LDAdamW (32.53% with EF vs. 28.66% without). This is a specific, unexplained failure mode that could reveal fundamental limits of the DCT approach. A targeted experiment would: (a) measure the magnitude and autocorrelation of the projection error $\Xi_t = G_t - g_t Q_{\text{crt}}^\top$ for both DCT-AdamW and LDAdamW at $r = 32$ over a fine-tuning run, (b) test whether the DCT projection error is systematically larger or less structured than LDAdamW's (which would explain why feeding it back adds noise rather than recoverable signal), and (c) evaluate whether a simple fix — scaling the EF term by a rank-dependent factor, or disabling EF when $r/d$ falls below a threshold — restores DCT-AdamW's competitiveness at low ranks. This addresses Limitation 6.4 and would provide practitioners with clear guidance on when to enable or disable error feedback.

5. End-to-end training with Makhoul's algorithm enabled via custom CUDA kernels for complex-bfloat16 FFT.

The paper identifies (Appendix D) that the asymptotic $O(n^2 \log n)$ advantage of DCT over generic orthogonal matrices is practically disabled because PyTorch lacks complex-bfloat16 support for FFT. A systems-focused follow-up would implement Makhoul's algorithm as a custom CUDA kernel that operates directly on bfloat16 data (either by emulating complex-bfloat16 or by restructuring the algorithm to avoid complex arithmetic entirely) and benchmark end-to-end training speedups on models with $d \geq 4096$ where the $O(n^3)$ matmul cost becomes meaningful. The paper's Table 4 suggests that for layers with $R < C$, a float32 Makhoul implementation is up to 50× faster than float32 matmul — if even a fraction of that speedup can be realized in bfloat16, it would make the DCT approach strictly faster than any matmul-based alternative (not just rank-independent, but asymptotically cheaper), which would strengthen the case for DCT over all other fixed orthogonal bases.

6. Stress-testing DCT projection on non-transformer architectures and non-language modalities.

All experiments use Llama-family transformer models on text data (C4 for pretraining, GSM-8k for fine-tuning). The DCT's theoretical motivation (Section 4.2) relies on the circulant-diagonal decomposition of $GG^\top$, which is a general matrix property and should apply to any architecture. But gradient structure varies across architectures: convolutional networks have spatially localized gradient structure, graph neural networks have structure induced by the adjacency matrix, and mixture-of-experts models have sparse, routed gradients. A strong negative result — demonstrating that DCT projection fails on, say, a ResNet on ImageNet or a vision transformer on video data — would be as informative as a positive one, because it would reveal the boundary conditions for the "DCT approximates the eigenbasis" argument. A concrete experiment would train a ViT-B/16 from scratch on ImageNet-1k with DCT-AdamW vs. LDAdamW, measuring both accuracy and the projection energy capture $E_t$ across layers, to test whether the DCT basis remains effective when gradient structure is qualitatively different from language model gradients.

---

Practical Applications and Downstream Use Cases

1. Memory-constrained fine-tuning of large language models on consumer GPUs.

The paper's fine-tuning results (Appendix H, Table 7) show that DCT-FRUGAL saves 8 GB (23.4%) of GPU memory vs. SVD-FRUGAL on Llama-2 7B at $r = 32$, while achieving comparable accuracy (32.53% vs. 32.75% on GSM-8k). This directly enables fine-tuning 7B-parameter models on GPUs with 24 GB of VRAM (e.g., RTX 3090/4090) where SVD-based low-rank methods might exceed memory limits, or fitting larger batch sizes within the same budget. A practitioner fine-tuning Llama-2 7B on a custom dataset using LoRA-style low-rank adaptation could replace the SVD step in FRUGAL or FIRA with DCT column selection and recover the 8 GB for larger context lengths or higher ranks, without changing the optimizer's API or the training script's structure. The 75% runtime reduction (35 minutes saved) is a meaningful secondary benefit for rapid experimentation cycles.

2. Large-scale pretraining with orthogonalized momentum (Muon-family) at higher ranks without proportional runtime penalties.

The paper's Table 1 shows that Trion's runtime is approximately constant across ranks 128 to 512, while Dion's grows by 8–18%. For teams running Muon-style pretraining at scale (e.g., training a 7B model with orthogonalized momentum), the practical implication is that they can increase rank — and thus the fidelity of the low-rank approximation to the momentum — without paying a proportional runtime cost. If a team is currently using Dion at $r = 256$ on a 1.3B model and finding the runtime acceptable, upgrading to Trion at $r = 512$ would give them a more accurate low-rank momentum (closer to full-rank Muon quality) while being faster (Trion at $r = 512$ is 8% faster than Dion at $r = 256$ based on Table 1's Llama-800M data, and 12% faster at $r = 256$ compared to Dion at the same rank). This changes the cost-benefit calculation for using orthogonalized momentum: previously, higher rank meant slower steps, so practitioners had to choose a compromise rank. Now, the rank can be set primarily based on accuracy needs.

3. Training infrastructure libraries seeking a single, universal low-rank projection primitive.

The paper demonstrates that the same DCT matrix and column selection procedure works across five different optimizers (Trion, DCT-AdamW, DCT-FRUGAL, DCT-FIRA, and implicitly Muon by reducing Newton-Schulz input size). For teams building training frameworks (e.g., PyTorch's distributed training utilities, Hugging Face's Transformers trainer, or internal training infrastructure at AI labs), this suggests that the DCT projection can be implemented once as a library primitive — dct_project(gradient, rank, dct_matrix) — and reused across all low-rank optimizers the framework supports, rather than maintaining separate SVD, QR, and Power-Iteration implementations for each optimizer. The memory savings from storing only the DCT matrix and integer indices per layer (instead of per-layer projection matrices) also simplify checkpointing and state serialization in distributed training, since the optimizer state for every layer can be reconstructed from the shared DCT matrix and per-layer index sets.

---

When to Prefer This Method

The paper provides enough quantitative comparisons against named alternatives (Dion, LDAdam, FRUGAL, FIRA) to support a conditional decision framework:

• Prefer Trion over Dion when: (1) You are using orthogonalized momentum (Muon-family) and want to use rank $r \geq 256$, where the QR decomposition in Dion's Power-Iteration becomes a meaningful fraction of step time (Table 1 shows 8–18% speedup at $r = 512$ vs. $r = 128$ for Dion); (2) you are training in a distributed setting (DDP/FSDP) and want to communicate only low-rank momentum components ($o_t$, size $R \times r$) rather than full orthogonal updates, which Trion enables by design; (3) you value consistent per-step runtime regardless of rank, for predictable job scheduling.

• Prefer DCT-AdamW over LDAdamW when: (1) Memory is the primary constraint and your model has many layers (deep network) with moderate embedding dimension, where storing two $d \times r$ projection matrices per layer in LDAdamW becomes expensive relative to the single $d \times d$ DCT matrix (the fine-tuning results in Table 7 show 23.4% memory savings at $r = 32$ on Llama-2 7B); (2) you are using moderate to high ranks ($r/d \geq 10\%$), where DCT-AdamW matches or exceeds LDAdamW's accuracy (Table 2: DCT-AdamW outperforms LDAdamW at $r/d \approx 20.8\%$ on Llama-800M; Table 7: comparable accuracy at $r = 512$); (3) you want runtime close to full-rank AdamW (Table 2: DCT-AdamW is only ~5% slower than AdamW vs. LDAdamW being ~44% slower).

• Prefer DCT-FRUGAL or DCT-FIRA over their SVD-based versions when: (1) The SVD step every 200 iterations dominates training time and you are willing to trade ~0.5 perplexity points for a ~23% runtime reduction (Table 6: DCT-FRUGAL degrades by 0.52 train perplexity, saves 1h 48m); (2) you are using FIRA specifically, where DCT actually outperformed SVD in the paper's experiments (Table 6: +0.24 train perplexity improvement).

• Prefer the original SVD/QR methods over DCT variants when: (1) You are using very low ranks ($r/d < 1\%$, e.g., $r = 32$ on a model with $d \geq 4096$), where the fine-tuning results suggest DCT-AdamW with error feedback may underperform LDAdamW (Table 7: 29.11% vs. 32.53% at $r = 32$); (2) your embedding dimension $d$ is very large relative to the number of layers $L$, such that $d^2 > L \times d \times r$ and the shared DCT matrix becomes more expensive than per-layer projection matrices (not yet observed in the paper's experiments but analytically possible at $d > 8192$ for typical $L$ and $r$); (3) you require the strongest possible theoretical guarantee on projection optimality (Eckart-Young) and cannot tolerate any approximation to the SVD basis, regardless of empirical performance.