Reproduce “Not All Bits Are Equal: Scale-Dependent Memory Optimization Strategies for Reasoning Models” on OpenResearch

About this paper

Junhyuck Kim, Ethan Ewer, Taehong Moon, Jongho Park, Dimitris Papailiopoulos

Researchers from KRAFTON, University of Wisconsin–Madison, UC Berkeley, and Microsoft Research empirically investigated memory optimization for reasoning models, discovering that optimal strategies for balancing model weights and KV cache compression are scale-dependent. Their findings indicate that 4-bit weight quantization is not universally optimal, particularly for mathematical reasoning, and highlight the critical role of KV cache compression in memory-constrained environments.

Key insights

• Optimal memory allocation between model weights and KV cache is scale-dependent: smaller models (effectively < 8-bit 4B) benefit more from increased weight memory, while larger models (effectively ≥ 8-bit 4B) benefit from longer generations.
• The most effective weight precision varies by task and model size; 4-bit quantization is memory-optimal for knowledge-intensive tasks, but mathematical reasoning often requires higher precision (8-bit or 16-bit), especially for smaller models.
• KV cache compression is crucial for reasoning models, consistently improving memory–accuracy trade-offs, with KV cache eviction being more effective for models smaller than an 8-bit 8B model, and KV cache quantization becoming competitive for larger models.

Summary by alphaXiv.