blank

Initials Atlas

2026-06-16T04:20:00+00:00

Notes on Variational Inference and Jensen’s Inequality

2026-04-17T13:05:00+00:00

Spatial Competence Benchmark

2026-03-07T23:30:00+00:00

List of open source contributions to Inspect AI

2026-03-04T23:30:00+00:00

Preview: Visual Geometry Bench

2025-11-13T00:00:00+00:00

What do we measure?

Visual Geometry Bench measures LLM spatial cognition through questions that compel humans to visualise in order to solve, aka test the “Geometric Intuition”. We probe how models reconstruct global structure from fragments, reason across spatial hierarchies, and infer topological relationships, all without visual input. All code can be found here: GitHub Repo.

Tasks

Below are the core tasks in Visual Geometry Bench and the aspects of spatial or topological reasoning they test. Each task is defined as a family of configurations with programmatic verifiers, giving us large diversity in instance structure and difficulty while keeping evaluation fully reliable.

Note: All images and GIFs in this post are for human comprehension only. They are never part of the model’s input. Models receive purely textual descriptions and must solve each instance using internal chain-of-thought reasoning, without tools, search, or external solvers.

Topology Edge Tasks

Corner labels are provided for a square whose interior may be partitioned by arbitrary curves. The model must predict the boundary edge pairs that must connect through the interior.

Challenge: Infer deterministic connectivity from minimal topological information and even correctly identify cases where no edge connections can be forced. Given configurations are not always in human-friendly order (e.g. in sample question 2: ‘bottom-right’, ‘top-right’, ‘top-left’, ‘bottom-left’), forcing deliberate positional reasoning.

Sample Question 1

Hill climb on MBPP using verifiers

2025-09-08T23:30:00+00:00

From scratch: SFT and GRPO on Qwen 2.5

2025-08-01T01:30:00+00:00

Intuiting Policy Gradient methods

2025-07-26T09:10:20+00:00

Recently, I found it imperative to grok Policy Gradient (PG) methods. As much as I enjoy entering rabbit holes of adjacent techniques, which are abundant in RL, I have refrained. The motivation is to think effectively about PG research in LLMs/foundational models.

Problem Setting

It comes naturally to me to pit anything which learns with the standard Supervised Learning framework.

In Supervised Learning, the world is static: given a fixed dataset, the task is to learn a parameterised function mapping inputs to the targets. The gradient of your loss function provides a direct, unambiguous error signal. Reinforcement learning breaks this assumption. Here, we have an agent interacting with an environment, only yielding sparse, indirect signals.

This interaction is formalised as a Markov Decision Process (MDP). At each timestep \(t\), we have: \(\begin{aligned} &\text{State:} \quad s_t \\ &\text{Action:} \quad a_t \sim \pi_\theta(a_t \mid s_t) \\ &\text{Reward:} \quad r_t = r(s_t, a_t) \\ &\text{Transition:} \quad s_{t+1} \sim P(s_{t+1} \mid s_t, a_t) \end{aligned}\)

\[\underbrace{\pi_\theta(a_t \mid s_t)}_{\text{Policy}} \qquad \underbrace{r(s, a)}_{\text{Reward function}} \qquad \underbrace{P(s' \mid s, a)}_{\text{Transition probability}}\]

LLM analogy: \(s_t\) is the prompt plus all tokens generated so far, \(a_t\) is the next token to generate, \(\pi_\theta(a_t \mid s_t)\) is the model’s policy (the probability of generating token \(a_t\) given the current context \(s_t\)), \(\tau\) is the complete response sequence, and \(\theta\) are the model weights.

Hence, in the Policy Gradient setting, our true objective is similar to Supervised Learning: to maximise the performance of our parameterised policy. However, since we do not possess golden labels, we do this indirectly by maximising the expected return over trajectories, \(J(\theta)\):

\[J(\theta) = \mathbb{E}_{\tau \sim \pi_\theta}[R(\tau)] = \sum_{\tau} P(\tau\mid \theta) R(\tau), \quad R(\tau) = \sum_{t=0}^T r_t \tag{1}\]

where a trajectory \(\tau = (s_0, a_0, \dots, s_{T+1})\) is a sequence of states and actions of length \(T\), and \(\pi_\theta\) is the policy parameterised by \(\theta\).

The policy \(\pi_\theta\) influences its own data distribution, creating a shifting landscape where we hope that the data gets better.

To optimise \(J(\theta)\), we use gradient ascent:

\[\theta \leftarrow \theta + \alpha \nabla_\theta J(\theta) \tag{2}\]

Also, the trajectory probability factors can be expanded into:

\[P(\tau \mid \theta) = \rho_0(s_0) \prod_{t=0}^{T} P(s_{t+1} \mid s_t, a_t) \pi_\theta(a_t \mid s_t) \tag{3}\]

i.e. the probability of a trajectory is the product, at each timestep, of the environment’s transition probability and the policy’s action probability, beginning from the initial state distribution \(\rho_0\). This factorisation is key in the next section.

Policy Gradient Theorem

However, there is a glaring issue with equation \((2)\). Computing \(\nabla_\theta J(\theta)\) directly from equation \((1)\) is infeasible because:

1. Exact Gradient Is Infeasible

Calculating the true gradient requires summing over all possible trajectories \(\tau\) in equation \((1)\). This is computationally intractable since the number of possible trajectories grows exponentially with trajectory length.

2. State Distribution Issue

Differentiating the objective would require differentiating the environment’s state distribution \(\rho_0(s_0)\) and transition probabilities \(P(s_{t+1}\mid s_t, a_t)\) in equation \((3)\). These environment dynamics are typically unknown or non-differentiable.

The Practical Solution

The policy gradient theorem provides an elegant solution:

Analytical Form: Derive the policy gradient as an expectation over trajectories:
\[\nabla_\theta J(\theta) = \mathbb{E}_{\tau \sim \pi_\theta} \left[ R(\tau) \sum_{t=0}^{T} \nabla_\theta \log \pi_\theta(a_t\mid s_t) \right] \tag{4}\]

The final form of the Policy Gradient Theorem saves us from differentiating through the environment dynamics. One could even interpret it as a common stoic lesson:

Make the best use of what is in our power, and treat the rest in accordance with its nature. ~ Epictetus

See Appendix A for complete treatment on how we got to equation \(4\).

Estimation: Estimate this expectation using sampled trajectories from agent-environment interactions.

REINFORCE

An example of estimation of the expectation is Monte Carlo REINFORCE. Specifically, we collect multiple trajectories and compute:

\[\hat{\nabla}_\theta J(\theta) = \frac{1}{m} \sum_{i=1}^{m} \sum_{t=0}^{T} \nabla_\theta \log \pi_\theta(a_t^{(i)}\mid s_t^{(i)}) R(\tau^{(i)}) \tag{5}\]

where:

\(m\) is the number of sampled trajectories
\(\tau^{(i)} = (s_0^{(i)}, a_0^{(i)}, \dots, s_{T+1}^{(i)})\) is the \(i\)-th trajectory
\(R(\tau^{(i)}) = \sum_{t=0}^{T} r_t^{(i)}\) is the total return for trajectory \(i\)
\(a_t^{(i)}, s_t^{(i)}\) are the action and state at time \(t\) in trajectory \(i\)

Intuition: Each trajectory gives a noisy gradient estimate; averaging \(m\) of them reduces variance. The total reward \(R(\tau^{(i)})\) weights each trajectory’s impact. This method is unbiased but still noisy.

You can probably deduce this method is pretty naive.

Example: Raw Returns

State	a1 reward	a2 reward
s1	12	7
s2	2	5

For clarity, we treat these as one-step episodes, so “reward” and “return” coincide: \(R(\tau) = r\). A stochastic policy samples:

\((s_1, a_2)\) (reward = 7) (sub-optimal)
\((s_2, a_2)\) (reward = 5) (optimal)

The raw REINFORCE update is:

\[\hat{\nabla}_\theta J(\theta) = 7\,\nabla_\theta\log\pi_\theta(a_2 \mid s_1) + 5\,\nabla_\theta\log\pi_\theta(a_2 \mid s_2)\]

Here, the wrong action in the high-reward state (\(s_1, a_2\)) receives a larger update than the correct action in the low-reward state, purely due to reward scale and MC randomness.

Variance Reduction

Most takes on Policy Gradient formulations try to resolve this nasty variance issue. One of the most common ways is using a baseline. A baseline \(b(s_t)\) is a function of the state that is subtracted from the return. The gradient estimator becomes:

\[\hat{\nabla}_\theta J(\theta) = \frac{1}{m} \sum_{i=1}^{m} \sum_{t=0}^{T} \nabla_\theta \log \pi_\theta(a_t^{(i)}\mid s_t^{(i)}) (R(\tau^{(i)}) - b(s_t^{(i)})) \tag{6}\]

Let’s apply this to our example. A simple baseline is the average reward for each state:

For \(s_1\): \(b(s_1) = \frac{12 + 7}{2} = 9.5\) (average of possible rewards)
For \(s_2\): \(b(s_2) = \frac{2 + 5}{2} = 3.5\) (average of possible rewards)

For our sampled actions:

\((s_1, a_2) \implies 7 - 9.5 = \mathbf{-2.5}\) (Correctly discourages below-average action)
\((s_2, a_2) \implies 5 - 3.5 = \mathbf{+1.5}\) (Correctly encourages above-average action)

Intuition: Raw returns can be huge or tiny depending on the state. Subtracting \(b(s)\) recentres the reward around “what’s normal” for that state, reducing variance by making updates reflect surplus/deficit rather than absolute scale.

The Quest for the Optimal Baseline

The optimal baseline, \(b(s_t)\), should minimise the variance of the gradient estimator in equation \((6)\) without introducing bias. A baseline is guaranteed to be unbiased if it only depends on the state \(s_t\). Hence, its expected contribution to the gradient is zero: \(\mathbb{E}_{a_t \sim \pi_\theta(\cdot\mid s_t)}[\nabla_\theta \log \pi_\theta(a_t\mid s_t)b(s_t)] = 0\). See the proof in Appendix B.

Now, minimising variance can be tricky, however, we can simplify the problem as follows:

We ignore correlations between timesteps, focusing on a single timestep’s gradient contribution.
We treat this contribution’s gradient term as a scalar. We call this term \(g_t\).

\[g_t = \nabla_\theta \log \pi_\theta(a_t\mid s_t)\]

The full derivation and intuition are provided in Appendix C. The practical result is that the optimal baseline is the state-value function:

\[b^*(s_t) = \frac{\mathbb{E}\left[ g_t\,R_t \mid s_t \right]}{\mathbb{E}\left[ g_t^2 \mid s_t \right]} \approx \mathbb{E}[R_t \mid s_t] = V^{\pi}(s_t) \tag{7}\]

Looking through the Advantage Function Lens

Given the action-value function: \(Q^{\pi}(s, a) = \mathbb{E}_{\pi}\left[ R_t \mid s_t = s, a_t = a \right]\)

In practice, we estimate \(Q^{\pi}(s, a)\) with the reward-to-go from a single trajectory: \(R_t^{(i)} = \sum_{k=t}^T \gamma^{k-t} r_k^{(i)}\) This uses only future rewards, as actions cannot affect the past.

We can define the advantage function as: \(A^{\pi}(s, a) = Q^{\pi}(s, a) - V^{\pi}(s)\) where \(V^{\pi}(s) = \mathbb{E}_{a \sim \pi(\cdot\mid s)}[Q^{\pi}(s, a)]\) is the expected value of the state under the policy.

Intuition: \(A^{\pi}(s, a)\) measures how much better (or worse) action \(a\) is compared to the average. Effectively, highlighting outlier actions.

The baseline-subtracted term from equation \((6)\) is exactly our empirical advantage estimate:

\[\hat{A}^{\pi}(s_t^{(i)}, a_t^{(i)}) = R_t^{(i)} - b(s_t^{(i)})\]

Hence, the final policy gradient estimator uses the advantage function:

\[\nabla_\theta J(\theta) = \mathbb{E}_{\tau \sim \pi_\theta} \left[ \sum_{t=0}^T \nabla_\theta \log \pi_\theta(a_t\mid s_t) A^{\pi}(s_t, a_t) \right] \tag{8}\]

Appendix

A. Policy Gradient Theorem Proof

The policy gradient theorem gets us over the direct gradient calculation hump. However, it is important to understand how we shed certain terms to get to the final form. Here’s the complete derivation:

1: Start from the objective definition

\[\nabla_\theta J(\theta) = \nabla_\theta \sum_{\tau} P(\tau\mid \theta) R(\tau)\]

Why: We start with equation \((1)\), \(J(\theta) = \sum_{\tau} P(\tau\mid \theta) R(\tau)\), and take its gradient.

2: Linearity of differentiation

\[\nabla_\theta J(\theta) = \sum_{\tau} \nabla_\theta P(\tau\mid \theta) R(\tau)\]

3: Apply the log-derivative trick

\[\nabla_\theta J(\theta) = \sum_{\tau} P(\tau\mid \theta) \nabla_\theta \log P(\tau\mid \theta) R(\tau)\]

Why: The basic log trick: \(\nabla_\theta \log f = \frac{\nabla_\theta f}{f}\)

4: Rewrite as expectation

\[\nabla_\theta J(\theta) = \mathbb{E}_{\tau \sim \pi_\theta} \left[ R(\tau) \nabla_\theta \log P(\tau\mid \theta) \right]\]

Why: The sum over trajectories weighted by their probabilities is equivalent to an expectation under the policy.

5: Factorise trajectory probability

Let’s bring back equation \((3)\) into focus, we expand the trajectory probability:

\[P(\tau \mid \theta) = \rho_0(s_0) \prod_{t=0}^{T} P(s_{t+1} \mid s_t, a_t) \pi_\theta(a_t \mid s_t)\]

6: Take logarithm of the product

\[\log P(\tau \mid \theta) = \log \rho_0(s_0) + \sum_{t=0}^{T} \log P(s_{t+1} \mid s_t, a_t) + \sum_{t=0}^{T} \log \pi_\theta(a_t \mid s_t)\]

7: Differentiate and eliminate environment terms

\[\nabla_\theta \log P(\tau \mid \theta) = \nabla_\theta \log \rho_0(s_0) + \sum_{t=0}^{T} \nabla_\theta \log P(s_{t+1} \mid s_t, a_t) + \sum_{t=0}^{T} \nabla_\theta \log \pi_\theta(a_t \mid s_t)\]

Since \(\rho_0(s_0)\) and \(P(s_{t+1}\mid s_t, a_t)\) don’t depend on policy parameters \(\theta\):

\[\nabla_\theta \log \rho_0(s_0) = 0\] \[\nabla_\theta \log P(s_{t+1}\mid s_t, a_t) = 0\]

Therefore:

\[\nabla_\theta \log P(\tau\mid \theta) = \sum_{t=0}^{T} \nabla_\theta \log \pi_\theta(a_t\mid s_t)\]

Why: Only policy terms carry \(\theta\) dependence. Environment dynamics are constant with respect to policy parameters.

8: Substitute back to obtain final result

\[\nabla_\theta J(\theta) = \mathbb{E}_{\tau \sim \pi_\theta} \left[ R(\tau) \sum_{t=0}^{T} \nabla_\theta \log \pi_\theta(a_t\mid s_t) \right]\]

In this formulation we relinquish control over the environment. However, I believe, high-agency humans do exert control over their environment to get to their goals.

B. Proof of Unbiasedness for State-Dependent Baselines

The proof relies on the fact that the expectation of the score function is zero.

Proof of Zero-Expectation Score:

\[1 = \int_{a_t} \pi_\theta(a_t\mid s_t) da_t\]

Any probability distribution must integrate to 1.

\[\nabla_\theta(1) = \nabla_\theta \int_{a_t} \pi_\theta(a_t\mid s_t) da_t\] \[0 = \int_{a_t} \nabla_\theta \pi_\theta(a_t\mid s_t) da_t\] \[0 = \int_{a_t} \pi_\theta(a_t\mid s_t) \nabla_\theta \log \pi_\theta(a_t\mid s_t) da_t\]

Apply the identity \(\nabla_\theta f = f \nabla_\theta \log f\) (same as Appendix A).

\[0 = \mathbb{E}_{a_t \sim \pi_\theta(\cdot\mid s_t)}[\nabla_\theta \log \pi_\theta(a_t\mid s_t)]\]

The integral defines the expectation, proving the expected score is zero.

Unbiased Baseline Proof:

Since \(b(s_t)\) is constant with respect to the expectation over actions, it can be factored out.

\[\begin{aligned} \mathbb{E}_{a_t \sim \pi_\theta(\cdot \mid s_t)}[\nabla_\theta \log \pi_\theta(a_t \mid s_t) b(s_t)] &= b(s_t) \mathbb{E}_{a_t \sim \pi_\theta(\cdot \mid s_t)}[\nabla_\theta \log \pi_\theta(a_t \mid s_t)] \\ &= b(s_t) \cdot 0 \\ &= 0 \end{aligned}\]

This proves the baseline is a control variate.

C. Optimal Baseline Derivation

To find the variance-minimising baseline \(b(s_t)\) under the simplified model, we analyse the variance of a single gradient sample, \(Z_t = (R_t - b(s_t))g_t\), where \(g_t\) is treated as a scalar: \(g_t = \nabla_\theta \log \pi_\theta(a_t\mid s_t)\).

Step 1: Set Up the Variance Minimisation

The goal is to minimise \(\text{Var}[Z_t] = \mathbb{E}[Z_t^2] - (\mathbb{E}[Z_t])^2\). Since the baseline is unbiased, \(\mathbb{E}[Z_t]\) is a constant with respect to \(b(s_t)\), so we only need to minimise \(\mathbb{E}[Z_t^2]\). For a given state \(s_t\):

\[\min_{b(s_t)} \mathbb{E}_{a_t, R_t | s_t} \left[ g_t^2 (R_t - b(s_t))^2 \right]\]

Why is \(\mathbb{E}[Z_t]\) constant?
Because \(\mathbb{E}[Z_t] = \mathbb{E}[g_t(R_t-b(s_t))] = \mathbb{E}[g_t R_t] - b(s_t)\,\mathbb{E}[g_t]\), and by the result in Appendix B, \(\mathbb{E}[g_t] = 0\). Thus, the baseline-dependent part vanishes, leaving a term independent of \(b(s_t)\).

Step 2: Differentiate and Solve

We find the minimum by differentiating with respect to \(b(s_t)\) and setting the result to zero:

\[\frac{\partial}{\partial b(s_t)} \mathbb{E}\left[g_t^2(R_t - b(s_t))^2 | s_t\right] = \mathbb{E}\left[-2g_t^2(R_t-b(s_t))|s_t\right] = 0\]

This simplifies to \(\mathbb{E}[g_t^2 R_t \mid s_t] = b(s_t) \mathbb{E}[g_t^2 \mid s_t]\), which gives the optimal baseline:

\[b^*(s_t)=\frac{\mathbb{E}[g_t^2 R_t \mid s_t]}{\mathbb{E}[g_t^2\mid s_t]}\]

Since \(b(s_t)\) is a function of state \(s_t\) only, it is constant with respect to the expectation over actions and rewards, conditioned on that state.

Step 3: One more approximation

This weighted baseline is still impractical. The final leap is to assume the policy sensitivity term, \(g_t^2\), can be factored out and cancelled.

Why is this reasonable?
While \(g_t\) itself varies by action, its expected squared norm, \(\mathbb{E}[g_t^2\mid s_t]\), can be seen as a measure of the policy’s overall “instability” at a state \(s_t\). By assuming this instability is not strongly correlated with the rewards \(R_t\), we can approximate \(\mathbb{E}[g_t^2 R_t \mid s_t] \approx \mathbb{E}[g_t^2\mid s_t]\mathbb{E}[R_t\mid s_t]\). It’s important to note that this is not an exact result, but rather a pragmatic workaround.

The sensitivity term now cancels out, leaving the simple expected return: \(b(s_t) \approx \frac{\mathbb{E}[g_t^2 \mid s_t] \mathbb{E}[R_t \mid s_t]}{\mathbb{E}[g_t^2 \mid s_t]} = \mathbb{E}[R_t \mid s_t] \equiv V^{\pi}(s_t)\)

Acknowledgments

Excellent resources I referenced to build my understanding:

Daniel Takeshi’s blog post on Policy Gradient Fundamentals
OpenAI Spinning Up
RLHF Book: Policy Gradient Algorithms
Lilian Weng’s Policy Gradient Methods
Sutton, R. S., & Barto, A. G. (2018). Reinforcement Learning: An Introduction (2nd ed.). MIT Press.

Citation

Please cite this work as:

Vira, Jash. “Intuiting Policy Gradient methods.” jashvira.com (Jul 2025). https://jashvira.com/blog/2025/policy_gradients/

Or use the BibTeX citation:

@article{ vira2025intuitingpolicygradientmethods,
  title="Intuiting Policy Gradient methods",
  author={Vira, Jash},
  journal={jashvira.com},
  year={2025},
  month={Jul},
  url={https://jashvira.com/blog/2025/policy_gradients/}
}

LLM Agent ~ Reddit Consensus

2025-07-25T23:30:00+00:00

Neural Network precision pitfalls in the wild

2025-05-05T12:02:20+00:00

In my work, I draw on concepts from computational geometry and graphics, often demanding high-precision guarantees. A typical case is sampling an object’s implicit function to generate a parameterised form for downstream use.

Neural networks tend to be a perfect fit for such problems. They are great at expressing complex, nonlinear structures, are differentiable, and can be infinitely precise. In theory, the precision part is well-supported, but in practice, let’s explore the challenges that arise. I’ll take a top-bottom approach, grounding the problem scenarios in the real-world, and then get into pitfalls to lookout for and even establish certain error bounds.

Problem setting

To reflect on an actual use-case (from my work), let \(F\colon\mathbb{R}^3\to\Delta^1\subset\mathbb{R}^2\) be our neural network. For illustration we restrict to a one‐dimensional slice in \(\mathbb{R}^3\). Fix two coordinates \(x_2^*,x_3^*\) and let

\[L \;=\;\bigl\{(t,x_2^*,x_3^*) : t\in[t_{\min},\,t_{\max}]\bigr\}\]

for some suitably chosen interval \([t_{\min},t_{\max}]\). \(F\) is trained by sampling the two ground-truth curves shown in the plot below.

\[p_0(t) \quad\text{and}\quad p_1(t)\;=\;1 - p_0(t)\]

Figure 1. Plot of \(p_0(x)=0.5+0.5\sin(45\pi x)e^{-300x^2}\) and its complement \(1 - p_0(x)\).

We shall use JAX/FLAX for all experiments in this work.

Pitfalls

Now, given we have a well trained Neural Network \(F\) that approximates a function. How reliably can we sample \(F\)?

Error accumulation in a forward pass

A network can fit such data very easily, given decent complexity in it’s architecture. However, there is also an unintended effect on the stability of its predictions as the numbers flow through a forward pass. Let’s analyse this through a gridsearch over the depth and width of a simple ResNet.

Figure 2. Grid search for accumulated error over depth and width of network. Default matmul settings used here.

In the above image we sample a neural network in a small interval, where the two classes overlap. For many geometry applications this is an area of interest, since it contains a zero for \(F_1-F_0\). However, because of noise accumulated, there exist multiple zeroes here.

This raises a fair few questions!

Why does this noise occur?
Why is it exaggerated for bigger networks?
Given a network, can we find bounds for it?

These questions force us to quantify, how errors can be injected in a forward pass. Some of the relevant pieces being:

\[u \;=\; 2^{-23}\;\approx1.19\times10^{-7}, \qquad \gamma_n \;=\;\frac{n\,u}{1 - n\,u}\;\approx\;n\,u.\]

Here \(u\) is the unit round-off for IEEE-754 float32 (not what is used by default!), and \(\gamma_n\) is the worst-case relative error of an \(n\)-term dot product.

Note: In most famous Deep Learning libraries, the default matmul settings use a mix of float32 and TF32 calculations depending on the kernel employed. This increases the instability! More on this later.

Component	Fan-in / ops	Worst-case relative error bound	Notes / mitigation
Dense layer	\(n\) multiplies + adds	\(\le\gamma_n\)	narrower layers ⇒ smaller \(\gamma_n\)
ReLU / skip-add	1 compare + 1 add	\(\le u < \gamma_1\)	negligible
BatchNorm / LayerNorm	≈3 ops	\(\le3\,u < \gamma_3\)	normalises activations ⇒ later \(\gamma\) shrink
TF32 matmul (default of most libraries)	same count, 10-bit mantissa	\(\gamma_n^{\rm TF32}\approx n\,2^{-10}\)	Can be made more precise, pushing the error down

Because each layer \(\ell\) injects at most a \(\gamma_{n_\ell}\) relative slip, chaining \(L\) layers gives (dropping \(O(u^2)\) terms):

\[\hat f(x) = f(x)\,\prod_{\ell=1}^L(1+\delta^{(\ell)}) \;=\; f(x)\,\bigl(1 + \underbrace{\sum_{\ell=1}^L\delta^{(\ell)}}_{\displaystyle\varepsilon_{\rm acc}(x)}\bigr), \quad \bigl|\delta^{(\ell)}\bigr|\le\gamma_{n_\ell} \;\Longrightarrow\; \bigl|\varepsilon_{\rm acc}(x)\bigr|\le\sum_{\ell=1}^L\gamma_{n_\ell}.\]

That is, the total forward-pass error is bounded by the sum of each layer’s \(\gamma\). See Appendix for the detailed proof.

Kernel Drift for different batch sizes

When you compare “single-input” (loop) inference to “multi-input” (batch) inference, you’re driving two different reduction pipelines, meaning they can produce slightly different sums in floating-point. Hence producing non-bitwise identical results for the exact same point!

In geometric or scientific applications, these \(10^{-3}\) scale jitters can introduce false zero crossings, spurious topology changes, or just plain wrong decisions at interfaces where two probabilities are nearly equal.

Figure 3a: Default (with TF32) kernel drift error.

Figure 3b: FP32 kernel drift error.

Empirical observation

Across different widths, we measured (light blue background)

\[\Delta(x)\;=\;P_0^{\rm loop}(x)\;-\;P_0^{\rm batch}(x)\]

and summarised each configuration by \(\mathrm{RMSE}(\Delta)\). we see (for the 4060ti 16GB):

Drift starts small at very narrow widths.
It peaks at intermediate widths (a few hundred).
Then it falls off again at large widths, where both loop and batch paths happen to use the same kernel.

This non-monotonic pattern arises because XLA’s choice of GPU reduction kernels and their internal tiling depend on matrix dimensions. At certain sizes, the two pipelines diverge and coincide at others.

Remedies

Uniform reductions: Force a single, full-precision FP32 path (no mixed-precision or fast-math shortcuts) so both modes use identical kernels. However, this still does not guarantee bitwise-identical results.
Higher precision: Run critical inference in FP64 or a deterministic 64-bit-mantissa format, this also like the first remedy only pushes down the error further.
Pad Batches: If we pad all our calls and use the same batch, we can atleast make the error reproducible.

Appendix

Bounds for error accumulation

First let’s get started with the smallest unit. Since most frameworks have the default setting of float32 we’ll assume that for now. For IEEE‑754 single‑precision arithmetic every real number \(z\) is rounded by

\[\operatorname{fl}(z)=z\,(1+\varepsilon),\qquad|\varepsilon|\le u,\qquad u:=2^{-23}\approx1.19\times10^{-7}. \tag{1}\]

The relative error is bounded by \(u\); the absolute error scales with \(\lvert z\rvert\).

Now let’s extend this to the dot product. Let

\[z=\sum_{k=1}^n a_k b_k,\qquad \widehat z=\operatorname{fl}(z).\]

Giving us:

\[\widehat z = z\,(1+\delta),\quad |\delta|\le\gamma_n,\quad \gamma_n := \frac{n\,u}{1 - n\,u}\approx n\,u. \tag{2}\]

Proof for \(\gamma_n\)

By compounding the worst‐case rounding factor \((1+u)\) over \(n\) operations, the relative error satisfies \(|\delta_n|\le (1+u)^n - 1 \;\le\;\sum_{k=1}^n\binom{n}{k}u^k \;\le\;\sum_{k=1}^n(nu)^k \;=\;\frac{n\,u}{1 - n\,u} \;=\;\gamma_n.\)

Specifically:

Binomial theorem: \((1+u)^n = \sum_{k=0}^n\binom{n}{k}u^k\;\Longrightarrow\;(1+u)^n - 1 = \sum_{k=1}^n\binom{n}{k}u^k\)
Bound coefficients: \(\binom{n}{k}\le n^k\;\Longrightarrow\;\binom{n}{k}u^k\le(nu)^k\)
Geometric‐series sum: \(\sum_{k=1}^n(nu)^k = (nu) + \cdots + (nu)^n\)
Closed‐form: \(\displaystyle\sum_{k=1}^n(nu)^k = \frac{n\,u}{1 - n\,u} = \gamma_n\)

Intuition: \((1+u)^n - 1\) is the multiplicative snowball after \(n\) roundings

Error in the linear map

Let

\[y = W x \in \mathbb{R}^m,\qquad \hat y = \operatorname{fl}(Wx).\]

By \((2)\) each length-\(n\) dot-product satisfies

\[\hat y_i = y_i\,(1+\delta_i),\quad |\delta_i|\le \gamma_n, \qquad \gamma_n := \frac{n\,u}{1 - n\,u}\approx n\,u.\]

Hence for each neuron \(i\)

Side note: \(\|\cdot\|_\infty\) simply picks out the worst‐case single‐neuron error, so we get one uniform bound across all \(m\) neurons.

Full layer with 1-Lipschitz activation

Define

\[y = W^{(\ell)}x^{(\ell)},\quad \hat y = \operatorname{fl}(y),\quad x^{(\ell+1)} = \sigma(y),\quad \hat x^{(\ell+1)} = \sigma(\hat y),\]

where \(\sigma\) (ReLU, tanh, …) obeys

\[\|\sigma(v)-\sigma(w)\|_\infty \le \|v-w\|_\infty.\]

Then

\[\|\hat x^{(\ell+1)} - x^{(\ell+1)}\|_\infty \;\le\;\|\hat y - y\|_\infty \;\le\;\gamma_{n_\ell}\,\|y\|_\infty \;\le\;\gamma_{n_\ell}\,\|x^{(\ell+1)}\|_\infty.\]

So for each neuron \(i\) in layer \(\ell+1\)

\[|\hat x_i^{(\ell+1)} - x_i^{(\ell+1)}|\le \gamma_{n_\ell}\,|x_i^{(\ell+1)}|,\]

and we write

\[\hat x^{(\ell+1)} = x^{(\ell+1)}\,(1+\delta^{(\ell)}),\quad |\delta^{(\ell)}|\le \gamma_{n_\ell}. \tag{3}\]

Side note: \(\gamma_{n_\ell}\) uses the layer’s incoming dimensionality \(n_\ell\) (e.g. 256 inputs ⇒ \(\gamma_{256}\)) to bound every neuron’s relative slip.

How the next layer sums up errors.

The \(\infty\)-norm bound above compresses all per-neuron slips into one scalar \(\delta^{(\ell)}\). When these perturbed activations \(\hat y\) enter the next layer’s dot-product, \((2)\) on a length-\(n_{\ell}\) dot-product automatically sums whatever errors remain:

\[|\widehat z_j - z_j| \;=\;\Bigl|\sum_{i}w_{ji}\,\hat y_i - \sum_{i}w_{ji}\,y_i\Bigr| \;\le\;\gamma_{n_{\ell+1}}\sum_i|\hat y_i| \;\le\;\gamma_{n_{\ell+1}}\,n_{\ell+1}\,\|\hat y\|_\infty.\]

Thus the “sum over all neurons” happens inside each subsequent dot-product bound, without ever unpacking the single \(\infty\)-norm scalar.

Accumulated Error Across \(L\) Layers

Chaining these per-layer relative slips gives

\[\hat f(x) = f(x)\prod_{\ell=0}^{L-1}(1+\delta^{(\ell)}) = f(x)\,\bigl(1+\varepsilon_{\rm acc}(x)\bigr),\]

Direct‐expansion
\[\prod_{\ell=0}^{L-1}(1+\delta^{(\ell)})\]
by distributing out each factor, you get one term for each single \(\delta^{(\ell)}\), plus terms involving products of two or more \(\delta\)s. Because each \(\delta^{(\ell)}=O(u)\), any product \(\delta^{(i)}\delta^{(j)}\) is \(O(u^2)\) and hence negligible at single precision. Dropping those higher-order terms leaves exactly
\[1 \;+\;\sum_{\ell=0}^{L-1}\delta^{(\ell)} + O(u^2).\]

where to first order and using \((3)\)

\[\varepsilon_{\rm acc}(x) = \sum_{\ell=0}^{L-1}\delta^{(\ell)} + O(u^2), \quad |\varepsilon_{\rm acc}|\le \sum_{\ell=0}^{L-1}\gamma_{n_\ell}.\]

To reinforce: Each \(\gamma_{n_\ell}\) already folded in the sum of that layer’s neuron-level errors, so summing \(\gamma_{n_\ell}\) over \(\ell\) neatly captures the network’s total forward error.

Note on Conditioning

A more rigorous analysis replaces each \(\gamma_{n_\ell}\approx n_\ell\,u\) by a data-dependent bound \((\kappa_\ell,\,n_\ell,\,u)\), where the dot-product condition number \(\kappa_\ell \;=\; \frac{\sum_{k=1}^n |a_k b_k|}{\bigl|\sum_{k=1}^n a_k b_k\bigr|}\) captures any cancellation or scaling in that layer. Including \(\kappa_\ell\) refines the bound but does not alter its overall form.

Citation

Please cite this work as:

Vira, Jash. “Neural Network precision pitfalls in the wild.” jashvira.com (May 2025). https://jashvira.com/blog/2025/nn_precision_pitfalls/

Or use the BibTeX citation:

@article{ vira2025neuralnetworkprecisionpitfallsinthewild,
  title="Neural Network precision pitfalls in the wild",
  author={Vira, Jash},
  journal={jashvira.com},
  year={2025},
  month={May},
  url={https://jashvira.com/blog/2025/nn_precision_pitfalls/}
}