Lecture 10: Sequence-space DEQNs

Course: Deep Learning for Solving and Estimating Dynamic Models in Economics and Finance
Script reference: §6.7 (Sequence-space DEQNs) -- Brock-Mirman warm-up
Notebook role: core
Author: Simon Scheidegger

Sequence-Space DEQN: Brock–Mirman Warm-Up¶

Status: in-class walkthrough — Day 4, sequence-space slot. This notebook is the simplest setting in which the sequence-space idea (network input = recent history of exogenous shocks instead of current endogenous aggregate state) can be made visible. The companion notebook 06_SequenceSpace_KrusellSmith.ipynb (TF2 teaching) and the JAX reference KrusellSmith_Tutorial_CPU.ipynb apply the same template to the heterogeneous-agent benchmark.

Learning goals¶

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

1. State the sequence-space formulation of a recursive equilibrium and explain the ergodicity-and-truncation argument that makes the truncated history a valid input for the policy network.

2. Recognise the single conceptual change versus a state-space DEQN: the network’s input changes (a shock-history window replaces the current endogenous aggregate state); the loss, the equilibrium conditions, and the gradient step are unchanged.

3. Read the four ingredients of the sequence-space cloud method as implemented here: a savings-rate network with sigmoid output, a Gauss–Hermite quadrature for the conditional expectation, a forward simulator that rolls the shock history one step at a time, and a per-episode SGD update on the squared Euler residual.

4. Validate the trained policy against the Brock–Mirman closed form $K_{t+1} = \alpha\beta\, A_t K_t^{\alpha}$ (available under $\delta = 1$, log utility), and read the diagnostics: loss curve, ergodic state cloud, and policy / Euler-error histograms.

What this notebook is not¶

• It is not a heterogeneous-agent solver. There is one representative household and one aggregate state. Sequence space here is a teaching device, not a dimensionality reduction. The payoff appears in the HA notebooks.

• It does not use shape-preserving I-spline outputs. The sigmoid-savings-rate parameterisation is enough for the BM warm-up; I-splines enter in 06_…KrusellSmith.

• The loss is the direct residual minimisation algorithm of Azinovic-Yang & Žemlička (2025). The alternative time iteration with EGM algorithm is not used here.

References¶

• Azinovic-Yang, M., & Žemlička, J. (2025). Deep learning in the sequence space. arXiv:2509.13623. Companion JAX repository: azinoma/DeepLearningInTheSequenceSpace.

• Brock, W. A., & Mirman, L. J. (1972). Optimal economic growth and uncertainty: the discounted case. Journal of Economic Theory 4(3), 479–513.

1. The Brock–Mirman model¶

A representative household maximises

subject to the resource constraint

with full depreciation $\delta = 1$ and an exogenous AR(1) for log productivity

Key analytic fact (used as a benchmark)¶

With $\gamma = 1$ and $\delta = 1$ the planner’s policy is closed-form:

We use this as a ground truth in the diagnostic plots (§14). The notebook also runs in regimes where the closed form does not apply — those skip the policy-comparison plot.

2. State-space vs sequence-space — the same equilibrium, two domains¶

The classical state-space representation of the recursive equilibrium reads:

The sequence-space representation expresses the same equilibrium policy as a function of the history of exogenous shocks instead of the current endogenous state:

Same equilibrium, different domain of approximation. The two formulations describe the same set of equilibrium policies. They differ in what the network takes as input. State-space inputs may be high-dimensional once a distribution $\mu_t$ enters; sequence-space inputs live on a (potentially long but) purely exogenous domain.

3. Ergodicity and truncation — why a finite history is enough¶

For a large class of economies the equilibrium policy satisfies the ergodicity property

i.e. the influence of long-past shocks decays. Truncating the history after $T$ lags gives a truncated sequence-space solution

Brock–Mirman makes this concrete¶

Iterating the BM transition (with $\delta = 1$, log utility, $s = \alpha\beta$) gives

so the lag-$j$ shock $A_{t-j}$ enters with weight $\alpha^{j-1}$. With $\alpha = 1/3$, the lag-25 weight is $\alpha^{24} \approx 3.5 \times 10^{-12}$ — truncating at $T = 25$ leaves negligible error. This is the entire justification for using a length-25 shock-history window as the network input.

4. Algorithm sketch — direct residual minimisation with a cloud method¶

The training loop has three pieces. We will see each one in code below.

1. Initialisation. Sample $N_{\text{cloud}}$ initial state-history pairs $(x_j^{(0)}, \mathcal{E}_j^{(0)})$. We initialise all of them at the deterministic steady state with empty histories.

2. Per-episode training (mini-batch SGD).

   • Draw a mini-batch of cloud entries.

   • Forward-pass the network $\mathcal{N}_\theta(\mathcal{E})$ to get a savings rate $s$, then $c$, $\mu = u'(c)$, and the implied $k_{t+1}$.

   • For each Gauss–Hermite quadrature node, build the next-period shock history and run the network again to get $\mathbb{E}_t[u'(c_{t+1})\,\mathrm{MPK}_{t+1}]$.

   • Compute the squared Euler residual; backprop and Adam step.

3. Cloud refresh. After every episode, draw fresh shocks for the entire cloud and roll one period forward (forward-simulator step, no gradients). The cloud now lives one period later — the ergodic distribution under the current policy emerges over many episodes.

The cloud-method sampling concentrates training data on the ergodic set: states the model actually visits under the policy being trained, rather than uniformly drawn states which may include economically irrelevant regions.

5. Imports and TF/GPU setup¶

We use TensorFlow 2 / Keras. All later cells assume the imports and the symbol USE_GPU are in scope.

6. Calibration¶

The choice $\gamma = \delta = 1$ is what unlocks the closed-form benchmark $s = \alpha\beta$. Change either to a non-unit value and the analytical comparison plots in §14 are simply skipped.

State and policy index conventions¶

The state vector is 5-wide here, not 2-wide: in addition to $(k_t, z_t)$ we pre-compute $(Z_t = e^{z_t},\, y_t,\, \mathrm{MPK}_t)$ in ExpandState so the loss function can read them off directly without recomputing each time. The policy vector is 3-wide: savings $s$, consumption $c$, marginal utility $\mu = u'(c)$.

7. Algorithm hyperparameters and Gauss–Hermite quadrature¶

The Gauss–Hermite nodes and weights are rescaled so that they integrate against the standard-normal density: nodes are multiplied by $\sqrt{2}$, weights divided by $\sqrt{\pi}$, so that

8. Deterministic steady state and initial state¶

At $z = 0$ (so $A = 1$), the deterministic steady state $k_{\text{ss}}$ solves

With $\delta = 1$ this gives $k_{\text{ss}} = (\alpha\beta)^{1/(1-\alpha)}$. We use this point as the initial condition for every track of the cloud at episode 0; the policy-driven simulator then disperses the cloud to its ergodic distribution.

9. Economic primitives — utility and the BM analytic policy¶

The helper mksp (“make positive”) clips its argument away from zero by adding a tiny eps_safe. This avoids c < 0 blowing up u'(c) = c^{-\gamma} during the early, untrained iterations.

Marginal utility and its inverse¶

• u_c(c, params) returns $u'(c) = (1-\beta)\, c^{-\gamma}$. The factor $(1-\beta)$ scales the optimisation problem so that the value function is on a similar scale to per-period utility (bounded, normalised problem). It does not affect the policy.

• u_c_i(u, params) is the inverse: given a marginal utility level it returns the consumption that produces it. Used in the Euler residual to convert $\beta\,\mathbb{E}[u'(c_{t+1})\,\mathrm{MPK}_{t+1}]$ into “implied consumption today”.

Brock–Mirman closed form¶

BrockMirman(x, params) returns the analytic savings $s_{\text{BM}} = \alpha\beta\,y$ for the case $\gamma = \delta = 1$. We evaluate this on the trained cloud at the end to verify the network learned the right policy.

10. ExpandState: precompute derived state quantities¶

Given the core state $(k, z)$, this returns the expanded 5-wide state $(k, z, Z = e^z, y = Z\, k^{\alpha} + (1-\delta)\,k, \mathrm{MPK} = \alpha Z k^{\alpha-1} + (1-\delta))$.

The reason for storing these explicitly: the Euler residual computation later needs $\mathrm{MPK}_{t+1}$ at each quadrature node. Building it inside the loss and watching TensorFlow trace the same algebra repeatedly would be wasteful; precomputing keeps the loss function clean.

11. The sequence-space network¶

The network architecture is intentionally minimal:

• Input. A length-25 vector of past innovations $\mathcal{E}_t^{(25)} = (\varepsilon_t, \varepsilon_{t-1}, \ldots, \varepsilon_{t-24})$. At episode 0 the cloud has all-zero histories; histories grow as the simulator runs.

• Hidden layers. Three dense layers of width 64 with gelu activation.

• Output. A single scalar passed through sigmoid — interpreted as the savings rate $s \in (0, 1)$. This is a hard constraint in the spirit of Day 4 nb 02: feasibility ($0 < s < 1$, hence $c > 0$ and $k_{t+1} > 0$) is built into the architecture, never a loss term.

Why predict the savings rate instead of $k_{t+1}$ directly? Two reasons:

1. the savings rate is bounded and a smoother function of the state (much easier for the network to learn than the raw level $k_{t+1}$, which has a wide range);

2. cumulating $s\, y$ recovers $k_{t+1}$ exactly without an extra layer.

12. From network output to a full policy¶

Policy_sequence(x, history, model, params) reads the precomputed $y$ from the expanded state, runs the network on the shock history to get the savings rate $s_{\text{rate}} \in (0,1)$, and returns the 3-wide policy vector

The 3-D-input branch in the function (if len(history.shape) == 3) handles the batched evaluation across Gauss–Hermite quadrature nodes inside the Euler-residual computation: there, each batch entry has 8 next-period histories (one per quadrature node), so the input tensor has shape [batch, 8, 25] and we reshape to [batch*8, 25] to feed Keras, then reshape the output back.

13. The Euler residual in sequence space — the loss core¶

The household’s stochastic Euler equation (with full depreciation $\delta = 1$) is

Writing $\mu_{t+1} = u'(c_{t+1})$ and integrating the conditional expectation by Gauss–Hermite,

The residual we minimise is the relative-error form

i.e. implied consumption today / actual consumption today − 1. Squaring and averaging over the cloud gives the loss.

What the function E_sequence does, step by step¶

1. Read $z_t$, $c_t$, $s_t$ from the current state and policy (Policy_sequence).

2. Build the next-period shock $z_{t+1} = \rho_z z_t + \sigma_z\, \xi_q$ for each of 8 quadrature nodes — broadcasted over the batch (shape becomes [batch, 8, 1]).

3. Build the next-period expanded state at every quadrature node by feeding $(k_{t+1} = s_t,\, z_{t+1})$ through ExpandState.

4. Build the next-period shock history by prepending each quadrature node and dropping the oldest entry of the current history (shape: [batch, 8, 25]).

5. Re-run the network at every quadrature node to get $\mu_{t+1}$ (this is the “second forward pass” — analogous to the second tape.gradient call in nb 02–04).

6. Form the GH-weighted expectation, invert via u_c_i, divide by $c_t$, subtract 1.

Note: the network is called twice in this function — once at $(\mathcal{E}_t)$ and once at every $(\mathcal{E}_{t+1}, \xi_q)$. Both calls use the same network weights $\theta$, and TensorFlow correctly backpropagates through both.

14. Forward simulator: rolling the shock history one period¶

ForwardStep_sequence is not used inside the loss — it is the cloud-refresh step. For each cloud entry it draws a single innovation $\varepsilon_{t+1}$ (the epsilon argument), computes:

• $z_{t+1} = \rho_z z_t + \sigma_z \varepsilon_{t+1}$,

• $k_{t+1} = s_t$ (the savings rate already returned $s = s_{\text{rate}}\, y$),

• new expanded state via ExpandState,

• new history $(\varepsilon_{t+1}, \varepsilon_t, \ldots, \varepsilon_{t-23})$ — prepend the fresh innovation, drop the oldest.

Because the function is decorated with @tf.function and receives only constants from outside the graph, TensorFlow compiles it once and reuses the trace each episode — fast.

15. The single SGD step¶

train_step wraps the loss in a tf.GradientTape, computes the gradient of the mean squared Euler residual w.r.t.\ the network’s trainable variables, and applies one Adam step. Standard pattern; nothing model-specific here.

16. The cloud-method training loop¶

CloudMethod_sequence is the outer loop. Each episode:

1. Build a tf.data.Dataset from the current cloud (x_cloud, history_cloud), shuffle, batch into mini-batches of size n_minib = 256.

2. Iterate over mini-batches, accumulate per-episode loss.

3. Print progress every 20 episodes.

4. Draw fresh shocks for the entire cloud (one shock per cloud entry, via tf.random.normal), and roll forward one period with ForwardStep_sequence.

The cloud entries’ histories grow by one entry each episode, until they fill the length-25 window after 25 episodes. From episode 25 onwards every entry has a “warm” history.

Runtime note. With the defaults (cloud_steps = 2048, n_cloud = 4096, n_minib = 256, GH order 8) the loop does $2048 \times 16 = 32{,}768$ SGD steps, each evaluating the network on 256 + $256 \times 8 = 2304$ inputs. CPU runtime is in the 20-40 minute range; GPU compresses it to a few minutes.

17. Diagnostics¶

Three diagnostic plots:

1. Loss curve ($\log_{10}$ of the per-episode mean squared Euler residual) — should decrease monotonically (with noise) and reach $\sim 10^{-6}$ on the BM problem.

2. Final state cloud $(k, z)$ — the cloud should concentrate around the stochastic steady state, with spread set by the AR(1) dynamics.

3. Policy comparison and error histograms — only shown when $\gamma = \delta = 1$ (closed-form available). The network’s savings $s_{\text{NN}}$ should overlay the analytic $\alpha\beta\, y$; the Euler-error histogram should sit at small relative error.

17.2 Ergodic state cloud¶

The scatter shows the final positions of all $n_{\text{cloud}} = 4096$ tracks at the end of training. We expect a tight blob around $(k_{\text{ss}}, 0)$ with spread along both axes set by the unconditional variance of the AR(1) productivity process.

17.3 Policy and Euler errors (only when $\gamma = \delta = 1$)¶

If we are in the analytic regime, we overlay the network’s savings policy against $\alpha\beta\, y$ on a few cloud samples, and plot the histogram of:

• Policy error $|s_{\text{NN}} - s_{\text{analytic}}| / |s_{\text{analytic}}|$.

• Euler error $|\mathrm{err\_eul}|$ (relative form).

Both should be tightly concentrated at small values.

18. Take-away¶

The single conceptual change versus the state-space DEQNs of Day 4 nb 02–03:

In the BM model the input swap is larger, not smaller — sequence space is a teaching device here, not a dimensionality reduction. The motivation only pays off in the heterogeneous-agent setting (06_SequenceSpace_KrusellSmith.ipynb and KrusellSmith_Tutorial_CPU.ipynb), where the alternative input would be a 100-bin histogram of agents’ wealth. There the shock history is genuinely smaller and exogenous (does not move with the policy update) — which is the source of the stability gains documented in Azinovic-Yang & Žemlička (2025).

Suggested follow-on¶

1. Change n_history and re-run. The error should bottom out as soon as $\alpha^{T}$ falls below the loss floor.

2. Set $\delta < 1$ and observe that the policy is no longer $\alpha\beta\, y$, but the algorithm still works — only the analytic comparison plot is skipped.

3. Open 06_SequenceSpace_KrusellSmith.ipynb for the heterogeneous-agent extension.