RetainCode: One-Bit Exact Gated-MLP Activation Compression from Retained Products

Abstract

Gated Transformer MLPs save two expanded projection branches for backward, although the following down projection separately retains their rounded elementwise product. We turn this product into side information. RetainCode keeps the gate branch and replaces the BF16 up branch with a one-bit disambiguator for a divide-based predictor. If the bit is set, exactly one adjacent BF16 value remultiplies to the retained product. Non-normal paths use a fixed raw buffer whose positions are inferred from retained values. A single post-forward guard retries the ordinary fused graph before backward if that buffer overflows. We prove that normal BF16 multiplication followed by FP32 division can move the recovered multiplicand by at most one adjacent BF16 value and that a retained product has at most two normal preimages. Thus one bit is sufficient. An explicit collision proves it is necessary. Exhaustive enumeration verifies all 3.17 billion normal finite pairs. We also trace retention through dropout and obtain a three-bit exact GeGLU code for FLAN-T5. Against fused Liger kernels, RetainCode reduces saved activations by 12.70% on OLMo2-1B, 7.48% on Qwen3-0.6B, and 13.50% on Qwen3-8B. It lowers peak memory by 7.42% and 4.16% on the two smaller dense models, and raises measured forward-backward capacity from 12,288 to 14,336 OLMo2 tokens and from 3,264 to 3,616 Qwen3-8B tokens on one 39.49 GiB A100. Same-checkpoint paired measurements give 0.982$\times$, 0.984$\times$, and 0.978$\times$ Liger throughput. On the 6.9B-parameter OLMoE-1B-7B, it codes all 1,024 expert MLPs and saves 7.09% of activations at 0.957$\times$ throughput. FLAN-T5 saves 2.73% of activations and 2.6–2.7% of peak memory at 0.959$\times$ throughput. Two optimizer-inclusive four-A100 OLMo2 runs reduce per-rank peak memory by 6.36% with DDP and ZeRO-sharded AdamW, and by 6.83% with fully sharded AdamW. The DDP configuration retains 0.975$\times$ throughput. With fused linear cross entropy, the fully sharded reduction rises to 8.88% at 0.980$\times$ throughput and increases the largest tested fitting global-token step from 70,656 to 77,824. A same-checkpoint frontier shows that the codec, partial MLP checkpointing, and their hybrids occupy distinct memory–throughput points. Every reported single-GPU loss matches bitwise, all performance-run overflow counters are zero, and complete pretrained audits across all five checkpoints match all 17.44 billion parameter-gradient elements.

Main results

Theorem. Let $s$ and $b$ be normal finite nonzero precision-$p$ values, and let $c=\operatorname{round}_p(sb)$ also be normal finite and nonzero. Compute $q=c/s$ exactly and let $\widehat b=\operatorname{round}_p(q)$. Then $\widehat b$ is $b$ or one of its two adjacent representable values. For fixed $s,c$, at most two normal values of $b$ round to $c$. Consequently, if $b\ne\widehat b$, exactly one normal neighbor of $\widehat b$ is the true $b$, so a single bit suffices to identify it.

Keywords

activation compression, gated MLP, memory-efficient training, one-bit encoding, exact gradients, floating-point, SwiGLU

Files

• LICENSE
• SUPPLEMENT_README.md
• fancyhdr.sty
• iclr2027_conference.bst
• iclr2027_conference.sty
• main.bbl
• main.pdf
• main.tex
• math_commands.tex
• method_figure.tex
• natbib.sty
• references.bib
• requirements.txt
• supplementary.zip