🇬🇧 English · 🇪🇸 Español

Parameter Gradients via the Implicit Function Theorem

puremacro.vfi.analytic_gradients provides Jacobians of representative-agent continuous projection solutions, and of the deterministic steady-state aggregates they imply, with respect to structural model parameters \(\theta = (\beta, \alpha, \delta, \sigma, \dots)\) using the Implicit Function Theorem (IFT). It operates uniformly across Chebyshev orthogonal polynomial collocation, Finite Element Method (FEM) Galerkin projections, and Cubic/Schumaker spline systems.

Scope (4.3.0 and later). The IFT step is one linear solve with the converged residual Jacobian (no re-solve of the model); the residual Jacobians \(\mathbf{J}_c\) and \(\mathbf{J}_\theta\) it consumes are built by central finite differences with steps step_c and h (a custom residual_fn replaces the built-in residual evaluator and is differenced the same way). Only representative-agent solutions are supported: compute_ift_gradients and equilibrium_parameter_jacobian raise NotImplementedError for any solution that carries a distribution attribute (for example AiyagariContinuousEquilibrium). The adjoint stationary-distribution sensitivity described in section 1.3 is not implemented.

In structural econometric estimation (GMM, SMM) and Bayesian posterior sampling (HMC, NUTS), gradient evaluations of dynamic economic models have historically relied on numerical finite differences. However, finite differences suffer from a severe step-size dilemma (\(\epsilon \sim 10^{-5}\) balances truncation and round-off error), exhibit numerical chatter near non-linear borrowing constraints, and require \(2 \times \dim(\theta)\) full non-linear model re-solves. The puremacro analytic gradient engine removes the \(2 \times \dim(\theta)\) re-solves by computing the coefficient sensitivities from a single LU factorization of the converged residual Jacobian. The repository test suite asserts at least a 5x speedup over two-sided finite differences on the worked example of section 4; larger ratios are typical on bigger bases, but the timings quoted in section 3 are indicative rather than a contract.

---

1. Theoretical & Mathematical Foundation

1.1 The Implicit Function Theorem on Continuous Projection Systems

Consider a continuous dynamic programming problem approximated using an \(N\)-dimensional basis (Chebyshev polynomials, piecewise linear finite elements, or cubic B-splines). The equilibrium policy is parametrized by coefficient vector \(c^* \in \mathbb{R}^N\) satisfying a system of non-linear residual equations evaluated at collocation or quadrature nodes:

\[\mathbf{R}(c^*(\theta); \theta) = \mathbf{0} \in \mathbb{R}^N\]

where \(\theta \in \mathbb{R}^p\) is the vector of structural parameters (such as the subjective discount factor \(\beta\), capital share \(\alpha\), or depreciation rate \(\delta\)).

Assuming the residual operator \(\mathbf{R}\) is continuously differentiable with respect to both \(c\) and \(\theta\), the total derivative of the equilibrium identity with respect to \(\theta\) is:

\[\frac{d \mathbf{R}(c^*(\theta); \theta)}{d \theta} = \nabla_c \mathbf{R}(c^*; \theta) \cdot \nabla_\theta c^*(\theta) + \nabla_\theta \mathbf{R}(c^*; \theta) = \mathbf{0}\]

Provided that the coefficient Jacobian \(\mathbf{J}_c \equiv \nabla_c \mathbf{R}(c^*; \theta) \in \mathbb{R}^{N \times N}\) is non-singular at the converged solution \(c^*\), the Implicit Function Theorem guarantees that \(c^*(\theta)\) is locally differentiable, with parameter Jacobian:

\[\nabla_\theta c^*(\theta) = - \left[ \nabla_c \mathbf{R}(c^*; \theta) \right]^{-1} \nabla_\theta \mathbf{R}(c^*; \theta)\]

where: - \(\mathbf{J}_c = \nabla_c \mathbf{R} \in \mathbb{R}^{N \times N}\) is the Jacobian of the residual system with respect to basis coefficients. - \(\mathbf{J}_\theta = \nabla_\theta \mathbf{R} \in \mathbb{R}^{N \times p}\) is the direct sensitivity of the residual equations to the parameter vector.

1.2 Single LU Factorization vs Finite Differences

Rather than re-solving the non-linear dynamic programming problem \(2p\) times (as required by two-sided numerical finite differences), the IFT computes the entire \(N \times p\) coefficient sensitivity matrix by solving a linear system with \(p\) right-hand sides:

\[\mathbf{J}_c \cdot \mathbf{X} = - \mathbf{J}_\theta\]

Using an LU factorization \(\mathbf{J}_c = \mathbf{P} \mathbf{L} \mathbf{U}\): 1. Factorization: The \(N \times N\) matrix \(\mathbf{J}_c\) is factored once in \(\frac{2}{3} N^3\) floating-point operations. 2. Back-Substitution: For each parameter \(\theta_j\) (\(j = 1, \dots, p\)), the sensitivity column \(\nabla_{\theta_j} c^*\) is obtained via forward and back substitution in \(2 N^2\) operations.

This reduces the computational complexity from \(O(2p \cdot N_{\text{iter}} \cdot N^3)\) for non-linear re-solving down to \(O(N^3 + p N^2)\), accelerating gradient evaluation by 1 to 2 orders of magnitude.

1.3 Policy Sensitivities & Adjoint Aggregate Jacobians

Because continuous projection approximations represent policy functions as linear combinations of basis functions \(g(k; c) = \sum_{j=1}^N c_j \phi_j(k) = \mathbf{\Phi}(k) c\), the continuous policy gradient at any arbitrary state \(k\) is obtained analytically:

\[\nabla_\theta g(k; \theta) = \mathbf{\Phi}(k) \cdot \nabla_\theta c^*(\theta) \in \mathbb{R}^{1 \times p}\]

General Equilibrium Macroeconomic Aggregates

In a representative-agent continuous projection economy, the deterministic steady-state capital stock \(k^*\) satisfies the fixed-point condition \(g(k^*; \theta) = k^*\). Differentiating implicitly yields:

\[\nabla_\theta k^* = \left( 1 - \frac{\partial g(k^*)}{\partial k} \right)^{-1} \nabla_\theta g(k^*)\]

Equilibrium factor prices \(r^* = \alpha z (k^*)^{\alpha - 1} - \delta\) and \(w^* = (1 - \alpha) z (k^*)^\alpha\) have sensitivities:

\[\nabla_\theta r^* = \frac{\partial r^*}{\partial k} \nabla_\theta k^* + \left. \nabla_\theta r^* \right|_{\text{direct}}, \quad \nabla_\theta w^* = \frac{\partial w^*}{\partial k} \nabla_\theta k^* + \left. \nabla_\theta w^* \right|_{\text{direct}}\]

These representative-agent steady-state sensitivities are what grad_aggregates returns (keys K, C, r, w). They are semi-analytic: the steady-state fixed point is located numerically and its derivative is then propagated analytically through the formulas above.

Not implemented: heterogeneous-agent economies. In incomplete-markets economies (Aiyagari 1994) aggregate capital supply is \(K^* = \int k \, d\mu^*(k; \theta)\) and its sensitivity would have to propagate through the adjoint of the stationary-distribution (lottery) operator,

\[\nabla_\theta K^* = \sum_{i, m} k_i \nabla_\theta \mu^*(k_i, z_m) = \mathbf{k}^\top (\mathbf{I} - \mathbf{T}^*)^{-1} \nabla_\theta \mathbf{T}^* \boldsymbol{\mu}^*\]

together with the general-equilibrium price response. puremacro 4.3.0 removed the earlier fixed-distribution shortcut because it was not a general-equilibrium sensitivity; since then compute_ift_gradients, equilibrium_parameter_jacobian and the internal aggregate routine raise NotImplementedError for any solution with a distribution attribute. The formula is recorded here as the target of future work, not as a shipped feature.

1.4 Fast Structural Estimation (SMM & GMM)

Structural estimation seeks parameter vector \(\hat{\theta}\) minimizing the weighted distance between simulated model moments \(m(\theta)\) and empirical target moments \(\hat{m}\):

\[Q(\theta) = \left( m(\theta) - \hat{m} \right)^\top \mathbf{W} \left( m(\theta) - \hat{m} \right)\]

Applying the chain rule, the gradient of the GMM objective is:

\[\nabla_\theta Q(\theta) = 2 \left[ \nabla_\theta m(\theta) \right]^\top \mathbf{W} \left( m(\theta) - \hat{m} \right)\]

where \(\nabla_\theta m(\theta) = \nabla_c m(c^*) \nabla_\theta c^* + \partial_\theta m\). Because \(\nabla_\theta c^*\) comes from one linear solve of the converged system rather than from \(2p\) noisy re-solves, gradient-based quasi-Newton optimizers (L-BFGS-B, SLSQP) converge reliably in a fraction of the time required under finite differences.

---

2. Methodological & Solver Options

Feature	Single LU (solver="lu")	Tikhonov Regularization (solver="tikhonov")	Truncated SVD (solver="svd")
Mathematical Basis	\(\mathbf{P} \mathbf{L} \mathbf{U} = \mathbf{J}_c\)	\((\mathbf{J}_c^\top \mathbf{J}_c + \lambda \mathbf{I})^{-1} \mathbf{J}_c^\top\)	\(\mathbf{V} \mathbf{\Sigma}^+ \mathbf{U}^\top\) with cutoff \(\sigma_i > \epsilon \sigma_1\)
Condition Threshold	\(\text{cond}(\mathbf{J}_c) \le 10^{12}\)	\(10^{12} < \text{cond}(\mathbf{J}_c) \le 10^{15}\)	Singular or ill-posed systems
Accuracy	Backward-stable LU solve (error of order \(\text{cond}(\mathbf{J}_c) \cdot \epsilon_{\text{mach}}\)) on top of the central-difference error of \(\mathbf{J}_c\) and \(\mathbf{J}_\theta\)	Damped regularized gradient	Filtered subspace gradient
Cost	\(\frac{2}{3} N^3\) flops	\(O(N^3)\) flops	\(O(N^3)\) (full singular value spectrum)
Selection	Automatic default ("auto")	Triggered on ill-conditioned bases	High-degree collinear polynomials

---

3. Comparative Benchmarks & Precision Analysis

Dimension	IFT gradients (puremacro)	Two-Sided Finite Differences
Gradient Accuracy	Backward-stable IFT solve given \(\mathbf{J}_c\), \(\mathbf{J}_\theta\) (themselves central differences with steps step_c, h); the shipped test checks agreement with a finite-difference re-solve to relative \(10^{-5}\)	Discretization error \(O(\epsilon^2) \approx 10^{-5}\) plus re-solve tolerance noise
Step-Size Dilemma	Reduced: steps step_c, h difference the residual at the fixed converged coefficients, with no re-solve inside the difference	Requires careful tuning (\(\epsilon = 10^{-4}\) vs \(10^{-6}\))
Model Evaluations	\(1\) (converged steady state)	\(2 \times p\) full non-linear re-solves
Execution Time (\(p = 4\))	one Jacobian build and one LU solve (indicative: milliseconds on the section 4 example)	\(2p\) non-linear re-solves (the shipped test asserts the IFT path is at least 5x faster)
Numerical Chatter	No re-solve tolerance noise (the residual is differenced at fixed coefficients); round-off of the residual differences remains	Severe round-off noise near borrowing kinks
Optimizer Stability	Robust Hessian updates (BFGS/L-BFGS)	Spurious gradient reversals trigger premature termination

---

4. Runnable Worked Examples

The following script solves a neoclassical growth model via Chebyshev collocation, computes IFT parameter sensitivities, and evaluates macroeconomic aggregate derivatives:

import numpy as np
from puremacro.vfi.collocation import CollocationProblem
from puremacro.vfi.analytic_gradients import (
    compute_ift_gradients,
    policy_parameter_jacobian,
    equilibrium_parameter_jacobian,
)

# 1. Solve Continuous Projection Problem (Chebyshev Collocation)
alpha, beta, delta, sigma = 0.36, 0.96, 1.0, 1.0
k_ss = float((alpha * beta) ** (1.0 / (1.0 - alpha)))

def euler_res(k, kp, kpp, s, sp, beta=beta):
    c = k**alpha - kp
    c_next = kp**alpha - kpp
    R = alpha * kp**(alpha - 1.0)
    return 1.0 / c - beta * (1.0 / c_next) * R

prob = CollocationProblem(
    domain=(0.5 * k_ss, 1.5 * k_ss),
    orders=6,
    method="euler",
    params={"alpha": alpha, "delta": delta, "sigma": sigma},
    beta=beta,
    options={"tol": 1e-12, "max_iter": 500},
)
sol = prob.solve(backend="numpy")

# 2. IFT Jacobians (J_c, J_theta by central differences; one LU solve)
ift_res = compute_ift_gradients(sol, prob, params=["alpha", "beta", "delta"])
assert ift_res.condition_number < 1e6
assert ift_res.grad_coefficients.shape == (7, 3)

# 3. Continuous Policy Gradient at Steady State
d_pol = ift_res.policy_gradient(k_ss)
assert d_pol[1] > 0.0  # d(k') / d(beta) > 0: patient households accumulate more capital

# 4. Macroeconomic Aggregate Sensitivities
grad_aggs = ift_res.grad_aggregates
print(f"Condition number cond(J_c): {ift_res.condition_number:.2e}")
print(f"d(K*)/d(beta): {grad_aggs['K'][1]:.4f}, d(r*)/d(beta): {grad_aggs['r'][1]:.4f}")

---

5. Full API Specification

compute_ift_gradients(
    solution: Any,
    problem: Any,
    params: Sequence[str] | None = None,
    h: float = 1e-5,
    step_c: float = 1e-6,
    residual_fn: Callable | None = None,
    cond_max: float = 1e12,
    backend: str = "numpy",
    **kwargs: Any,
) -> AnalyticGradientResult

policy_parameter_jacobian(
    solution: Any,
    problem: Any,
    s: float | Sequence[float] | np.ndarray,
    params: Sequence[str] | None = None,
    h: float = 1e-5,
    ift_result: AnalyticGradientResult | None = None,
    **kwargs: Any,
) -> np.ndarray

equilibrium_parameter_jacobian(
    solution: Any,
    problem: Any | None = None,
    params: Sequence[str] | None = None,
    h: float = 1e-5,
    ift_result: AnalyticGradientResult | None = None,
    **kwargs: Any,
) -> dict[str, np.ndarray]

gmm_objective_and_gradient(
    theta_vals: Sequence[float],
    problem: Any,
    empirical_moments: Sequence[float],
    moment_fn: Callable | None = None,
    weighting_matrix: np.ndarray | None = None,
    param_names: Sequence[str] | None = None,
    backend: str = "numpy",
    **kwargs: Any,
) -> tuple[float, np.ndarray]

Parameters:

• solution: Converged representative-agent solution object from CollocationProblem.solve(), FEMProblem.solve() or SplineCollocationProblem.solve(). Heterogeneous-agent solutions (any object with a distribution attribute, such as AiyagariContinuousEquilibrium) raise NotImplementedError.
• problem: Associated economic problem definition.
• params: Sequence of parameter names to differentiate (e.g. ['alpha', 'beta', 'delta']). If None, automatically inspects model parameters.
• h: Step size for parameter differentiation \(\nabla_\theta \mathbf{R}\) (default \(10^{-5}\)).
• step_c: Step size for coefficient Jacobian \(\mathbf{J}_c = \nabla_c \mathbf{R}\) (default \(10^{-6}\)).
• cond_max: Condition number threshold for triggering regularized linear solves (default \(10^{12}\)).
• s: Continuous state coordinates at which to evaluate \(\nabla_\theta g(s)\).
• theta_vals: Candidate parameter vector in structural estimation.
• empirical_moments: Target empirical moment vector \(\hat{m}\).

---

6. Result Interface & Manuscript Export

AnalyticGradientResult stores the full sensitivity matrices and diagnostic metadata:

Attributes

• grad_coefficients: \(N \times p\) Jacobian matrix \(\nabla_\theta c^*\) of basis coefficients.
• grad_aggregates: Dictionary of representative-agent steady-state aggregate derivatives {"K": (p,), "C": (p,), "r": (p,), "w": (p,), "mu": None} (semi-analytic, section 1.3). The "mu" entry is always None: no stationary-distribution sensitivity is computed.
• param_names: Ordered list of differentiated structural parameter names.
• jacobian_resid_c: Residual Jacobian \(\mathbf{J}_c = \nabla_c \mathbf{R}\) of shape \((N, N)\).
• jacobian_resid_theta: Parameter Jacobian \(\mathbf{J}_\theta = \nabla_\theta \mathbf{R}\) of shape \((N, p)\).
• condition_number: 2-norm condition number of \(\mathbf{J}_c\).
• elapsed_time: Wall-clock evaluation time in seconds.

Methods

• res.policy_gradient(s) -> np.ndarray: Evaluates continuous policy sensitivity \(\nabla_\theta g(s)\) at scalar or array state coordinates \(s\).
• res.summary() -> pd.DataFrame: Summary table displaying parameter sensitivities of capital, consumption, interest rates, and wages.
• res.to_frame() -> pd.DataFrame: Alias returning the summary table.
• res.to_markdown(digits=4) -> str: GitHub-flavored Markdown table.
• res.to_latex(digits=4) -> str: Publication-grade LaTeX tabular environment.
• res.to_typst(digits=4) -> str: Formatted Typst table.
• res.plot(figsize=(10, 4.5)) -> matplotlib.figure.Figure: Multi-panel visualization displaying the singular value spectrum of \(\mathbf{J}_c\), continuous policy gradients across the state space, and aggregate sensitivities.

---

References

• Judd, K. L. (1998). Numerical Methods in Economics. MIT Press.
• Maliar, L., & Maliar, S. (2014). "Numerical methods for large-scale dynamic economic models." In Handbook of Computational Economics (Vol. 3, pp. 325–477). Elsevier.
• Rust, J. (1994). "Structural estimation of Markov decision processes." In Handbook of Econometrics (Vol. 4, pp. 3081–3143). North-Holland.