Fast and Exact: Asymptotically Linear KL-Optimal Frequency Normalization

1. Problem Statement

In entropy coding (Range Coding, ANS, rANS), real-valued probability distributions must be converted into integer frequencies that sum to a fixed target $M$ (often a power of two, $2^k$) to enable efficient bit-shift operations during encoding and decoding.

The core optimization problem is to find integer frequencies $m_a$ for each symbol $a$ such that:

1. $\sum m_a = M$
2. $m_a \ge 1$ for all symbols with non-zero empirical counts.
3. The Kullback-Leibler (KL) divergence between the empirical distribution $p$ and the quantized distribution $q$ ($q_a = m_a/M$) is minimized.

Minimizing $D(p \parallel q)$ is equivalent to minimizing the expected per-symbol redundancy (in nats or bits) incurred by using the integer model $q$ to code a source generated by $p$.

The Challenge:
Existing fast normalization methods used in production codecs (e.g., Giesen's cumulative rescaling, Bloom's geometric-mean rounding, Collet's FSE normalization) are heuristic. While they run in $O(r)$ or $O(r \log r)$ time, they do not guarantee the global KL optimum. They often rely on:

• Order-dependent rounding errors.
• One-directional adjustments (only increasing or only decreasing frequencies).
• Initialization heuristics that may land "below" the true optimum, making recovery impossible for decrement-only algorithms.

The paper addresses the need for an algorithm that is both provably exact (guaranteed to find the global KL minimum) and fast enough to sit on the "hot path" of an entropy coder.

---

2. Methodology & Theoretical Foundation

The paper establishes that the problem is a separable concave integer allocation problem. The global optimum is characterized by a marginal-exchange condition (Lemma 1):
A frequency vector $m$ is optimal if and only if the minimum cost of removing a unit from any symbol (decrement ticket $\Delta^-$) is greater than or equal to the maximum gain of adding a unit to any symbol (increment ticket $\Delta^+$).

$$\min_{a: m_a \ge 2} \Delta^-_a(m_a) \ge \max_{b} \Delta^+_b(m_b)$$

Where:

• $\Delta^+_a(j) = c_a \log(\frac{j+1}{j})$ (Gain of increasing count $j \to j+1$)
• $\Delta^-_a(j) = c_a \log(\frac{j}{j-1})$ (Loss of decreasing count $j \to j-1$)

The author proposes three exact algorithms driven by this marginal-ticket principle, trading off between simplicity, asymptotic complexity, and implementation overhead.

Algorithm 1: Bottom-Up Archetype

• Approach: Starts with the minimal feasible allocation ($m_a=1$ for all active symbols) and greedily adds $M-r$ units to the symbol offering the highest marginal gain ($\Delta^+$).
• Mechanism: Uses a max-heap to track the next available gain for each symbol.
• Complexity: $O(M \log r)$.
• Pros: Simplest proof of optimality; no envelope assumptions required.
• Cons: Too slow for practical use when $M \gg r$ (typical in entropy coding where $M=2^{14}$ to $2^{20}$).

Algorithm 2: Bidirectional Bloom Repair

• Approach: An exact correction of Bloom's heuristic. It starts with Bloom's geometric-mean initialization. If the sum is not $M$, it performs a bidirectional exchange.
• Mechanism:
  1. Adjusts the total to $M$ using a one-directional heap (if sum $>M$, decrement; if sum $<M$, increment).
  2. Crucial Step: Once the sum is $M$, it enters a "Phase 2" exchange loop. It repeatedly swaps a unit from the symbol with the lowest decrement cost to the symbol with the highest increment gain, provided the gain exceeds the cost.
• Complexity: $O(M \log r)$ in the worst case, but $O(r \log r)$ on "benign" inputs where the initialization is close to optimal.
• Significance: Fixes the flaw in Bloom's original algorithm where a one-directional correction fails if the initialization lands on the wrong side of the optimum for multiple coordinates simultaneously.

Algorithm 3: Linear Top-Down Window (The Main Contribution)

• Approach: A top-down selection method that avoids the $O(M)$ dependency.
• Mechanism:
  1. Bounding: Proves that the optimal frequency $m^*_a$ for any symbol lies within a tight window $[L_a, U_a]$ around the scaled ideal frequency $s_a = M \cdot c_a / N$. The window width is bounded by $O(r)$.
  2. Selection: Instead of starting from $M$ and decrementing, it starts from an upper bound $U_a$ and selects the $D$ cheapest decrement tickets within this window.
  3. Optimization: Uses a linear-time selection algorithm (like Quickselect) to find the $D$ smallest tickets without fully sorting them.
• Complexity: $O(r)$ (asymptotically linear in the alphabet size).
• Significance: This is the first algorithm to achieve exact KL optimality with linear time complexity relative to the alphabet size, independent of the target frequency $M$.

---

3. Key Contributions

1. Exactness Proof: Provides the first set of algorithms that are provably globally KL-optimal, correcting the sub-optimality of widely used heuristics (Giesen, Bloom, Collet).
2. Asymptotically Linear Algorithm: Introduces the "Top-Down Window" method (Algorithm 3), which runs in $O(r)$ time, making it theoretically superior to heap-based methods ($O(M \log r)$) for large $M$.
3. Bidirectional Correction: Identifies and fixes the specific failure mode of Bloom's algorithm where one-directional repair cannot recover from multi-coordinate initialization errors.
4. Implementation & Validation:
   • Provides portable C implementations and AVX2 vectorized versions.
   • Validates correctness against specific "witness" counterexamples where existing heuristics fail.
   • Demonstrates that the exact algorithms match the speed of the fastest heuristics in practice.

---

4. Experimental Results

The authors benchmarked the algorithms on an AMD Ryzen 9 5900X across various distributions (Uniform, Geometric, Zipfian, Gaussian, Sparse) and alphabet sizes ($r$ up to 65,536).

• Correctness: All proposed algorithms passed strict marginal-exchange checks and matched the theoretical KL optimum. Existing heuristics (Giesen, FSE fast pass) showed redundancy gaps ranging from $10^{-3}$ to $0.5$ nats per symbol on adversarial inputs.
• Performance (Scalar):
  • Window-smart (Alg 3): Consistently fast ($\approx 36\text{--}49$ cycles/symbol) regardless of input distribution shape. It is the "safe default" for unknown or heavy-tailed inputs.
  • Bloom-smart (Alg 2): Extremely fast ($\approx 13$ cycles/symbol) on uniform-like inputs where the geometric initialization is close to optimal. However, performance degrades significantly (up to $119$ cycles/symbol) on skewed distributions (e.g., Geometric $p=0.95$) as the initialization drifts further from the optimum.
  • Bottom-up: Significantly slower (orders of magnitude) due to $O(M)$ iterations, serving only as a correctness benchmark.
• Performance (AVX2):
  • Window-super: Achieved a $1.4\times$ to $2.2\times$ speedup on heavy-tailed inputs by using a Lagrangian threshold path (bisection) to avoid materializing tickets, effectively vectorizing the selection process.
  • Bloom-super: Saw limited gains because the heap operations (Phase 1 and 2) are serial and data-dependent, resisting vectorization.

---

5. Significance and Conclusion

This paper bridges the gap between theoretical optimality and practical performance in entropy coding.

• Theoretical Impact: It resolves the open question of whether a polynomial-time, exact KL-optimal normalizer exists that is fast enough for real-time coding. The answer is yes, via the $O(r)$ window method.
• Practical Impact:
  • Codec developers can now replace heuristic normalizers with exact ones without sacrificing speed.
  • The paper provides a "complementary" strategy: Use Bloom-smart for uniform/symmetric data (fastest) and Window-smart for skewed/unknown data (most robust).
  • The redundancy reduction (though small in bits per symbol) is mathematically guaranteed, which is critical for high-efficiency compression standards (like zstd, CRAM, rANS) where every bit counts.
• Future Work: The authors suggest investigating tighter window bounds ($o(r)$), integer-arithmetic implementations to avoid floating-point issues, and end-to-end codec integration studies to measure file-size reductions in production.

In summary, Szewczyk presents a rigorous mathematical framework and highly optimized implementations that make exact KL-optimal frequency normalization a viable, high-performance reality for modern entropy coders.