HANK on Speed:
Robust Nonlinear Solutions using
Automatic Differentiation
Gregor Boehl
Uni Bonn
SED, June 2024
Motivation
Q: How to solve heterogeneous-agent GE models with many aggregate equations?
- Reiter: linearize aggregates, solve linear state space system
- SSJ: solve equilibrium as linear system in sequence space
Problem:
What about strong nonlinearities (OBCs, assymetries, ...)?
This paper
Solve for
nonlinear
sequence space equilibrium:
- Fast: nonlinear transition dynamics within seconds
- General: applies to broad class of HA models
- Accessible: high level open-source implementation (like dynare): the Econpizza package
HOW:
new way to solve linear system of each Newton step
new way to solve linear system of each Newton step
HANK
Detour: RANK
\begin{align}
(a_t, w_t) &= E_t \: W(w_{t+1}, d_{t+1}, x_{t-1}, x_t, x_{t+1}) \\
d_t &= D(a_t, d_{t-1}) \\
\mathbf{0} & = E_t \: f(x_{t-1}, x_t, x_{t+1}, d_t, a_t)
\end{align}
\begin{align}
(a_t, w_t) &= E_t \: W(w_{t+1}, \cancel{d_{t+1}}, x_{t-1}, x_t, x_{t+1}) \\
d_t &= D(a_t, d_{t-1}) \\
\mathbf{0} & = \cancel{E_t} f(x_{t-1}, x_t, x_{t+1}, d_t, a_t)
\end{align}
$$
\mathbf{0} = f(x_{t-1}, x_t, x_{t+1}, \cancel{d_t}, \cancel{a_t})
$$
- $x_t$: aggregate variables
- $w_t$: (marginal) values
- $a_t$: agent's actions
- $d_t$: distribution of agents
Detour: RANK
General representation of RANK models: $ \mathbf{0} = f(x_{t-1}, x_t, x_{t+1}) $Solve by finding $\mathbf{x}=\{x_t\}_0^T = \{x_0, x_1, \dots, x_T\} \in \mathbb{R}^{n \times T}$ s.t.
$$
\small
\begin{align}
\mathbf{0} & = f(x_{0}, x_1, x_{2}) \\
\mathbf{0} & = f(x_{1}, x_2, x_{3}) \\
&\vdots \\
\mathbf{0} & = f(x_{T-1}, x_{T}, \bar{x})
\end{align}
$$
$$
\small
\begin{rcases}
\mathbf{0} & = f(x_{0}, x_1, x_{2}) \\
\mathbf{0} & = f(x_{1}, x_2, x_{3}) \\
&\vdots \\
\mathbf{0} & = f(x_{T-1}, x_{T}, \bar{x})
\end{rcases} =F(\mathbf{x})
$$
Detour: RANK
Find $\mathbf{x}: F(\mathbf{x}) = \mathbf{0}$.
Newton's method:
$ \mathbf{x}_{i+1} = \mathbf{x}_i - J(\mathbf{x}_i)^{-1}F(\mathbf{x}_i) $
Newton's method:
$ \mathbf{x}_{i+1} = \mathbf{x}_i - J(\mathbf{x}_i)^{-1}F(\mathbf{x}_i) $
|
|
$$
\tiny
J(\mathbf{x}) =
\begin{bmatrix}
\partial f_1/\partial x_1 & \partial f_1/\partial x_2 & & & & \\
\partial f_2/\partial x_1 & \partial f_2/\partial x_2 & \partial f_2/\partial x_3 & & & \\
& \partial f_3/\partial x_2 & \partial f_3/\partial x_3 & \partial f_3/\partial x_4 & & & \\
& & \ddots & \ddots & \ddots & & \\
& & & \partial f_{T-2}/\partial x_{T-2} & \partial f_{T-2}/\partial x_{T-1} & \partial f_{T-2}/\partial x_T \\
& & & & \partial f_{T}/\partial x_{T-1} & \partial f_{T}/\partial x_{T} \\
\end{bmatrix}
$$
Back to HANK
\begin{align}
(a_t, w_t) &= W(w_{t+1}, x_{t-1}, x_t, x_{t+1}) \\
d_t &= D(a_t, d_{t-1}) \\
\mathbf{0} & = f(x_{t-1}, x_t, x_{t+1}, d_t, a_t)
\end{align}
|
|
$F(\mathbf{x})$ still exists (Auclert et al., 2021):
- given $w_T$: $\quad \mathbf{x} \underset{W}{\to} \mathbf{a}$
- given $d_0$: $\quad \mathbf{a} \underset{D}{\to} \mathbf{d}$
- $(\mathbf{x}, \mathbf{d}, \mathbf{a}) \underset{f}{\to} \mathbf{z} \overset{!}{=} 0$
Newton's method:
$
\mathbf{x}_{i+1} =
\mathbf{x}_i -
J(\mathbf{x}_i)^{-1}F(\mathbf{x}_i)
$
$$
\tiny
J(\mathbf{x}) =
\begin{bmatrix}
{\partial f_1/\partial x_1}^{\color{red}{*}} & {\partial f_1/\partial x_2}^{\color{red}{*}}
& \color{red}\text{stuff} & \color{red}\text{spam} & \cdots & \color{red}\text{loot} \\
{\partial f_2/\partial x_1}^{\color{red}{*}} & {\partial f_2/\partial x_2}^{\color{red}{*}} & {\partial f_2/\partial x_3}^{\color{red}{*}}
& \color{red}\text{junk} & & \color{red}\text{flotsam} \\
\color{red}\text{things} & {\partial f_3/\partial x_2}^{\color{red}{*}} & {\partial f_3/\partial x_3}^{\color{red}{*}} & {\partial f_3/\partial x_4}^{\color{red}{*}}
& & \color{red}\text{trinket} \\
\vdots & & \ddots & \ddots & \ddots & \\
\color{red}\text{jigger} & \color{red}\text{debris} & \color{red}\text{trash} &
{\partial f_{T-2}/\partial x_{T-2}}^{\color{red}{*}} & {\partial f_{T-2}/\partial x_{T-1}}^{\color{red}{*}} & {\partial f_{T-2}/\partial x_T}^{\color{red}{*}} \\
\color{red}\text{beer} & \color{red}\text{much} & \color{red}\text{more} & \color{red}\text{stuff} &
{\partial f_{T}/\partial x_{T-1}}^{\color{red}{*}} & {\partial f_{T}/\partial x_{T}}^{\color{red}{*}} \\
\end{bmatrix}
$$
Newton step for HANK
$$
\begin{align}
\mathbf{x}_{i+1} &= \mathbf{x}_i - J(\mathbf{x}_i)^{-1}F(\mathbf{x}_i) \\
0 &= J(\mathbf{x}_i) \mathbf{y} - F(\mathbf{x}_i) \qquad \qquad \text{with } \mathbf{y} = \mathbf{x}_{i} - \mathbf{x}_{i+1}
\\
0 &= \alpha \bar{J}^{-1}[F(\mathbf{x}_i) - J(\mathbf{x}_i) \mathbf{y}]\\
\mathbf{y} &= \mathbf{y} + \alpha \bar{J}^{-1}[F(\mathbf{x}_i) - J(\mathbf{x}_i) \mathbf{y}]
\end{align}
$$
- but convergence?
- requires all eigenvalues $\left|\lambda_k\left(I - \alpha\bar{J}^{-1} J(\mathbf{x}_i)\right)\right| < 1$
- only works close to steady state with $\alpha <2$
As an iterative scheme:
$$
\mathbf{y}_{j+1} = \mathbf{y}_j + \alpha \bar{J}^{-1}\left[F(\mathbf{x}_i) - J(\mathbf{x}_i) \mathbf{y}_j\right]
$$
As an iterative scheme:
$$
\mathbf{y}_{j+1} = \mathbf{y}_j + \alpha \bar{J}^{-1}\left[F(\mathbf{x}_i) - J(\mathbf{x}_i) \mathbf{y}_j\right]
$$
but for (with $R(M, \mathbf{z}) = \frac{\mathbf{z}^\intercal M \mathbf{z}}{\mathbf{z}^\intercal \mathbf{z}}$):
\begin{align}
\mathbf{y}_{j+1} &= \mathbf{y}_j + \alpha_j \bar{J}^{-1}\left[F(\mathbf{x}_i) - J(\mathbf{x}_i)\mathbf{y}_j\right] \\
\alpha_{j} &= \min\left\{\alpha_{j-1}, \gamma\left|R\left(\bar{J}^{-1}J(\mathbf{x}_i), \mathbf{y}_j\right)\right|^{-1}\right\}
\end{align}
Automatic Differentiation
- AD: computational technique to efficiently compute derivatives
- AD does not provide $J(\cdot)$ cheaply!
- But: AD provides $J(\mathbf{x}_i) \mathbf{y} = \Lambda(\mathbf{x}_{i}, \mathbf{y})$ gratis
|
|
Iterative scheme:
\begin{align}
\mathbf{y}_{j+1} &= \mathbf{y}_j + \alpha_j \bar{J}^{-1}\left[F(\mathbf{x}_i) - J(\mathbf{x}_i)\mathbf{y}_j\right] \\
\alpha_{j} &= \min\left\{\alpha_{j-1}, \gamma\left|R\left(\bar{J}^{-1}J(\mathbf{x}_i), \mathbf{y}_j\right)\right|^{-1}\right\}
\end{align}
Final iterative scheme:
\begin{align}
\mathbf{y}_{j+1} &= \mathbf{y}_j + \alpha_j \bar{J}^{-1}(F(\mathbf{x}_i) - \Lambda(\mathbf{x}_i, \mathbf{y}_j)) \\
\alpha_{j} &= \min\left\{\alpha_{j-1}, \gamma\left|R\left(\bar{J}^{-1}\Lambda(\mathbf{x}_i, \mathbf{y}_j), \mathbf{y}_j\right)\right|^{-1}\right\}
\end{align}
The Proposition
Given an initial guess $\mathbf{y}_0$, a scaling parameter $\gamma \in (1,2)$ and initializing $\alpha_0=1$, the iterative scheme \begin{align} \mathbf{y}_{j+1} &= \mathbf{y}_j + \alpha_j \bar{J}^{-1}(F(\mathbf{x}_i) - \Lambda(\mathbf{x}_i, \mathbf{y}_j)),\\ \alpha_{j} &= \min\left\{\alpha_{j-1}, \gamma/\left|R\left(\bar{J}^{-1}J(\mathbf{x}_i), \mathbf{y}_j\right)\right|\right\} \end{align} converges to \begin{equation} \lim_{j \to \infty} \mathbf{y}_{j} = J(\mathbf{x}_i)^{-1} F(\mathbf{x}_i) \end{equation} if all generalized eigenvalues of $\bar{J}$ and $J(\mathbf{x}_i)$ are real, positive and finite.
Speed!
| T=200 | T=300 | |||
| first run | again | first run | again | |
| RANK | 2.10s | 0.09s | 2.19s | 0.18s |
| HANK1 | 10.97s | 1.29s | 16.55s | 3.72s |
| HANK2 (small) | 26.43s | 2.13s | 57.39s | 10.70s |
| HANK2 | 28.38s | 1.53s | 80.56s | 8.38s |
Example: Model
2-Asset HANK + medium-scale DSGE
- HANK:
- cross sectional distribution of households
- precautionary savings
- portfolio choice
-
DSGE Features:
capital, capital adjustment, capital utilization, sticky prices á la Rotemberg, etc...
Example: permanent increase in transfers
Example: permanent increase in transfers
Implementation
- The Econpizza package
- Easy-to-use Python implementation relying on Google's JAX
- Exhaustive documentation
- High level interaction with dynare-like syntax:
- Plus: black box steady state cracking using pseudoinverse
Thanks!
- Solve nonlinear HANK models fast-as-light ✓
- Econpizza: easy to use software package ✓