Low-precision optimizer-state methods are commonly designed and evaluated for dense Adam-style first and second moments. Memory-efficient optimizers depart from this setting: Adafactor factorizes second moments, CAME adds factored confidence states, and APOLLO maintains statistics in a projected gradient space. Consequently, an equal amount of state reconstruction error can induce different update errors depending on state topology and update semantics. We first characterize this heterogeneity in optimizer-state traces from language model pre-training. We then introduce Adaptive Log-Space (AL) quantization, a block-wise representation for non-negative states that adapts its nonzero range per block and enforces the exact-zero invariant $q = 0 \Leftrightarrow x = 0$. AL8 and AL16 are combined with independent signed-momentum encodings and state-specific precision choices rather than a single policy for every state. Across 96 runs totaling 214.7 GPU-hours, we evaluate dense, factored, confidence, and projected states in AdamW, Adafactor, CAME, and APOLLO paths. On a 20K-step TinyLlama-1.1B pre-training benchmark, an AdamW configuration with AL8 second moments and uniform 8-bit momentum reaches 72.90 perplexity, compared with 72.48 for FP32 AdamW and 73.54 for an 8-bit dynamic-quantization baseline, while reducing measured optimizer-state storage from 8392.7 to 2119.2 MiB. CAME exposes a different precision regime: promoting its non-negative states to AL16 recovers 86.16 perplexity versus 86.68 for the full-precision reference, whereas all-AL8 reaches 90.19. A 100K-step GPT-2 experiment further shows that topology-aware parameter protection reduces the late-loss gap of quantized Adafactor from +0.1185 to +0.0159 in the evaluated setup. These results support a state- and topology-aware view of optimizer quantization.
Scaling modern large language models (LLMs) to long contexts is limited by the quadratic computation cost, and poor length extrapolation of dense attention. Chunk-wise sparse attention offers a promising alternative, but all existing methods fall short of full attention because of their inaccurate chunk selection. We propose Hierarchical Landmark Sparse (HiLS) Attention, a chunk-wise sparse attention mechanism that learns chunk selection end-to-end under the language-modeling (LM) loss. HiLS factorizes attention hierarchically: each query performs attention independently with each retrieved chunk to extract chunk-specific information, and the resulting outputs are fused according to chunk retrieval scores. By incorporating retrieval scores into the forward attention computation, HiLS optimizes them directly with the LM loss, enabling end-to-end retrieval learning and native sparse training. Experimental results show that HiLS-Attention achieves performance comparable to, and in some cases better than, full attention at in-domain context lengths. Meanwhile, HiLS-Attention extrapolates more than $64\times$ the training context length with 90% retrieval accuracy, far beyond full attention. Moreover, existing full-attention models can be converted to HiLS-Attention with lightweight continued pretraining, preserving in-domain performance while acquiring ultra-long-context extrapolation. Together with its sparse KV access and computation, HiLS-Attention breaks the usual efficiency-performance trade-off, enabling long-context LLMs that are both more efficient and more effective on general long-context tasks than their full-attention counterparts.
This work proposes COMPACT-VA, a planning-aligned working memory framework built on conditional VQ-VAE, compressing extended context into bounded representations, and evaluates on high-signal dynamic scenarios where historical context is most critical for behavior correctness, and accordingly design behavioral metrics.
Zhixuan Liang, Yuxiao Chen, Yurong You et al.· arXiv.org· 1 citation