Week 5: Deep Equilibrium Networks II — Stochastic Models

Deep Learning for Macroeconomics — Honours, The University of Edinburgh
Instructor: Juan Zurita · juan.zurita@ed.ac.uk

---

Learning objectives

By the end of this notebook you will be able to:

1. Extend the DEQN framework to models with stochastic productivity shocks
2. Use Gauss–Hermite quadrature to compute expectations in the Euler equation
3. Compare different loss functions (MSE, Huber, log-cosh) for DEQN training
4. Simulate the solved model and compute impulse response functions

---

# Setup
import numpy as np
import matplotlib.pyplot as plt
import torch
import torch.nn as nn

torch.manual_seed(42)
np.random.seed(42)
device = 'cuda' if torch.cuda.is_available() else 'cpu'
print(f"PyTorch {torch.__version__} on {device}")

1. Adding uncertainty to the growth model

The deterministic Brock–Mirman model from Week 4 is a useful testbed, but real economies have shocks. The stochastic version adds a productivity shock \(z_t\):

\[\max \mathbb{E}_0 \sum_{t=0}^{\infty} \beta^t u(c_t) \quad \text{s.t.} \quad k_{t+1} = z_t k_t^\alpha - c_t\]

where \(\ln z_t = \rho \ln z_{t-1} + \varepsilon_t\), \(\varepsilon_t \sim N(0, \sigma^2)\).

The state space is now two-dimensional: \((k_t, z_t)\). The policy network maps \((k, z) \to c\).

The stochastic Euler equation

\[u'(c_t) = \beta \, \mathbb{E}_t \big[ u'(c_{t+1}) \, \alpha z_{t+1} k_{t+1}^{\alpha-1} \big]\]

The expectation is over tomorrow’s shock \(z_{t+1}\), conditional on today’s \(z_t\).

2. Gauss–Hermite quadrature for expectations

We need to compute \(\mathbb{E}_t[g(z_{t+1})]\) where \(z_{t+1}\) depends on \(z_t\) and a normal shock. Gauss–Hermite quadrature approximates this integral using a weighted sum:

\[\mathbb{E}[g(\varepsilon)] \approx \sum_{j=1}^{Q} w_j \, g(\varepsilon_j)\]

where \((\varepsilon_j, w_j)\) are the quadrature nodes and weights.

# Gauss-Hermite quadrature

from numpy.polynomial.hermite import hermgauss

# 5-point quadrature (sufficient for smooth integrands)
Q = 5
nodes, weights = hermgauss(Q)

# Adjust for standard normal: ε ~ N(0, σ²)
sigma_eps = 0.02
quad_nodes = nodes * np.sqrt(2) * sigma_eps
quad_weights = weights / np.sqrt(np.pi)

print("Quadrature nodes and weights:")
for j in range(Q):
    print(f"  ε_{j+1} = {quad_nodes[j]:+.4f},  w_{j+1} = {quad_weights[j]:.4f}")
print(f"\nWeights sum to: {quad_weights.sum():.6f} (should be 1.0)")

# Verify: E[ε²] should equal σ²
E_eps2 = np.sum(quad_weights * quad_nodes**2)
print(f"E[ε²] = {E_eps2:.6f}, σ² = {sigma_eps**2:.6f}")

3. The stochastic DEQN

The policy network now takes two inputs: \(c = \text{NN}(k, z; \theta)\).

The loss function replaces the deterministic Euler equation with its expectation:

\[\mathcal{L}(\theta) = \mathbb{E}_{k,z} \left[ \left( u'(c_t) - \beta \sum_{j=1}^Q w_j \, u'(c_{t+1}^j) \, \alpha z_{t+1}^j (k')^{\alpha-1} \right)^2 \right]\]

where \(z_{t+1}^j = z_t^\rho \exp(\varepsilon_j)\) for each quadrature node \(j\).

# Stochastic DEQN implementation

RHO = 0.9      # persistence of productivity
SIGMA = 0.02   # std dev of shock
ALPHA = 0.33
BETA = 0.95

# Convert quadrature to torch
quad_nodes_t = torch.tensor(quad_nodes, dtype=torch.float32).to(device)
quad_weights_t = torch.tensor(quad_weights, dtype=torch.float32).to(device)

class StochasticPolicyNet(nn.Module):
    """Maps (k, z) → c, ensuring 0 < c < z·k^α."""
    def __init__(self, hidden=64):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(2, hidden), nn.Tanh(),
            nn.Linear(hidden, hidden), nn.Tanh(),
            nn.Linear(hidden, 1)
        )
    
    def forward(self, k, z):
        x = torch.cat([k, z], dim=1)
        raw = self.net(x)
        share = torch.sigmoid(raw)
        output = z * k ** ALPHA
        return share * output * 0.99

def stochastic_euler_loss(policy_net, k, z):
    """Euler-equation loss with stochastic shocks via quadrature."""
    c_t = policy_net(k, z)
    k_next = z * k ** ALPHA - c_t
    
    # Expected marginal utility tomorrow (quadrature)
    rhs = torch.zeros_like(c_t)
    for j in range(len(quad_nodes_t)):
        ln_z_next = RHO * torch.log(z) + quad_nodes_t[j]
        z_next = torch.exp(ln_z_next)
        c_next = policy_net(k_next, z_next)
        mpk_next = ALPHA * z_next * k_next ** (ALPHA - 1)
        rhs += quad_weights_t[j] * BETA * mpk_next / c_next
    
    lhs = 1.0 / c_t
    residual = lhs - rhs
    return torch.mean(residual ** 2)

print("Stochastic DEQN ready. State space: (k, z) → c")
print(f"Quadrature points: {Q}")

4. Loss function design

Different loss functions penalise residuals differently. For economic models, the choice matters:

Loss	Formula	Behaviour
MSE	\(r^2\)	Penalises large residuals heavily
Huber	\(r^2\) if \(|r| < \delta\), else \(\delta|r|\)	Robust to outliers
Log-cosh	\(\ln(\cosh(r))\)	Smooth approximation to Huber

Recommendation: start with MSE; switch to Huber or log-cosh if training is unstable.

5. Impulse response functions

Once the model is solved, we can simulate it forward and compute IRFs — how the economy responds to a one-time productivity shock.

---

Exercises

Exercise 1: Train the stochastic DEQN for 10,000 epochs. Plot the policy function \(c(k, z)\) as a surface and compare slices at \(z = 0.95, 1.0, 1.05\).

Exercise 2: Simulate the solved model for 1,000 periods. Plot the time series of \(k_t\), \(c_t\), and \(z_t\). Compute the standard deviation of output.

Exercise 3: Compute impulse response functions: starting from steady state, apply a 1% productivity shock and trace the response of capital and consumption over 40 periods.

---

Next week: Constraints and Real-World Models — borrowing limits and Fischer–Burmeister complementarity.