Most testing strategies miss rare edge cases until customers find them in production. Our system automatically generates targeted unit tests for rare bugs, including the one that would have caught Anthropic’s recent approximate top-K bug. In this blog post, we’ll provide a brief overview of how it works.
Catching the rare bug in top-K sampling
A bug in the TPU implementation of approximate top-K resulted in the most likely token sometimes being excluded. Rare bugs like this frequently slip through to production because covering every behavior with testing is infeasible in practice. After discovery, Anthropic provided a simple reproducer of the bug, but it is the sort of test you only manage to write after a laborious bug minimization process.
We used fractional proof decomposition to automatically generate the unit test without relying on Anthropic’s bug reproducer code. You can run the unit test on colab. For any code, if testing is done via fractional proof decomposition, bugs can be systematically found without the benefit of hindsight.
@given(k=st.integers(min_value=0, max_value=TOP_K_RANGE), arr=arr_strategy)
def test_approx_max_k(k, arr):
N = len(arr)
k = int(k % min(N - MIN_TOP_K, TOP_K_RANGE)) + MIN_TOP_K
approx_values, _ = lax.approx_max_k(arr, k=k)
assert jnp.max(approx_values) == jnp.max(arr), \
max(approximate_top_k(arr, k=k)) == max(arr). If the implementation of lax.approx_max_k is correct, we should expect the test to pass because the approximate top-K algorithm is implemented by dividing data points into L bins and computing the true max in each bin. L is chosen based on the desired average recall r as .Systematically generating tests via fractional proof decomposition
Step 1: Identify the theorem, which is the property that your implementation must satisfy. Then, encode it as a PBT using the Hypothesis framework. We call the top-level theorem an end-to-end PBT because it corresponds to the end-to-end behavior of the function.
The theorem for the top-K bug is:
$$\forall\ \text{prompt}, k,\ LLM_{\text{top-1}}(\text{prompt}) \in LLM_{\text{top-}k}(\text{prompt})$$For any $\text{prompt}$, $LLM_{\text{top-k}}(\text{prompt})$ represents the model’s prediction of the top-K most likely next tokens. Now, we can encode this as an end-to-end PBT. Since the end-to-end PBT does not need to run on TPU, we set up a different colab.
@given(prompt=st.text(min_size=1, max_size=MAX_SIZE), k=st.integers(1, VOCAB_SIZE))
@settings(max_examples=50, deadline=None)
def test_top_token_present(prompt, k):
greedy_params = SamplingParams(temperature=0.0, max_tokens=1, logprobs=1)
topk_params = SamplingParams(temperature=1.0, max_tokens=1, logprobs=MAX_LOGPROBS, top_k=k)
most_likely_token = llm.generate([prompt], greedy_params, use_tqdm=False)[0].outputs[0].token_ids[0]
logprobs = llm.generate([prompt], topk_params, use_tqdm=False)[0].outputs[0].logprobs[0]
assert most_likely_token in logprobs
Although the end-to-end PBT has comprehensive coverage, to catch rare bugs would take an excessively large number of tokens. The rarer the bug, the more compute is required.
Step 2: Recursively decompose the theorems into a collection of smaller theorems, which are also encoded as PBTs. These smaller theorems are intermediate results that compose together to establish the original end-to-end PBT.
Decomposing the end-to-end PBT for the top-K bug yields, by construction, three theorems:
max(approximate_top_k(arr, k=k)) == max(arr)(true max always included)- On any input tokens, the logits are finite (not ∞ and not NaN)
- In vLLM, the token ids are the same as the logprobs dict keys
You can think of PBTs as fractional components of the brute-force proof. Just as the brute-force proof is optimized by decomposing properties into logical sub-properties via reasoning, better known as partial evaluation, we’re applying reasoning to decompose the fractional brute-force proofs. The reasoning bootstraps trust in PBT coverage–so even though you don’t exhaustively check every single input like a formal proof, you get systematic understanding of your programs, and control over how you spend your testing compute. We’re calling this sampling technique fractional proofs.
Step 3: Continue decomposing until the input space is small enough to be compute efficient at catching rare bugs. You stop decomposing when:
- Each sub-test runs sufficiently quickly,
- Each sub-test tests a sufficient fraction of its input distribution so as to catch most bugs that would end up in code of similar complexity to that being tested, and
- The sub-properties provably compose to cover the full end-to-end property
While we found the top-K bug in about 10m of sampling, we were able to find the XLA:TPU bug (also discussed in Anthropic’s post), involving an issue with excess-precision, in just a few seconds.
Fractional proofs for efficient oversight
Systematic decomposition catches rare bugs without sacrificing developer speed or compute efficiency. Instead of scaling compute proportional to the rarity of the bug, fractional proofs scale compute as the logarithm of rarity.
We can straightforwardly extend the approximate top-K example in this post to real world codebases. For example, top-K can be decomposed into a sequence of PBTs testing how libtpu implements the algorithm described in its reference paper. Or, we can use this reasoning to establish how single-TPU behavior composes into cluster behavior.
At Theorem, we’re training models that can automatically reason about program correctness. If you want to catch bugs earlier and make your devs happy, send me an email.