Guest Post: A Hidden Symmetry, Found Through Formalization

Recently, we used Prove2Me to formalize a paper that David Simchi-Levi, Chonghuan Wang, and I wrote in 2024 on A/B testing in Markov decision processes. In the course of formalization, the AI agents found a substantially shorter proof of one of our main theorems. The argument drew on Rao–Blackwellization, an idea from classical point estimation with which we were unfamiliar. It revealed a more fundamental reason for a result we had established through a lengthy covariance calculation.

What interested us most was that formalization had led the agents to discover a connection we had overlooked. The resulting proof was not only simpler, but also formally verified, giving us confidence that this was a genuine mathematical insight rather than another plausible-sounding AI argument.

Here is the mathematical story.

The original problem

Consider a company testing a "we miss you" coupon. Customers move among a finite collection of states. The control policy takes the usual action, while the treatment policy offers a coupon only in a particular state $s^1$. Everywhere else, the two policies behave identically. We want to estimate the difference in their long-run discounted values,

$$V^t(s)-V^c(s).$$

We observe a single trajectory, where a fresh fair coin determines which policy to follow at each time step. Each observation is

$$X_i=(s_i,\gamma_i,s_{i+1},r_i), \qquad \gamma_i\in\{t,c\},$$

consisting of the current state, the randomized treatment/control label, the next state, and the reward.

A standard A/B estimator splits the observations by their labels and estimates the two policies separately. But outside $s^1$, both policies behave identically, so why not share those observations between the two arms?

This is what we called information sharing.

Formally, let $L(X)$ encode the observed transition and reward. For either arm $a\in\{t,c\}$, ordinary A/B splitting uses

$$u^a_{\mathrm{AB}}(X)=2\mathbf 1\{\gamma=a\}L(X),$$

while information sharing uses

$$u^a_{\mathrm{IS}}(X)= \begin{cases} 2\mathbf 1\{\gamma=a\}L(X), & s=s^1,\\ L(X), & s\ne s^1. \end{cases}$$

One of our main theorems establishes that information sharing can only reduce the asymptotic variance of a broad class of estimators.

Our original proof was quite involved. Because the observations are temporally dependent, we had to compute covariance matrices block by block, show that various cross-time terms cancel, and eventually establish positive semidefiniteness.

We expected this to be particularly painful to formalize, since Lean would require every omitted calculation to be justified. Instead, the agents found a surprisingly simple proof.

The hidden symmetry

The key was a transformation called armFlip in the generated Lean code.

Given an observation $X=(s,\gamma,s',r)$, define $\tau X$ by flipping the treatment/control label whenever $s\ne s^1$, leaving everything else unchanged.

Crucially, $\tau$ is an involution ($\tau^2=\mathrm{id}$) that preserves the distribution of $X$. This is because the label comes from a fair coin, and outside $s^1$, both policies take the same action.

Now consider the following elementary result, a symmetric version of Rao–Blackwellization.

Theorem (Symmetric Rao–Blackwell). Suppose $\tau$ is a distribution-preserving involution. For any square-integrable statistic $f$, define

$$\bar f=\frac{f+f\circ\tau}{2}, \qquad d=\frac{f-f\circ\tau}{2}.$$

Then

$$\boxed{ \mathbb E[\bar f]=\mathbb E[f],\qquad \operatorname{Var}(f) =\operatorname{Var}(\bar f)+\mathbb E[d^2] \ge\operatorname{Var}(\bar f). }$$

The proof is immediate from symmetry: $\bar f$ is unchanged by $\tau$, while $d$ changes sign. Distribution preservation therefore implies $\mathbb E[d]=0$ and $\mathbb E[\bar f d]=0$.

Here comes the beautiful observation made by the agents:

$$\boxed{ u_{\mathrm{IS}}(X) =\frac12\left( u_{\mathrm{AB}}(X)+u_{\mathrm{AB}}(\tau X) \right). }$$

Information sharing is exactly the Rao–Blackwellization of ordinary A/B splitting!

There is one subtlety: the theorem above concerns individual observations, whereas our estimators depend on a Markovian trajectory. Fortunately, the same symmetry resolves this too. Flipping the label at any single time $i$, outside $s^1$, preserves the distribution of the entire trajectory, because that label has no effect on the action or subsequent state.

Writing $u_{\mathrm{AB}}(X_i)=u_{\mathrm{IS}}(X_i)+d_i$, this symmetry implies

$$\mathbb E[d_i u_{\mathrm{IS}}(X_j)^\top]=0, \qquad \mathbb E[d_i d_j^\top]=0\quad(i\ne j).$$

Thus all nonzero-lag covariance terms remain unchanged, and the difference between the asymptotic covariance matrices is simply

$$\boxed{ \Sigma_{\mathrm{AB}}-\Sigma_{\mathrm{IS}} =\mathbb E[d_i d_i^\top]\succeq0. }$$

Under the regularity assumptions of our paper, the delta method extends this comparison to our class of differentiable estimators.

That is essentially the entire proof. What originally required lengthy covariance calculations now follows from a single symmetry. Even better, it reveals why information sharing works: it removes randomness introduced by an experimental label that the underlying system itself ignored.

Beyond verification

We had not recognized this connection when writing the original paper. The agents did not discover a new theorem, but they gave us a much better understanding of one we had already proved.

This is what excites us about Prove2Me: formalization is not merely about verifying existing mathematics. It can also help AI agents discover simpler proofs and unexpected connections, while formal verification ensures that these insights are mathematically trustworthy.