Research map / Architectures and efficiency / RNNs and normalization
Memory-Efficient Training with In-Place FFT Implementation
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes rdFFT, a real-domain fully in-place FFT/IFFT that packs conjugate-symmetric spectra into n real slots, and applies it to circulant-structured fine-tuning to cut training memory.
Problem
Standard FFT and rFFT cannot run truly in-place on real tensors because rFFT maps n reals to n/2+1 complex values (n+2 reals), causing size mismatch, extra allocations, and no bfloat16 support in libraries, which hurts memory-efficient training with FFT-based operators.
Approach
Exploit conjugate symmetry and butterfly structure so each complex coefficient's real part sits at index k and imaginary part at index r-k, letting the whole FFT and IFFT run in place within an n-element real buffer, with bf16 support, and use it inside circulant-matrix fine-tuning.
Key results
- Single-layer training with D=4096, B=1: ours with p=128 uses 1.06 MB peak versus 144.33 MB for full fine-tuning and 3.14 MB for rfft (p=128).
- On LLaMA2-7B, ours_p=512 reaches 17.93 GB total peak memory versus 19.28 GB for rfft_p=512 and 26.90 GB for full fine-tuning.
- On MRPC with RoBERTa-large, ours_p=1024 reaches 90.3% accuracy versus 90.9% for full fine-tuning and 89.7% for rfft_p=1024, but throughput is lower (0.93k vs 1.77k tokens/sec for rfft).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Replace rFFT with in-place real-domain FFT | Circulant fine-tuning of LLaMA2-7B · GSM8K | peak GPU memory (GB): 19.28 -> 17.93 | Helped |
| Replace rFFT with in-place real-domain FFT | Circulant fine-tuning of RoBERTa-large · MRPC | peak GPU memory (GB): 6.13 -> 5.77 | Helped |
| Replace rFFT with in-place real-domain FFT | Circulant fine-tuning of LLaMA2-7B · GSM8K | throughput (k tokens/sec): 1.77 -> 0.92 | Hurt |
| Replace rFFT with in-place real-domain FFT | RoBERTa-large fine-tuning · MRPC | accuracy (%): 90.7 -> 90.0 (p=512); 89.7 -> 90.3 (p=1024) | Mixed |
| Replace rFFT with in-place real-domain FFT | Standalone FFT operator · FFT p=512 forward | runtime (ms): 0.0195 -> 0.0279 | Hurt |
| Replace rFFT with in-place real-domain IFFT | Standalone inverse FFT operator · FFT p=512 inverse | runtime (ms): 0.0450 -> 0.0233 | Helped |
| Replace rFFT with in-place real-domain FFT | Standalone FFT operator · FFT p=4096 forward | runtime (ms): 0.0199 -> 0.0687 | Hurt |
| Vary circulant block size from 128 to 4096 (ablation) | rdFFT single-layer training · D=4096, B=1 | peak GPU memory (MB): 1.06 -> 0.09 | Helped |
| Vary circulant block size from 128 to 4096 (ablation) | rdFFT single-layer training · D=4096, B=256 | peak GPU memory (MB): 20.50 -> 20.02 | No clear effect |
| Use in-place real-domain encoding instead of standard FFT | Single-layer circulant training · D=4096, B=1, block size 128 | peak GPU memory (MB): 3.65 -> 1.06 | Helped |
| Use in-place real-domain encoding instead of rFFT | Single-layer circulant training · D=4096, B=1, block size 128 | peak GPU memory (MB): 3.14 -> 1.06 | Helped |
Surprising findings
- The proposed in-place method has roughly half the throughput of rfft-based circulant training (0.92-0.93k vs 1.77k tokens/sec on LLaMA-2-7B).
- At large batch size (B=256) with D=4096, the in-place method uses about 20 MB versus 547 MB for rfft at p=128, but the gap narrows for some configurations, and at p=4096 the runtime overhead grows due to CUDA thread-block synchronization.
Limitations
- Frequency-domain information is encoded implicitly, so explicit complex spectrum access (e.g., spectral filtering) requires casting that breaks in-place memory.
- Throughput is lower than rfft and fft circulant baselines; runtime overhead grows at large block sizes due to synchronization.
- Evaluation is limited to RoBERTa-large on MRPC and LLaMA2-7B on GSM8K with few configurations.
Experimental setup
Datasets: GSM8K MRPC
Benchmarks: MRPC GSM8K
Metrics: peak GPU memory (MB/GB) throughput (k tokens/sec) accuracy (%) operator runtime (ms) absolute and relative numerical error
Baselines: full fine-tuning LoRA torch.fft.fft circulant torch.fft.rfft circulant block circulant adapter
Builds on: Cooley-Tukey FFT block circulant adapter LoRA FourierFT
Closest papers by meaning
- Parameter-Efficient Fine-Tuning with Discrete Fourier Transform (ICML 2024 · similarity 0.85)
- FlashFFTConv: Efficient Convolutions for Long Sequences with Tensor Cores (ICLR 2024 · similarity 0.84)
- Learning Spatially Collaged Fourier Bases for Implicit Neural Representation (AAAI 2024 · similarity 0.84)
- S$^{2}$FT: Efficient, Scalable and Generalizable LLM Fine-tuning by Structured Sparsity (NeurIPS 2024 · similarity 0.84)
- Learning High-Frequency Functions Made Easy with Sinusoidal Positional Encoding (ICML 2024 · similarity 0.84)
- Sparse Spectral Training and Inference on Euclidean and Hyperbolic Neural Networks (ICML 2025 · similarity 0.83)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.