Streaming Entropy with Numerical Stability

Streaming Entropy with Numerical Stability

In machine learning pipelines, monitoring the Shannon entropy of token distributions (e.g., model predictions, vocabulary usage, or feature categories) is essential for detecting distribution shift and model collapse in real time.

Design and implement a class StreamingEntropy that tracks the Shannon entropy of an unbounded stream of categorical tokens in base-2 (bits).

Mathematical Definition

Given a stream with total token count N and distinct token counts c_i, the empirical probability of token i is p_i = c_i / N. The Shannon entropy H in bits is defined as:

H = - sum(p_i * log2(p_i)) for all unique tokens i where p_i > 0.

By convention, if N == 0, H = 0.0.

Requirements

1. add(token: str) -> None: Ingests a new token into the stream. Must run in O(1) amortized time.
2. get_entropy() -> float: Returns the current Shannon entropy in bits in O(1) time.
3. Numerical Stability: Standard naive accumulation using log2(N) - (1/N) * sum(c_i * log2(c_i)) suffers from catastrophic cancellation when H approaches 0 (e.g., when one token dominates the distribution) and precision degradation when N is large. Your solution must maintain numerical accuracy even for streams with millions of tokens and skewed distributions.

Interface

class StreamingEntropy:
    def __init__(self):
        pass

    def add(self, token: str) -> None:
        pass

    def get_entropy(self) -> float:
        pass

Examples

Example 1:

stream = StreamingEntropy()
stream.add("A")
stream.get_entropy()  # Returns 0.0 (only one unique token)
stream.add("B")
stream.get_entropy()  # Returns 1.0 (two tokens, each with p = 0.5)
stream.add("A")
stream.add("B")
stream.get_entropy()  # Returns 1.0 (counts: A=2, B=2)

Example 2:

stream = StreamingEntropy()
for _ in range(3):
    stream.add("A")
stream.add("B")
# Counts: A=3, B=1, N=4
# H = -(3/4 * log2(3/4) + 1/4 * log2(1/4))
# H = -(0.75 * (-0.415037) + 0.25 * (-2.0)) = 0.811278...
stream.get_entropy()  # Returns approximately 0.8112781244591328

Constraints

• The total number of tokens N can exceed 10^7.
• The number of unique tokens K can be up to 10^5.
• Memory consumption must be O(K), proportional to the number of unique tokens seen.