JHCodec: Generalized Ray Tracing Sampler

• JHCodec is a generalized ray tracing framework that reinterprets MCMC sampling by modeling parameter space as an optical medium with a refractive index tied to the likelihood.
• It enforces fair sampling through conservation of optical invariants, unifying methods such as HMC, Gibbs, and Metropolis under one geometric formalism.
• The method utilizes constant-speed ray propagation to mitigate stochastic gradient noise, enabling scalable sampling in high-dimensional settings like large neural networks.

Below is a focused technical summary of the paper’s generalized ray tracing framework for sampling.

---

1. Core idea: sampling by tracing rays in parameter space

The paper’s central idea is to reinterpret MCMC sampling geometrically:

• The parameter vector $x\in \mathbb{R}^D$ is a point in a $D$-dimensional “optical medium.”
• The target density is represented by an unnormalized likelihood $\mathcal{L}(x)$.
• One evolves trajectories not as particles under forces, but as light rays traveling through a medium whose refractive index depends on $\mathcal{L}(x)$.

The key design principle is to choose the refractive index so that the geometric concentration of rays in space exactly matches the target density. The paper argues that in $D$ dimensions, optical refraction changes ray density by powers of $n^{D-1}$, so fair sampling is obtained with

$n(x) = \mathcal{L}(x)^{\frac{1}{D-1}}.$

This is the basic ray tracing sampler. But the paper goes further: by allowing arbitrary sample weights $W(x)$ and/or variable speed $v(x)$, the method becomes a generalized framework that includes several existing samplers as special cases.

So the contribution is not just “a new sampler”; it is a unifying geometric formalism in which many local MCMC methods correspond to different choices of:

• path geometry,
• refractive index $n(x)$,
• speed $D$0,
• and sample weight $D$1.

---

2. Optical invariants that make the method work

2.1 Basic radiance

The paper derives that, in $D$2 dimensions, for refraction without reflection, the conserved optical quantity is

$D$3

where $D$4 is the radiance, interpreted here as the geometric density of ray paths, i.e. the sampling density in position/direction space.

This is derived from power conservation across an interface plus Snell’s law. For incoming/outgoing bundles at indices $D$5,

$D$6

and with differential Snell relations

$D$7

one gets

$D$8

Hence $D$9 is invariant.

2.2 Etendue

The paper also defines the conserved $\mathcal{L}(x)$0-dimensional optical phase-space quantity (analog of Liouville volume in Hamiltonian mechanics) as

$\mathcal{L}(x)$1

This is the generalized etendue. Conservation of etendue explains how cross-sectional area and solid angle transform under ray propagation. The consequence for sampling is that the ray-bundle phase-space volume is compressed by a factor $\mathcal{L}(x)$2.

---

3. Fair sampling from likelihood via refractive index

The basic fair-sampling construction is:

$\mathcal{L}(x)$3

Because radiance transforms as $\mathcal{L}(x)$4, this makes

$\mathcal{L}(x)$5

So the geometric density of ray paths is directly proportional to the target likelihood. This is the stationarity mechanism.

The paper gives a local detailed-balance style derivation: emit rays isotropically from a differential volume around $\mathcal{L}(x)$6 with intensity proportional to $\mathcal{L}(x)$7,

$\mathcal{L}(x)$8

with $\mathcal{L}(x)$9 independent of position and direction. Using conservation of etendue between $\mathcal{L}(x)$0 and $\mathcal{L}(x)$1, one gets that the power density arriving at $\mathcal{L}(x)$2 is

$\mathcal{L}(x)$3

for the choice $\mathcal{L}(x)$4. Because ray paths are reversible, forward and reverse transitions satisfy detailed balance.

Resulting stationary distribution

For the unweighted constant-speed ray tracer, the stationary sampling density in parameter space is proportional to

$\mathcal{L}(x)$5

In the generalized weighted case, the samples may carry nonuniform weights $\mathcal{L}(x)$6, but the effective weighted stationary measure remains the desired target $\mathcal{L}(x)$7.

---

4. Ray dynamics: constant-speed propagation

The standard ray equation used is

$\mathcal{L}(x)$8

where $\mathcal{L}(x)$9 is path length. Rewriting yields the evolution of the unit direction vector:

$D$0

This is important:

• only the gradient of $D$1 matters,
• only the component of $D$2 perpendicular to the direction of motion bends the ray.

If $D$3 is the angle between the direction of travel and $D$4, then

$D$5

The paper integrates this exactly over one kick step to maintain reversibility:

$D$6

Constant speed vs HMC

This is the paper’s main dynamical distinction from HMC.

• Ray tracing: rays move at constant speed through parameter space; the likelihood only bends the trajectory.
• HMC: a particle accelerates/decelerates under the potential $D$7; both direction and speed change because kinetic and potential energy exchange.

The paper emphasizes that constant speed means stochastic gradient errors perturb only the direction, not the kinetic energy. That is the basis for its robustness claims.

---

5. Acceptance rule and exactness with imperfect integrators

For an imperfect numerical integrator, the path will not preserve basic radiance exactly. The paper derives a local radiance boost for a direction update $D$8:

$D$9

Accumulating this over a trajectory yields a predicted ratio $n^{D-1}$0. Then the Metropolis-like correction is

$n^{D-1}$1

In the generalized weighted framework, this becomes

$n^{D-1}$2

This is exactly the inverse ratio of excess basic radiance at the endpoint, and it restores exact sampling provided:

1. the path map is reversible,
2. the Metropolis correction is applied,
3. the chain is ergodic.

---

6. Generalized ray tracing with weights and variable speed

The framework is generalized by allowing explicit sample weights $n^{D-1}$3 and path speed $n^{D-1}$4.

The paper states that the effective sample density scales as

$n^{D-1}$5

To preserve fair sampling, this must equal the target likelihood:

$n^{D-1}$6

Since $n^{D-1}$7, one obtains the generalized refractive index

$n^{D-1}$8

This is the master equation of the generalized framework.

Interpretation:

• changing $n^{D-1}$9 changes path geometry,
• changing $n(x) = \mathcal{L}(x)^{\frac{1}{D-1}}.$0 changes how much time is spent along the path,
• changing $n(x) = \mathcal{L}(x)^{\frac{1}{D-1}}.$1 changes the weight assigned to visited points.

This separation is what lets the framework subsume HMC, microcanonical HMC, Gibbs, Metropolis, Monte Carlo integration, and tempered variants.

---

7. Ergodicity, scattering, barriers, and holes

Ergodicity

Ray propagation by itself is deterministic given initial direction. To guarantee exploration, the direction must be refreshed. The paper proposes periodic full or partial momentum refreshes.

Partial momentum refresh / scattering

A partial refresh is treated as a random rotation of direction, preserving Jacobian. The paper recommends performing it between completed integration steps. If no refresh occurs between Metropolis tests, failed proposals should reverse direction to maintain reversibility.

Crossing likelihood barriers

A major claim is that ray tracing can cross arbitrary likelihood barriers, unlike ordinary HMC at fixed energy.

Why:

• In basic ray tracing, the speed is fixed and only curvature changes.
• There is no fixed energy ceiling preventing access to low-likelihood regions.
• Rays can pass through regions of arbitrarily low likelihood; they are merely refracted according to $n(x) = \mathcal{L}(x)^{\frac{1}{D-1}}.$2.

This is contrasted with HMC, where a fixed total energy forbids entering regions with $n(x) = \mathcal{L}(x)^{\frac{1}{D-1}}.$3.

Crossing holes / disallowed regions

The paper further claims ray tracing can cross “holes” or disconnected allowed regions if boundaries are treated as blackbody radiators:

• boundaries absorb incoming rays,
• then re-emit them with Lambertian emissivity proportional to $n(x) = \mathcal{L}(x)^{\frac{1}{D-1}}.$4.

This gives a mechanism to move between disconnected components (“parameter space islands”), something standard Hamiltonian integration cannot do.

Conditions for correct sampling

Correct sampling requires:

1. reversibility of the path integrator,
2. detailed balance via conserved basic radiance or Metropolis correction,
3. ergodicity via direction refresh/scattering,
4. either connected support, or special blackbody-boundary handling for islands.

---

8. Why stochastic gradients affect ray tracing differently from HMC

The paper’s main practical claim is much stronger resilience to stochastic gradients.

HMC problem

Stochastic gradients inject energy into the trajectory, causing path heating. In SGHMC, this requires friction/diffusion terms tuned to the noise covariance.

Ray tracing claim

Because ray tracing keeps path speed fixed, stochastic gradient errors do not accumulate as kinetic-energy heating. They mainly perturb direction, which is actually similar to the deliberate scattering needed for ergodicity.

The residual issue is wrong acceptance probabilities due to noisy likelihood estimates, not runaway heating.

The paper proposes a stochastic acceptance correction:

$n(x) = \mathcal{L}(x)^{\frac{1}{D-1}}.$5

where $n(x) = \mathcal{L}(x)^{\frac{1}{D-1}}.$6 is the variance of $n(x) = \mathcal{L}(x)^{\frac{1}{D-1}}.$7.

But the paper also notes that when stochastic likelihood noise dominates integration error, a Metropolis test can become meaningless; then threshold/masking strategies are more useful.

---

9. Relationship to prior methods under the generalized framework

This is one of the paper’s most interesting claims: many samplers are recovered by choosing $n(x) = \mathcal{L}(x)^{\frac{1}{D-1}}.$8, $n(x) = \mathcal{L}(x)^{\frac{1}{D-1}}.$9, and $W(x)$0 appropriately.

9.1 Hamiltonian Monte Carlo as ray tracing

For HMC with potential $W(x)$1, and specific kinetic energy $W(x)$2, total energy conservation implies

$W(x)$3

Geometrically, the transverse component of velocity is conserved across potential jumps:

$W(x)$4

which has Snell-law form. Hence HMC paths are identical to ray paths with

$W(x)$5

and specifically

$W(x)$6

So HMC is a generalized ray-tracing method where:

• refractive index depends on energy through $W(x)$7,
• speed is not constant,
• sample weighting is effectively inverse velocity.

At low likelihood where $W(x)$8, $W(x)$9, so HMC has an impassable barrier. This explains geometrically why ordinary HMC cannot cross arbitrary likelihood barriers.

The paper also notes Langevin methods follow because they can be written as single-step HMC.

9.2 Microcanonical HMC

The paper identifies microcanonical HMC as the choice

$v(x)$0

which leads to

$v(x)$1

This is close to standard ray tracing $v(x)$2 in large $v(x)$3, so the path dynamics are similar asymptotically. The paper’s distinction is that ray tracing decouples speed from likelihood, which improves stochastic-gradient robustness.

9.3 Gibbs sampling

In this framework, Gibbs sampling corresponds to

$v(x)$4

Since $v(x)$5 is constant:

• rays are straight lines,
• $v(x)$6,
• the phase-space Jacobian is trivial.

The weights $v(x)$7 then enforce the conditional distribution along the chosen coordinate/subspace.

9.4 Metropolis(-Hastings)

A Metropolis step is described as:

• a single ray-tracing step,
• integrated by Euler,
• with $v(x)$8,
• $v(x)$9,
• and proposal length $n(x)$0 drawn from the chosen proposal distribution (often Gaussian).

So again the path is straight, and the usual accept/reject rule is recovered as the weighted-sampling correction.

9.5 Monte Carlo integration

Monte Carlo integration appears as a limiting case of Gibbs sampling with masked samples:

• weights are zero everywhere along the path,
• except at ray endpoints / refresh points, where they are $n(x)$1.

So raw Monte Carlo integration is a degenerate straight-ray weighted sampler.

9.6 Goodman–Weare stretch move

The paper also sketches an interpretation of the stretch move:

• $n(x)$2 outside a narrow cone from target walker to current walker,
• inside the cone, weighting $n(x)$3 for $n(x)$4. This is less central than the other correspondences but supports the unification claim.

---

10. Tempering inside the generalized framework

Because path geometry and weighting are separated, one can define a tempered family while still sampling the original target correctly.

For effective temperature $n(x)$5, the paper gives

$n(x)$6

So:

• increasing $n(x)$7 flattens the refractive-index landscape,
• low-likelihood regions are traversed faster / with lower weight.

If $n(x)$8 depends on $n(x)$9, then

$D$00

This preserves alignment with the $D$01 gradient direction up to a scalar.

---

11. Algorithmic implementation

Inputs

The pseudocode GRTSample takes roughly:

• initial point $D$02,
• dimension $D$03,
• number of integration steps $D$04,
• path step size $D$05,
• refresh parameter $D$06,
• target likelihood $D$07,
• optional weight function $D$08 (default 1).

Required gradient

Only the gradient of the log refractive index is needed. In the basic weighted constant-speed implementation:

$D$09

So in practice one needs $D$10 (or gradient of loss / log posterior), not the Hessian.

Update scheme

The pseudocode uses a drift-kick-drift leapfrog structure:

1. partial momentum refresh,
2. half drift:

$D$11

1. kick using UpdateV,
2. another half drift.

The velocity magnitude is irrelevant for the path because only $D$12 matters in drifting.

Partial refresh

The direction update is

$D$13

with fresh Gaussian $D$14, i.e. Ornstein–Uhlenbeck-style partial refresh.

Direction update

Inside UpdateV, one computes:

• $D$15,
• $D$16,
• initial angle $D$17,
• final angle $D$18 via the exact reversible formula.

Then the new velocity is

$D$19

where

$D$20

The cumulative radiance change is tracked as

$D$21

Metropolis test

At the end,

$D$22

triggers rejection.

Numerical integrators

The paper discusses:

• KDK and DKD leapfrog,
• random alternation between KDK and DKD,
• Omelyan second-order,
• Forest–Ruth / Yoshida fourth-order.

It finds Omelyan best over a broad practical range in a 10,000D Gaussian test.

Complexity

Per integration step:

• one gradient of $D$23 (or loss) at a midpoint,
• $D$24 vector operations.

So computational cost is similar in order to HMC: dominated by gradient evaluations. The key gain is not lower per-step cost but better behavior under noisy gradients and large-model minibatching.

---

12. Practical recommendations for neural networks

Approximate likelihood from a loss

When no proper probabilistic likelihood is available, the paper proposes

$D$25

where:

• $D$26 is the training loss,
• $D$27 is desired tolerance above optimum,
• $D$28 is an effective number of constrained parameters.

This is exact for Gaussian targets and suggested as a practical approximation for high-dimensional models.

Recipe

The paper’s recommended workflow:

1. burn in with low-temperature ray tracing or Adam,
2. choose likelihood scaling / $D$29,
3. tune step size (start around $D$30),
4. tune refresh rate,
5. run many walkers,
6. assess convergence in function space, not necessarily weight space.

High-dimensional caveats

• Neural-network posteriors in weight space are often filamentary and non-Gaussian.
• Convergence in weight space may be poor or absent even when function-space uncertainty is useful.
• The sampler is not generically separable; if subspaces have very different gradient scales, independent ray dynamics by subspace may help.

---

13. Experimental evidence relevant to the framework

13.1 Gaussian benchmarks

For a 10,000D Gaussian:

• with exact gradients, ray tracing and HMC have roughly similar performance when tuned well;
• Metropolis correction matters for imperfect integrators;
• Omelyan integrator performed best across a broad step-size range.

So the paper does not claim huge gains over exact-gradient HMC in benign settings.

13.2 Stochastic-gradient robustness

This is the headline empirical result.

For stochastic Gaussian gradients defined by

$D$31

the paper finds:

• ray tracing remains accurate up to $D$32,
• HMC degrades by $D$33.

The median-log-likelihood error obeys

$D$34

with empirical resilience

• HMC: $D$35,
• ray tracing: $D$36.

So ray tracing tolerated about $D$37 larger noise amplitude, or $D$38 larger variance, at comparable bias.

Scaling with step size:

• HMC robustness: $D$39,
• ray tracing robustness: $D$40.

This implies smaller minibatches can improve total efficiency for ray tracing, but not for HMC.

13.3 Neural-network applications

• MLP, 1433 params: both HMC and ray tracing work; similar posterior predictions, somewhat better autocorrelation for ray tracing.
• ResNet-34, 22M params: ray tracing feasible on one consumer GPU; HMC not attempted because minibatch stochasticity too high.
• GPT-2, 1.5B params: approximate posterior sampling on a single consumer GPU; stable validation loss around target; median autocorrelation time for sampled token probabilities around $D$41 tokens.

The GPT-2 experiment is presented as approximate rather than exact, but it supports the paper’s claim that the framework scales to very large neural networks under stochastic gradients.

---

14. Limitations and caveats

The paper is fairly candid about several limitations:

1. $D$42: the basic formula

$D$43

is undefined; the ray-tracing formulation is meaningful only for $D$44.

1. Need for refresh/scattering: deterministic ray tracing alone is not ergodic.
2. Disconnected support: requires special blackbody-boundary treatment for islands.
3. Approximate likelihoods for neural nets: practical results depend on choosing $D$45 and loss scaling, often heuristically.
4. Weight-space convergence may fail even when function-space uncertainty is useful.
5. Very large models: reported GPT-2 sampling is approximate and affected by BF16 quantization spikes; the paper uses masking to skip bad samples.
6. No universal speed advantage over exact-gradient HMC in clean, standard targets; the strongest advantage is specifically for stochastic-gradient settings.

---

15. Why the generalized framework matters

The significance of the framework is threefold.

(i) A new geometric route to correct sampling

Instead of preserving Hamiltonian energy, it preserves an optical invariant: $D$46 This gives a distinct acceptance rule and a natural fair-sampling construction.

(ii) A principled separation of geometry from weighting

The generalized formulation

$D$47

separates:

• path geometry $D$48,
• traversal speed $D$49,
• sample weight $D$50.

That makes the approach a family of samplers, not just a single one.

(iii) Unification of prior samplers

Under this lens:

• HMC is ray tracing with $D$51,
• microcanonical HMC uses $D$52,
• Gibbs uses $D$53,
• Metropolis is a one-step straight-ray weighted update,
• Monte Carlo integration is masked Gibbs sampling.

So the framework provides a common language for comparing how different samplers trade off exploration of high-probability peaks versus low-probability tails.

---

16. Bottom-line summary

The paper proposes a sampling framework based on optical ray propagation in parameter space. The fair-sampling version chooses

$D$54

so that conservation of optical radiance implies ray density $D$55. Ray paths evolve according to

$D$56

or equivalently

$D$57

with constant path speed and reversible angle update

$D$58

Correctness is enforced through conservation of the “basic radiance”

$D$59

with Metropolis correction

$D$60

or in generalized form

$D$61

The generalized framework introduces weights and speeds: $D$62 which lets many standard methods appear as special cases.

Its main practical claim is that constant-speed rays are vastly more robust than HMC to stochastic-gradient heating, which the experiments support strongly. That robustness is what enables approximate posterior sampling for very large neural networks, including GPT-2-scale models, on modest hardware.