PufferAI
Create PMLL.py
Overview of the PMLL Compression Algorithm
The Persistent Memory Logic Loop (PMLL) architecture introduces a novel approach to memory-efficient inference in large language models (LLMs) by augmenting standard Transformers with an external, compressed persistent memory pool. A key innovation is the Recursive Memory Compression Algorithm, which dynamically reduces the memory footprint of this pool while minimizing accuracy loss. This algorithm achieves 59–60% memory reduction with less than 1.5% degradation in model performance, as validated on benchmarks like WikiText-2 and OpenWebText.
The algorithm is "recursive" because it iteratively applies compression in a hierarchical manner across multiple levels of the memory pool, re-evaluating and refining the compression until utilization targets are met. It combines importance scoring (to prioritize data), thresholding (for pruning), quantization (for precision reduction), and a feedback loop for recursion. This is triggered by the PMLL Memory Controller when pool utilization exceeds a threshold (e.g., 80%).
Below, I'll break it down step by step, including key equations and pseudocode from the PMLL architecture description.
1. Importance Scoring Function
Each entry in the persistent memory pool (e.g., key-value pairs from KV caches or embeddings) is assigned an importance score ( s_i ) to gauge its utility. This score balances multiple factors reflecting the entry's relevance and usage patterns:
[ s_i = \alpha_1 \cdot \text{recency}(i) + \alpha_2 \cdot \text{frequency}(i) + \alpha_3 \cdot \text{semantic_value}(i) ]
- Recency(( i )): A decay function (e.g., exponential) based on time since last access, favoring recent data.
- Frequency(( i )): Cumulative access count, emphasizing frequently retrieved entries.
- Semantic Value(( i )): Derived from contextual similarity (e.g., cosine similarity to current queries) or external validation (e.g., knowledge graphs).
The weights ( \alpha_1, \alpha_2, \alpha_3 ) are tunable hyperparameters, often learned via fine-tuning or set empirically (e.g., ( \alpha_1 = 0.4 ) for recency-heavy tasks like real-time chat). Scores are computed in a vectorized manner using SIMD intrinsics in the C backend for efficiency.
This step ensures that critical, high-utility data (e.g., core factual knowledge) is protected, while redundant or outdated entries are deprioritized.
2. Thresholding for Pruning
With scores computed for all ( n ) entries, a pruning threshold ( \tau ) is determined to decide which entries to retain:
[ \tau = \text{quantile}({s_i}_{i=1}^n, \rho) ]
- ( \rho ): The compression ratio (e.g., 0.1–0.25), representing the fraction of top-scored entries to keep uncompressed.
- The quantile operation sorts scores and selects the value at the ( \rho )-th percentile.
Entries with ( s_i < \tau ) are pruned (discarded or archived), while those above are candidates for lighter compression. This step alone can eliminate 70–80% of low-value data, directly tying into PMLL's promise queue semantics—pruned entries' "promises" (deferred operations) are resolved or expired based on TTL (time-to-live).
3. Quantization Process
Retained entries are further compressed via adaptive vector quantization, where the bit precision ( q ) is scaled by importance:
[ q = \begin{cases} 8 & \text{if } s_i > 0.8 \cdot \max({s_j}) \ 4 & \text{if } 0.4 \cdot \max({s_j}) \leq s_i \leq 0.8 \cdot \max({s_j}) \ \text{discard} & \text{otherwise (fallback to pruning)} \end{cases} ]
- For a vector entry ( v ) (e.g., a float32 embedding), quantization maps it to a lower-bit representation: [ v_q = \round\left( \frac{v - \min(v)}{\max(v) - \min(v)} \cdot (2^q - 1) \right) \cdot \frac{\max(v) - \min(v)}{2^q - 1} + \min(v) ] followed by casting to int8/int4, halving or quartering storage needs.
Higher ( q ) preserves fidelity for important data (e.g., float16 equivalent), while lower ( q ) aggressively compresses peripherals. Dequantization occurs on-the-fly during retrieval, with negligible latency due to C-optimized routines.
4. Recursion Mechanism
To handle varying loads, the algorithm recurses across a hierarchical memory structure (e.g., Level 0: uncompressed; Level 1: quantized; Level 2: pruned + archived). After one pass:
- The updated pool is re-scored and re-thresholded.
- Entries may "demote" to deeper levels (more compression) if their scores drop.
- Recursion halts when utilization < target (e.g., 60%) or max depth (e.g., 3 levels) is reached.
This creates a self-adaptive loop, integrated with PMLL's attention mechanism: during hybrid attention (local + persistent), dequantized entries blend seamlessly, with a blending factor ( \alpha ) computed via similarity norms.
Theoretical bounds ensure convergence: Accuracy loss ( \Delta L \leq C \rho^{\lambda - 1} ) (where ( \lambda > 1 ) from power-law score distributions), preventing over-compression.
Pseudocode
Here's the core algorithm in pseudocode (adapted from PMLL's Algorithm 2):
Algorithm: Recursive Memory Compression
Input: Memory pool M, ratio ρ, max_levels L
Output: Compressed pool M'
1: level ← 0
2: while level < L and utilization(M) > target:
3: Compute scores: {s_i} ← importance_scores(M) // Vectorized via C
4: τ ← quantile({s_i}, ρ)
5: M' ← empty pool
6: for each entry e_i in M:
7: if s_i ≥ τ:
8: q ← select_bits(s_i) // e.g., 8/4 based on score
9: e'_i ← quantize(e_i, q)
10: M' ← M' ∪ {e'_i}
11: end if
12: end for
13: M ← M' // Update pool
14: level ← level + 1
15: end while
16: Update metadata (e.g., dequantization flags)
17: return M
Integration with PMLL Architecture
In PMLL, compression runs asynchronously via the Promise Queue: Writes to persistent memory enqueue "promises" with initial scores, processed in batches. The Memory Controller (Python-orchestrated with C calls) triggers it on high utilization, syncing with Transformer forward passes. For example, in pml_attention, retrieved persistent KV pairs are dequantized before blending with local cache.
This yields KV cache savings of 60–62% for long sequences, enabling deployment on edge devices. Limitations include score computation overhead (mitigated by caching) and potential drift in extreme recursions, addressed via periodic full recompression.
For implementation details, see the PMLL paper's C extensions for SIMD-accelerated scoring and quantization.
So the one thing with time… this file did take time to cook. Give it grace with how raw it is.