Update README.md

README.md

@@ -72,45 +72,11 @@ Stay tuned!
 
 Nandi-Mini-500M introduces several efficiency-focused architectural optimizations designed for compact yet capable language models.
 
-
-### Shared KV KV-Cache Memory Comparison
-
-The following comparison illustrates the KV-cache memory reduction enabled by Shared KV mode.
-
-```python
-import matplotlib.pyplot as plt
-
-modes = ["Vanilla KV", "Shared KV"]
-memory = [100, 50]
-
-plt.figure(figsize=(5,4))
-bars = plt.bar(modes, memory)
-
-plt.ylabel("Relative KV Cache Memory")
-plt.title("KV Cache Memory Usage")
-
-for bar, val in zip(bars, memory):
-    plt.text(
-        bar.get_x() + bar.get_width()/2,
-        val + 2,
-        f"{val}%",
-        ha='center'
-    )
-
-plt.ylim(0, 120)
-plt.show()
-```
-
-Expected result:
-
-- Vanilla KV → 100% KV-cache memory
-- Shared KV → ~50% KV-cache memory
-
-Shared KV trades a small increase in compute overhead for significantly lower memory usage, since RoPE and Key normalization are applied dynamically during attention computation.
+#### Shared KV (Shared Key-Value Vectors)
 
 Shared KV is one of the core architectural ideas explored in Nandi-Mini. Instead of storing separate Key and Value vectors, both share the same underlying representation, while a lightweight Key normalization step is applied specifically for attention computation.
 
-This design reduces KV-cache memory usage by ~50% during inference with only a small increase in compute overhead, since Key normalization and RoPE transformations must be applied dynamically during attention computation.
+This design reduces KV-cache memory usage by ~50% during inference with only a small increase in compute overhead, since RoPE and Key normalization are applied dynamically during attention computation.
 
 Nandi supports two KV cache modes:
 
@@ -126,7 +92,21 @@ Uses Shared KV, reducing KV-cache memory by ~50% with slightly higher compute ov
 
 Uses standard separate Key-Value caching for maximum inference compatibility and lower compute overhead.
 
-Shared KV is part of our broader focus on deployable foundation models optimized for edge devices, on-premise AI systems, low-latency enterprise inference, and efficient multilingual serving.
+### KV-Cache Memory Comparison
+
+<p align="center">
+  <img src="./assets/shared_kv_cache_comparison_improved.png" width="650"/>
+</p>
+
+- Vanilla KV → Standard KV-cache memory usage
+- Shared KV → ~50% lower KV-cache footprint
+
+Shared KV is part of our broader focus on deployable foundation models optimized for:
+
+- Edge devices
+- On-premise AI systems
+- Low-latency enterprise inference
+- Efficient multilingual serving
 
 This remains an active research area within the Nandi model family, and we plan to share deeper technical details in upcoming engineering blogs.