Analytical Gradient Investigation

1. One-Paragraph Summary

We estimate a dynamic discrete-choice model (DDCM) by NFXP: an optimizer searches over parameters θ, and for each θ we solve the agent's full day by GPU-accelerated backward induction and compare it to observed travel diaries. To make the optimizer converge efficiently we need the gradient of the log-likelihood. No NFXP estimation run with the analytical gradient ever truly converged — the likelihood went flat after a few iterations and the parameter mu_home jammed against its bound every time. The cause was a wrong gradient. After eliminating a dozen suspects, we found it: the analytical gradient is exactly correct when processing one home-zone at a time and 38% wrong when batching multiple zones together on the Apple GPU (MPS) backend. The math, the features, and the recursion are all right; the defect is a backend implementation issue — most likely an int32 integer overflow at real graph scale, the same family as three prior MPS-specific failures in this codebase. CUDA has now been confirmed clean (batch=8 == batch=1 == FD, 0.0% difference), so the production estimation launches immediately with no workaround. The MPS issue is a local-development matter only.

2. The Story — Why It Was So Hard to Find

Why we have two devices

The production estimator runs on a Linux server with a single NVIDIA GPU (32 GB), shared with other lab members. Queue times, memory limits, and long iteration times made local debugging painful. The Mac on the desk has Apple Silicon with large unified memory (192 GB shared CPU+GPU). Moving local development to MPS gave us zero queuing and no memory cliffs. The design contract was byte-compatibility: same model, same parameters, same answer on CUDA, MPS, and CPU.

The symptom

Every analytical-gradient estimation run showed the same fingerprint:

• The optimizer takes a few steps, the likelihood improves a little, then goes flat.
• mu_home slides to its lower bound and jams there.
• BFGS reports "convergence" at iteration 2; L-BFGS-B stalls; gradient-free BOBYQA gave no better results.
• The final parameters are not credible.

What we tried that didn't work

We spent a long time treating this as an optimization problem and an identification problem:

• Different warm-starts → no effect
• Tighter/looser parameter bounds → no effect
• Workers-only vs. mixed sample → no effect
• Different optimizers (BFGS, L-BFGS-B, BOBYQA, CMA-ES) → no effect
• Switching between FD gradient and analytical gradient → FD had its own problem near convergence (GPU noise)

None helped, because the cause was never the landscape or the optimizer. When every knob changes nothing, the thing you're turning is not the problem.

Why the bug stayed invisible

Three layers of concealment:

The break

3. The Decisive Evidence

FD (ground truth)         : +704.60

MPS  zone_batch_size = 1 : +704.61  ( 0.0%)  ✓ exact
MPS  zone_batch_size = 8 : +436.24  (38.1%)  ✗ WRONG
CPU  zone_batch_size = 8 : +704.61  ( 0.0%)  ✓ exact

CUDA confirmation

Because the production server runs CUDA (NVIDIA GPU), we immediately ran the same test there (workers group, 8 home zones, same script):

FD (ground truth)          : +942.62
CUDA zone_batch_size = 1  : +942.68  ( 0.0%)  ✓ exact
CUDA zone_batch_size = 8  : +942.67  ( 0.0%)  ✓ exact

batch=8 vs batch=1: 0.0% difference — CUDA is FINE

Why a 38% GV error flips the sign of mu_home

Near a reasonable parameter value, the two halves of the gradient score nearly cancel:

A 38% error in the GV-part means the computed value is −14,900 instead of −24,260. The computed "net" becomes +23,090 − 14,900 = +8,190 — wrong sign. The optimizer then walks away from the optimum on mu_home, hits the bound, and jams. This near-cancellation is a structural property of any well-specified likelihood at a reasonable parameter point (near the optimum, obs_feat and GV-part cancel exactly, since ∂ℓ/∂θ = 0 is the first-order condition). The tighter the cancellation, the more sensitive the sign is to errors in GV.

4. What the Mechanism Is

First hypothesis: MPS operator bug

When processing Z_b home zones simultaneously, the gradient accumulation uses a 3-D tensor operation over shape (N_states, Z_b, P_params):

# GV_batch: (N, Z_b, P)
GV_dst = GV_batch[edge_tgt]           # gather destinations
GV_dst.add_(F_sl.unsqueeze(1))        # add features (broadcast over zones)
GV_dst.mul_(policy_norm.unsqueeze(2)) # weight by policy (broadcast over params)
GV_batch.index_add_(0, src_idx, GV_dst)  # ← suspected op

Our first guess was that MPS computes index_add_ wrong on 3-D tensors. But micro-tests showed MPS handles all these operations correctly at small scale.

Refined hypothesis: int32 overflow at scale (current best)

The corruption appears only at real graph scale. At timestep t with E edges, Z_b=8 zones, P=11 params, the flattened element count is:

This crosses 2³¹ ≈ 2.15 billion (the int32 limit) when E > ~24M edges per timestep. At Z_b=1 the count is 8× smaller — stays under the limit — which is exactly why batch=1 is exact and batch=8 is corrupted. CPU and CUDA use int64 indexing internally and never overflow. This is the same family as two prior MPS-specific failures in this codebase:

5. Why the Analytical Gradient Should Equal FD

The analytical gradient is the exact derivative of the same log-likelihood that FD approximates. The published literature is unambiguous about this (see §8 for references). The key results from the literature:

• Rust (1987) NFXP manual: analytic first and second derivatives of the likelihood are standard, obtained by differentiating the Bellman/EV equation. Supplying analytic gradients — rather than numeric — is what makes the optimizer converge reliably.
• Eberwein & Ham (2008): analytic derivatives of the NFXP likelihood can be obtained recursively at the same order of complexity as the likelihood itself (one extra backward sweep). In our finite-horizon case, the linear system they describe becomes a backward recursion — which is exactly our gradient_bi_and_obs_features.
• Aguirregabiria & Mira / Abbring & Klein: the score takes the "observed minus expected feature" form. The policy derivative never appears (envelope theorem). Near the optimum, FD is the less reliable reference because it amplifies GPU float32 noise; the analytic score is better conditioned there.

A correct analytical gradient must match FD to within ~0.1–1% (FD truncation + float32 noise floor). A 17–38% disagreement is unambiguously a bug. And at zone_batch_size=1 we measure 0.0% — FD-exact — confirming the math is right.

6. All Bugs Found and Fixed

7. Current Status and Next Steps

Main track — production estimation

1. ✅ [Done] CUDA zone-batch test — 0.0% diff. Bug is MPS-only.
2. → [Next] Validate all 11 params with cuda_debug_fd2.py on workers multi-zone group at zone_batch_size=8. One final sanity check.
3. → [Then] Launch full --analytical-gradient estimation on CUDA, workers-only, L-BFGS-B, warm-start from best checkpoint.
4. → [After convergence] Compute BHHH standard errors and assess behavioral validity.

Side track — MPS local-dev fix (non-blocking)

1. Confirm int32-overflow hypothesis with investigations/check_edge_chunk_fix.py.
2. Apply edge-chunking so per-chunk size stays under 2.15B — keeps computation on GPU, minimal slowdown.
3. Add regression guard: GV(batch=1) == GV(batch=N) on every device.

The key lesson

The deepest lesson from this investigation: a wrong gradient masquerades as a hard optimization landscape. "Flat LL, parameter jammed at its bound" was interpreted as an identification or optimizer problem. It was a gradient bug. When every optimizer and warm-start behaves the same way, suspect the gradient. And the defense against silent correctness failures is always the same: an independent oracle (CPU / FD) and varying one knob at a time. The day we compared batch=1 vs batch=8 against FD, the problem resolved in a single line of output.

8. References

1. Rust, J. (1987). "Optimal Replacement of GMC Bus Engines." Econometrica 55(5):999–1033.
2. Rust, J. (2000). NFXP Documentation Manual, Version 6. editorialexpress.com/jrust/nfxp.pdf
3. Eberwein, C. & Ham, J. (2008). "Obtaining analytic derivatives for a popular discrete-choice dynamic programming model." Economics Letters 101(3):170–172. — Analytic likelihood derivatives at the same cost as the likelihood; value-function derivative via backward recursion in finite horizon.
4. Aguirregabiria, V. & Mira, P. (2002). "Swapping the Nested Fixed Point Algorithm." Econometrica 70(4):1519–1543.
5. Aguirregabiria, V. & Mira, P. (2010). "Dynamic discrete choice structural models: A survey." Journal of Econometrics 156(1):38–67. — Score in observed-minus-expected-feature form.
6. Abbring, J. & Klein, T. Dynamic Discrete Choice teaching materials. ddc.abbring.org

Supporting code & scripts: investigations/check_gv_order.py · investigations/check_zone_batch_device.py · cuda_check_zone_batch.py · estimation/analytical_se.py