Hugging Face's logo

Join the conversation

Join the community of Machine Learners and AI enthusiasts.

Sign Up
Back to feed
Banaxi-Tech 
posted an update 2 days ago
Post
4105
Today we are releasing BananaMind-KV1-8M-2Bit-Experimental, a KV-cache-aware trained model that stores its generation KV cache in 2-bit precision instead of the usual 16-bit precision.

Result: 5.33x smaller KV cache vs FP16, with 0.0916 mean KLD against a 16-bit KV cache reference on WikiText-2.

Model: BananaMind/BananaMind-KV1-8M-2Bit-Experimental

The important part: this is not just post-training KV cache quantization.
Instead we take the BitNet approach.

KV1 is trained with a 2-bit-aware K/V path. Instead of training a normal model and quantizing the cache afterwards, the model learns during training to operate under the low-bit KV constraint, closer in spirit to the BitNet idea of training for the low-bit regime.

During generation, each K/V vector is quantized into 4 affine levels and packed into uint8 tensors, with four 2-bit values stored per byte.

WikiText-2 eval vs 16-bit KV cache reference:

Mean KLD: 0.0916 nats/token
Mean KLD: 0.1322 bits/token
Average KV cache shrink vs FP16: 5.33x
Evaluated positions: 372,675

If this actually gets used in models like Qwen or Gemma, then it may be possible to run 128K or even 256K Context on a Normal Machine!
Try it here: BananaMind/BananaMind-KV1-8M-2Bit-Experimental

Code: https://github.com/Banaxi-Tech/kv1
dipankarsarkar

5.33x shrink is the easy number. The one I'd chase is task-level, not KLD.

0.0916 nats/token is an average over 372k positions, and a KV cache lives or dies on the rare position that mattered: the token at 90k you have to attend back to. Average KLD can look tiny while long-context recall quietly falls off, because the easy majority of positions dominate the mean.

Before "128K on a normal machine," I'd want a needle-in-a-haystack sweep: same passkey at depth 8k / 32k / 128k, 2-bit KV vs 16-bit, pass/fail not perplexity. That's the eval that proves the claim.

Have you run anything retrieval-shaped over long context yet, or just WikiText so far?

·
Banaxi-Tech
•
edited 1 day ago

Great idea, I will try soon! But i am currently focusing on training my new Models.
dipankarsarkar

Since KV1 is training-aware not drop-in, you can't borrow a bigger model's competence for free. That's the real bind with going small.

So skip needle-in-haystack (it needs capability an 8M won't have) and use a long-range task an 8M CAN learn: induction/copy where the target token depends on one planted ~60 positions back. Train that, then 2-bit vs 16-bit K/V and measure recall of the planted token by depth.

Same size, isolates the cache under a real long-range dependency. KLD averages straight over that, a copy task can't hide it.

Would that fit your current training budget?
In this post