Paper deep dive
Training Transformers in Cosine Coefficient Space
Mohamed Amine Bergach
Intelligence
Status: succeeded | Model: google/gemini-3.1-flash-lite-preview | Prompt: intel-v1 | Confidence: 95%
Last extracted: 4/10/2026, 2:49:53 AM
Summary
The paper introduces a sparse-coefficient parameterization method for training transformer linear layers using 2D Discrete Cosine Transform (DCT) coefficients. By replacing dense weight matrices with a K-sparse representation in the DCT basis, the authors achieve performance comparable to dense baselines at 50% parameter count, outperforming LoRA at matched parameter counts. The mechanism is identified as 'rank flexibility,' where sparse-coefficient subspaces allow for higher-rank solutions compared to the rank-constrained nature of LoRA. The DCT is specifically highlighted for its separable fast transform, which enables a fused reconstruction kernel that optimizes on-chip memory bandwidth.
Entities (5)
Relation Signals (3)
DCT → enables → Fused Reconstruction Kernel
confidence 95% · the DCT is preferred because its separable fast transform admits a fused reconstruction kernel
Sparse-coefficient parameterization → exhibits → Rank Flexibility
confidence 95% · We identify rank flexibility as the mechanism.
LoRA → constrainedby → Rank
confidence 90% · A LoRA layer is bilinear: its image is the rank-≤r variety
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
Abstract:Linear layers hold most of a transformer's parameters. We replace each linear layer with one that stores $K$ out of $mn$ two-dimensional DCT coefficients per weight matrix and reconstructs the full matrix through an inverse DCT at every forward pass; the $K$ coefficients are the trainable parameters. A 4-layer, 128-dim transformer trained from scratch on character-level Shakespeare reaches validation loss $1.604$ at $K = mn/2$, against $1.580$ for a standard dense baseline -- a gap of $+0.024$ at half the trainable parameter count, within the terminal-epoch variation of the dense run. A rank-48 LoRA factorization at the same trainable parameter count reaches only $1.801$ ($+0.221$). The structural advantage of sparse-coefficient over low-rank parameterizations at matched $K$ is qualitative. We identify rank flexibility as the mechanism. A random orthonormal basis matches the DCT within noise at $K = mn/2$, and a compression sweep through $K = mn/10$ and $K = mn/20$ shows that subspaces that can host high-rank matrices keep the loss low, while subspaces that flatten into a low-rank block (zigzag-selection variants) converge onto the observed stable rank \emph{and} the loss line of the rank-48 LoRA reference in lock-step. Among these orthonormal bases, the DCT is preferred because its separable fast transform admits a fused reconstruction kernel: the materialized weight matrix never leaves on-chip memory, so the parameter saving translates into a bandwidth saving as well.
Tags
Links
- Source: https://arxiv.org/abs/2604.04440v2
- Canonical: https://arxiv.org/abs/2604.04440v2
Trouble viewing inline? Open PDF directly →
Full Text
21,705 characters extracted from source content.
Expand or collapse full text
Training Transformers in Cosine Coefficient Space Mohamed Amine Bergach Illumina, San Diego, CA, USA mbergach@illumina.com Abstract Linear layers hold most of a transformer’s parameters. We replace each linear layer with one that stores K out of mnmn two-dimensional DCT coefficients per weight matrix and reconstructs the full matrix through an inverse DCT at every forward pass; the K coefficients are the trainable parameters. A 4-layer, 128-dim transformer trained from scratch on character-level Shakespeare reaches validation loss 1.6041.604 at K=mn/2K=mn/2, against 1.5801.580 for a standard dense baseline—a gap of +0.024+0.024 at half the trainable parameter count, within the terminal-epoch variation of the dense run. A rank-48 LoRA factorization at the same trainable parameter count reaches only 1.8011.801 (+0.221+0.221). The structural advantage of sparse-coefficient over low-rank parameterizations at matched K is qualitative. We identify rank flexibility as the mechanism. A random orthonormal basis matches the DCT within noise at K=mn/2K=mn/2, and a compression sweep through K=mn/10K=mn/10 and K=mn/20K=mn/20 shows that subspaces that can host high-rank matrices keep the loss low, while subspaces that flatten into a low-rank block (zigzag-selection variants) converge onto the observed stable rank and the loss line of the rank-48 LoRA reference in lock-step. Among these orthonormal bases, the DCT is preferred because its separable fast transform admits a fused reconstruction kernel: the materialized weight matrix never leaves on-chip memory, so the parameter saving translates into a bandwidth saving as well. 1 Introduction The cost of training a transformer is dominated by its weight matrices, the majority of them in the linear projections of attention and feed-forward sub-layers (Vaswani et al., 2017). Two families of methods reduce that cost during training. Low-rank parameterization writes W=ABW=AB with A∈ℝm×rA ^m× r, B∈ℝr×nB ^r× n, capping the rank of W at r (Hu et al., 2021). Sparse-coefficient parameterization writes W=ΦcW= c for a fixed orthonormal basis Φ and a K-sparse coefficient vector c: only K of the mnmn basis coefficients are trainable, the remainder held at zero. Spectral fine-tuning methods (Gao et al., 2024; Shen and others, 2025) instantiate the second family on top of a pre-trained checkpoint, and earlier work used compressed Fourier weights for evolutionary training of small feedforward networks (Koutník et al., 2010). The transformer pretraining setting and a matched-parameter-count comparison against a low-rank baseline are missing from the existing record. We provide both. A 4-layer, 128-dim transformer trained on tinyshakespeare with each weight matrix parameterized as K=mn/2K=mn/2 two-dimensional DCT coefficients reaches validation loss 1.6041.604 against 1.5801.580 for a matched dense baseline—within the terminal-epoch variation of the dense run, at half the trainable parameter count. A LoRA factorization at rank 48—the rank that matches the spectral layer’s trainable-parameter count for our block shapes—reaches only 1.8011.801, +0.22+0.22 from dense and +0.20+0.20 from the DCT layer. Contributions. 1. From-scratch DCT pretraining matches a dense baseline at 2×2× compression, while a matched-parameter-count LoRA factorization trails by +0.22+0.22. The structural advantage of sparse-coefficient over low-rank parameterizations at matched K is qualitative rather than marginal. 2. The mechanism is rank flexibility. A random orthonormal basis matches the DCT at K=mn/2K=mn/2, and a compression sweep through K∈mn/10,mn/20K∈\mn/10,mn/20\ tracks both loss and stable rank: sparse-coefficient variants whose effective subspace can host high-rank matrices keep their stable rank flat at ∼42 42 and remain the best sparse cells at every compression, while zigzag-selection variants converge onto the observed stable rank and the loss line of the rank-48 LoRA reference within noise. 3. The DCT is the kernel-friendly member of the equivalence class. A fused sparse-IDCT kernel can reconstruct W inside on-chip memory and feed the downstream matmul without a round-trip through main memory, giving the DCT layer the statistical properties of any orthonormal basis and the kernel properties of a classical fast transform. 2 Method Fix an orthonormal basis Φ∈ℝmn×mn ^mn× mn for the space of m×nm× n weight matrices and a selection set S⊂1,…,mnS⊂\1,…,mn\ of size |S|=K|S|=K. A sparse-coefficient layer parameterizes its weight matrix as W=unvec(ΦembedS(c)),W=unvec\! ( \,embed_S(c) ), (1) where c∈ℝKc ^K is the trainable coefficient vector and embedSembed_S inserts c at the K selected positions of an mnmn-vector whose other entries are zero. Both Φ and S are fixed before training. The primary method uses Φ=ΦDCT = _DCT (the orthonormal 2D type-I DCT) with S=SzigzagS=S_zigzag (the K lowest-frequency indices under a zigzag scan of the 2D frequency grid, as in JPEG). Three further (Φ,S)( ,S) choices serve as ablations: random orthonormal basis with zigzag selection, DCT basis with a fixed random subset, and random basis with a fixed random subset. A rank-r LoRA layer W=ABW=AB is outside the family: it is a bilinear map whose image is the rank-≤r≤ r variety, an algebraic subvariety of ℝm×nR^m× n of dimension r(m+n−r)r(m+n-r) with empty interior for r<min(m,n)r< (m,n). The forward pass reconstructs W via Eq. 1 at O(mnlogmn)O(mn mn) cost using a separable fast DCT and applies the usual matrix-vector product; backpropagation through the linear reconstruction propagates gradients from ∂ℒ/∂W /∂ W to ∂ℒ/∂c /∂ c by projecting onto the selected basis columns. Coefficients are initialized i.i.d. Gaussian with σ=2/n⋅mn/Kσ= 2/n· mn/K, so that the reconstructed weight matrix has Kaiming variance. 3 Experiments 3.1 Setup We train a character-level language model on tinyshakespeare (Karpathy, 2015) (≈1≈ 1M characters, 90/1090/10 train/val). Architecture: 44 layers, 128128-dim embeddings, 44 attention heads, context length 128128; all linear layers inside the transformer blocks (QKV, attention output, two MLP layers) are parameterized via the method under test, while the embeddings and LM head are dense in every configuration. Training: 3030 epochs of 200200 SGD steps, AdamW with weight decay 0.010.01 and cosine LR schedule, batch size 3232, gradient clipping at 1.01.0. Learning rate 3×10−43× 10^-4 for dense and LoRA, 1×10−31× 10^-3 for the sparse-coefficient variants. Every variant trains on identical batches in identical order under a fixed seed. 3.2 Main result and basis ablation (K=mn/2K=mn/2) Table 1 reports the matched-parameter comparison and the 2×22× 2 basis × selection ablation. All sparse-coefficient and matched-LoRA rows hold 424,832424,832 trainable parameters (52%52\% of dense). Table 1: Character-level language modelling on tinyshakespeare at K=mn/2K=mn/2. Validation loss is the mean cross-entropy over 5050 held-out batches at the final epoch. Method Params % of dense Val loss Δ vs. dense standard 818,048 100% 1.580 +0.000+0.000 dct_zigzag (primary) 424,832 52% 1.604 +0.024+0.024 dct_random 424,832 52% 1.616 +0.036+0.036 rand_zigzag 424,832 52% 1.584 +0.004+0.004 rand_random 424,832 52% 1.593 +0.013+0.013 lora_r48 (matched K) 424,832 52% 1.801 +0.221+0.221 The DCT layer lands within 0.0240.024 of the dense baseline; the last five epochs of the dense run vary between 1.5761.576 and 1.5951.595, placing the gap inside the terminal-epoch variation of the dense baseline itself. The rank-48 LoRA layer at the same trainable parameter count is 0.220.22 behind under an identical training protocol. The four cells of the 2×22× 2 ablation lie within 0.040.04 of dense. A random orthonormal basis with zigzag selection reaches 1.5841.584, slightly ahead of the DCT cell at 1.6041.604; a fixed random subset of the basis (replacing the zigzag ordering) costs ∼0.01 0.01 in either basis. At K=mn/2K=mn/2, any orthonormal basis with a generic selection works; the DCT basis is one member of a statistical equivalence class. 3.3 High-compression sweep: K=mn/10K=mn/10 and K=mn/20K=mn/20 Table 2 pushes compression by retraining every variant at K=mn/10K=mn/10 (each layer stores 10%10\% of its dense parameter count) and K=mn/20K=mn/20 (5%5\% per layer), with the same architecture and training protocol. Table 2: High-compression sweep. The lora_r48 row is a rank-48 LoRA trained under the same protocol; its parameter count and loss are independent of the sparse-coefficient K, so at these compressions it holds ≥4×≥ 4× the trainable budget of the spectral cells and is included as a fixed reference line only. K=mn/10K=mn/10 K=mn/20K=mn/20 Method Params Val loss Δ Params Val loss Δ standard 818,048 1.580 +0.000+0.000 818,048 1.580 +0.000+0.000 dct_zigzag (primary) 110,260 2.048 +0.468+0.468 70,940 2.198 +0.618+0.618 dct_random 110,260 1.827 +0.247+0.247 70,940 1.950 +0.370+0.370 rand_zigzag 110,260 1.989 +0.409+0.409 70,940 2.052 +0.472+0.472 rand_random 110,260 1.837 +0.257+0.257 70,940 1.954 +0.374+0.374 lora_r48 (ref.) 424,832 1.801 +0.221+0.221 424,832 1.802 +0.222+0.222 At both compression points the four sparse-coefficient cells rank dct_random≈rand_random≪rand_zigzag<dct_zigzag, dct\_random≈ rand\_random\; \; rand\_zigzag\;<\; dct\_zigzag, a stable ordering across the two compressions and a clean reordering from K=mn/2K=mn/2. Random selection beats zigzag by 0.150.15 to 0.250.25 within either basis, and the gap grows with compression. Under random selection the DCT and random bases stay within 0.010.01 of each other at both points (1.8271.827 vs 1.8371.837 at r=10r=10; 1.9501.950 vs 1.9541.954 at r=20r=20). The sparse-coefficient family does not match dense at either high-compression point: the best sparse cells are +0.247+0.247 above dense at K=mn/10K=mn/10 and +0.370+0.370 above at K=mn/20K=mn/20. The crossover where parity with dense is lost sits somewhere in (mn/10,mn/2)(mn/10,mn/2); we do not resolve it here. The basis-agnostic behaviour at K=mn/2K=mn/2 does, however, survive through both high-compression points under random selection: the selection axis, not the basis axis, breaks first. 4 Mechanism: Rank Flexibility The basis ablation establishes that any orthonormal basis suffices at K=mn/2K=mn/2. The mechanism is therefore a structural property of K-sparse linear subspaces rather than of the DCT specifically, and we identify it as rank flexibility. Write the weight matrix of a sparse-coefficient layer as W=unvec(ΦSc)W=unvec( _Sc), where ΦS∈ℝmn×K _S ^mn× K is the orthonormal basis restricted to its K selected columns. The image of this parameterization is a K-dimensional linear subspace of ℝm×nR^m× n. A generic linear subspace of dimension K in matrix space contains matrices of every rank from 11 to min(m,n) (m,n): full-rank matrices are an open dense subset of ℝm×nR^m× n, so a generic K-dimensional subspace intersects them in an open K-dimensional piece. Sparse-coefficient parameterizations therefore constrain the K-dimensional linear subspace in which W lives, but place no bound on its rank. A LoRA layer is bilinear: its image is the rank-≤r≤ r variety, with empty interior for r<min(m,n)r< (m,n), and every W it can represent has rank at most r. Table 3 reports the stable rank ‖W‖F2/‖W‖22\|W\|_F^2/\|W\|_2^2 of the materialized W on the attention QKV projection, averaged over the four layers, for each parameterization at each compression. The pattern on the other three layer classes (attention output, MLP1, MLP2) is the same. Table 3: Stable rank of the materialized W on the attention QKV projection at the end of training. min(m,n)=128 (m,n)=128; the rank-4848 LoRA reference has a structural rank ceiling of 4848. The standard and lora_r48 columns are constants because neither parameterization depends on K. K standard dct_z dct_rnd rand_z rand_rnd lora_r48 mn/2mn/2 8.4 27.5 39.5 29.8 38.5 14.8 mn/10mn/10 8.4 18.8 47.3 18.3 46.5 14.8 mn/20mn/20 8.4 14.2 42.0 13.6 42.7 14.8 Figure 1: Validation loss (a) and stable rank of the materialized W at the QKV projection (b) across the compression sweep K∈mn/2,mn/10,mn/20K∈\mn/2,mn/10,mn/20\. The two panels mirror each other: random-selection variants (dashed) keep stable rank near 4242 and stay the lowest-loss sparse cells; zigzag variants (solid) drop to ∼14 14 at K=mn/20K=mn/20, where they meet the lora_r48 reference (red dotted) in both panels. Standard dense SGD collapses the QKV projection to stable rank 8.48.4, an instance of the intrinsic-dimensionality phenomenon (Aghajanyan et al., 2020). At K=mn/2K=mn/2, the sparse-coefficient cells reach the same loss at three to five times that stable rank (2424–5454 across the four layer classes; Table 3 shows the QKV column and the other three classes follow the same pattern): they land in different, higher-rank solutions in weight space, and those solutions are equally good. The compression sweep turns observation into prediction. Random-selection cells retain stable rank ∼42 42 across the sweep and remain the best sparse cells at every compression. Zigzag-selection cells collapse from stable rank ∼28 28 at K=mn/2K=mn/2 down to ∼14 14 at K=mn/20K=mn/20—the same stable rank that the rank-48 LoRA reference holds throughout the sweep—and their loss converges onto the rank-48 LoRA loss line within noise. Variants that reduce to LoRA’s rank reduce to LoRA’s loss. The reading consistent with our data is that the loss landscape is low-rank friendly rather than low-rank required: SGD on unconstrained parameters prefers the low-rank basin, the sparse-coefficient variants find equally good higher-rank basins, and a hard rank ceiling at r=48r=48 sits in neither cleanly. 5 Why DCT Specifically: A Fused Reconstruction Kernel Among the orthonormal bases that match dense pretraining at K=mn/2K=mn/2, the DCT is the one whose reconstruction is cheap: the 2D type-I DCT admits an O(mnlogmn)O(mn mn) separable fast transform, while a generic orthonormal basis requires an O(m2n2)O(m^2n^2) dense matrix-vector product that dominates the downstream matmul. A naïve DCT-layer forward pass scatters c into an m×nm× n tile, runs the inverse DCT to materialize W, and runs a standard matmul against the activations. The reconstructed weight matrix lives in main memory for the duration of the forward pass, which throws away the parameter saving that the sparse coefficients were meant to deliver: the bandwidth cost of the forward pass is still dominated by moving mnmn weights in and out of main memory, exactly as in a standard dense layer. A fused GPU compute shader removes that round-trip by folding the IDCT and the matmul into a single pass: it loads the K coefficients into on-chip shared memory, runs the row transform into registers, consumes each reconstructed row immediately via FMA against the downstream activations before the next row is produced, and then applies the column transform in a second register-local pass. The materialized W never leaves the on-chip memory hierarchy; off-chip traffic per layer drops from O(mn+BT(m+n))O(mn+BT(m+n)) bytes to O(K+BT(m+n))O(K+BT(m+n)) bytes. The separable structure of the 2D DCT is essential—a dense arbitrary-basis reconstruction would need O(mn⋅K)O(mn· K) arithmetic per layer, an order of magnitude more work than the matmul itself. We prototype the fused kernel on Apple Silicon using Metal, with a 3232 KiB per-threadgroup shared-memory budget; the same fusion pattern extends to GPUs with larger on-chip memories, which only relaxes the upper bound on the m,nm,n per block the kernel can handle in one pass. Standalone batched DCT kernels on the prototype reach ≈120≈ 120 GFLOPs at N=4096N=4096, and the sparse IDCT is strictly easier than the dense case because most of the input is zero. A measured benchmark of the fused kernel inside the pretraining loop is left to follow-on work. 6 Related Work Spectral fine-tuning of transformers. FourierFT (Gao et al., 2024) and sDCTFT (Shen and others, 2025) parameterize the LoRA-style update of a pre-trained transformer as a sparse spectral correction with far fewer trainable parameters than LoRA. Both start from a pre-trained checkpoint; we train from scratch and run the matched-K LoRA comparison that neither paper includes. The rank-ceiling diagnosis for LoRA. RandLoRA (Albert et al., 2025) parameterizes a fine-tuning update as a sum of fixed random low-rank bases scaled by learned coefficients, achieves full-rank updates at the parameter count of LoRA, and concludes that rank, not parameter count, is the bottleneck of LoRA fine-tuning. Shuttleworth et al. (2024) reach the same conclusion from an “intruder dimensions” diagnostic on LoRA fine-tuning. Both works study fine-tuning of a pretrained model; our contribution is the from-scratch counterpart, with the additional structured-vs.-random basis ablation at matched K. Random subspaces, low-rank pretraining, structured matrices. Li et al. (2018) trained a network in a fixed random low-dimensional subspace of its parameter space to measure the intrinsic dimension of the objective landscape; our rand_zigzag cell is a per-weight-matrix instance of the same construction. GaLore (Zhao et al., 2024) projects weight updates onto a rank-r subspace during pretraining. Monarch (Dao et al., 2022) demonstrates from-scratch GPT-2 pretraining with block-diagonal-structured sparse matrices; a matched-parameter comparison against Monarch is a natural next step. 7 Conclusion A 4-layer, 128-dim transformer trained from scratch with each weight matrix parameterized as K=mn/2K=mn/2 DCT coefficients matches a dense baseline within the dense run’s terminal-epoch variation, while a rank-48 LoRA factorization at the same trainable parameter count is +0.22+0.22 behind. The mechanism is rank flexibility: a generic K-sparse orthonormal subspace at K=mn/2K=mn/2 contains matrices of every rank up to min(m,n) (m,n), and a compression sweep through K∈mn/10,mn/20K∈\mn/10,mn/20\ shows that variants which collapse this property converge onto the rank-48 LoRA stable rank and the rank-48 LoRA loss line in lock-step. The DCT is preferred among the orthonormal bases in the equivalence class because its separable fast transform admits a fused reconstruction kernel that keeps the materialized weight matrix inside on-chip memory. Whether matched-K parity holds at the scale of modern language models, whether the rank-flexibility explanation carries to the multi-billion-parameter regime, and whether the fused kernel closes the wall-clock gap with a dense matmul in production are the three open questions this work leaves. A reference implementation of the batched FFT/DCT kernels is available at https://github.com/aminems/AppleSiliconFFT. References A. Aghajanyan, L. Zettlemoyer, and S. Gupta (2020) Intrinsic dimensionality explains the effectiveness of language model fine-tuning. arXiv preprint arXiv:2012.13255. Cited by: §4. P. Albert, F. Z. Zhang, H. S. Rodriguez, E. Abbasnejad, W. Buntine, and A. van den Hengel (2025) RandLoRA: full rank parameter-efficient fine-tuning of large models. In Proc. Int. Conf. Learning Representations (ICLR), Note: https://arxiv.org/abs/2502.00987 Cited by: §6. T. Dao, B. Chen, N. S. Sohoni, A. Desai, M. Poli, J. Grogan, A. Liu, A. Rao, A. Rudra, and C. Ré (2022) Monarch: expressive structured matrices for efficient and accurate training. In Proc. Int. Conf. Machine Learning (ICML), Cited by: §6. Z. Gao, Q. Qin, A. Li, L. Davis, J. Liang, and X. Chen (2024) Parameter-efficient fine-tuning with discrete Fourier transform. In Proc. Int. Conf. Machine Learning (ICML), Cited by: §1, §6. E. J. Hu, Y. Shen, P. Wallis, Z. Allen-Zhu, Y. Li, S. Wang, L. Wang, and W. Chen (2021) LoRA: low-rank adaptation of large language models. arXiv preprint arXiv:2106.09685. Cited by: §1. A. Karpathy (2015) The unreasonable effectiveness of recurrent neural networks. Note: http://karpathy.github.io/2015/05/21/rnn-effectiveness/ Cited by: §3.1. J. Koutník, F. Gomez, and J. Schmidhuber (2010) Evolving neural networks in compressed weight space. In Proc. Genetic and Evolutionary Computation Conf. (GECCO), Cited by: §1. C. Li, H. Farkhoor, R. Liu, and J. Yosinski (2018) Measuring the intrinsic dimension of objective landscapes. In Proc. Int. Conf. Learning Representations (ICLR), Cited by: §6. Y. Shen et al. (2025) Parameter-efficient fine-tuning via selective discrete cosine transform. In Proc. Assoc. Computational Linguistics (ACL), Cited by: §1, §6. R. Shuttleworth, J. Andreas, A. Torralba, and P. Sharma (2024) LoRA vs full fine-tuning: an illusion of equivalence. arXiv preprint arXiv:2410.21228. Cited by: §6. A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, Ł. Kaiser, and I. Polosukhin (2017) Attention is all you need. Advances in Neural Information Processing Systems 30. Cited by: §1. J. Zhao, Z. Zhang, B. Chen, Z. Wang, A. Anandkumar, and Y. Tian (2024) GaLore: memory-efficient LLM training by gradient low-rank projection. In Proc. Int. Conf. Machine Learning (ICML), Cited by: §6.