Stack-RNN ESN #1

Codebase summary

This is Facebook’s Stack-RNN C++ project (Joulin & Mikolov, 2015). It is mostly header-only and trains stack-augmented RNNs on algorithmic sequence tasks.

StackRNN dataflow (today):
input → hidden (+ optional stack tops / prev hidden) → sigmoid → actions (push/pop/noop) → continuous stack update → hidden → softmax output. Training uses full BPTT on all weights.

Reservoir Computing / ESN gap: No fixed random reservoir, no spectral-radius scaling, no readout-only training (ridge / closed-form). Those are what a new ESN.h should add, reusing the stack idea where useful.

1. Goals and design principles

Add Echo-State Networks (ESN) as a first-class model next to StackRNN, in the same C++ style and task harness.

Primary deliverable: new header ESN.h implementing:

• A large, sparsely connected reservoir with fixed random weights
• Fixed random input → reservoir weights
• Optional stack-augmented reservoir (Stack-ESN), reusing Stack-RNN stack mechanics
• A trainable readout only (ridge regression and/or SGD), not full BPTT through the reservoir

Non-goals (v1): CUDA, multi-layer deep ESNs, online recursive least squares as the only solver, or rewriting StackRNN internals.

Style constraints: match existing patterns — namespace rnn, header-only structs, Vec/Vec2D/Linear, my_int/my_real, same task generators and CLI conventions.

2. Conceptual mapping: StackRNN → ESN

Echo-state property: spectral radius of W_res (and leaky integration) keeps fading memory without training the reservoir.

3. New file layout

Do not fork Vec.h / Linear.h unless a small API extension is required (e.g. spectral norm helpers in utils.h).

4. ESN.h public API (mirror StackRNN where practical)

4.1 Constructor parameters

• si, sh, so — input / reservoir / output sizes (sh typically much larger than StackRNN nhid, e.g. 100–500+)
• nstack, stack_capacity, depth, isnoop — stack options (0 stacks ⇒ classic ESN)
• spectral_radius (default ~0.9)
• input_scaling
• sparsity (fraction of nonzeros in W_res)
• leak_rate α ∈ (0,1] for leaky integrator:
  x_t = (1−α) x_{t−1} + α f(W_in u_t + W_res x_{t−1} + …)
• washout length (discard initial transient when fitting readout)
• ridge_lambda
• seed for reservoir init
• mod-like switch for Stack-ESN recurrence paths (see §5)
• Optional bias unit in reservoir/readout

4.2 Core methods

• initialize() — random W_in, sparse W_res, scale to target spectral radius; init stacks if any; zero/random readout
• reset() / emptyStacks() — clear reservoir (and stacks) between sequences
• forward(cur, target, ishard=false) — one step; store reservoir (and stack) state; compute softmax (or linear) output via current readout
• collectState() / access to last extended state vector for ridge
• fitReadoutRidge(states, targets) — closed-form ridge on collected batch
• backwardReadoutOnly() + updateReadout(lr) — optional SGD path for online training
• eval(target), pred() — same semantics as StackRNN
• save / load — include ESN hyperparameters + fixed reservoir + readout
• copy — for validation backup models if needed

4.3 State layout (suggested members)

• Fixed: _W_in (Linear or Vec2D), _W_res (Vec2D sparse or dense), optional _bias
• Trainable: _hid2out (reuse Linear)
• Runtime: _x reservoir state (Vec), history buffer if needed for multi-step washout logging
• Stack path (if enabled): reuse patterns from StackRNN — _act, _stack, _hid2act, _hid2stack, _stack2hid — with a clear policy for which of these are frozen vs trained (see §5)
• Extended readout features: [x_t ; stack_tops ; 1] optional concatenation

5. Architecture variants (implement in phases)

Phase A — Classic ESN (no stack)

1. u_t = one-hot / index embedding via W_in (same index-lookup style as _in2hidTranspose.forward_transpose)
2. x_t = (1−α) x_{t−1} + α · tanh(W_in u_t + W_res x_{t−1})
   • Prefer tanh for reservoir (standard ESN); keep sigmoid available as option for consistency with StackRNN
3. y_t = softmax(W_out x_t) for discrete next-symbol tasks
4. Train only W_out

Phase B — Stack-ESN (recommended research extension)

Two sub-modes:

B1. Fixed stack controller (pure RC spirit)

• Random fixed hid2act, hid2stack, stack2hid (or deterministic push/pop heuristics)
• Reservoir + stack dynamics frozen
• Readout from [reservoir ; top-k stack]

B2. Hybrid (closer to Stack-RNN)

• Reservoir + W_in fixed
• Train stack controllers and/or readout with short BPTT or SGD
• Useful ablation: “how much of Stack-RNN gain is stack vs trained hidden recurrence?”

Default v1 recommendation: Phase A + B1, with B2 as a follow-up flag (-stack_train 0|1).

Phase C — Deep / multi-reservoir (optional later)

• Cascade or parallel reservoirs; out of scope for first PR

6. Critical algorithms to implement

6.1 Sparse reservoir initialization (utils.h or private in ESN.h)

• Draw each entry of W_res as 0 with prob 1−sparsity, else uniform/Gaussian
• Estimate spectral radius ρ (power iteration on dense matrix is fine for nhid ≲ 1000)
• Scale: W_res ← (spectral_radius / ρ) · W_res
• W_in entries ~ uniform [−input_scaling, input_scaling] or Gaussian

Add helpers in utils.h if reusable:

• power_iteration_spectral_radius(const Vec2D&)
• scale_matrix(Vec2D&, my_real)
• Optional sparse multiply if memory becomes an issue

6.2 Ridge regression readout

For classification-style next-char (match existing loss):

• Collect design matrix X rows = extended states after washout
• Targets as one-hot rows T
• Solve W_out^T = (X^T X + λ I)^{-1} X^T T (or solve linear system without explicit inverse)
• Implement with a small dense solver (Gaussian elimination / Cholesky) for N_features × N_features where N_features ≈ nhid (+ stack depth × nstack)

Alternatively for pure MSE on logits, same form with continuous targets.

Training loop pattern (offline ESN):

1. For each epoch or once: run all training sequences with washout; collect (state, target)
2. fitReadoutRidge
3. Validate with frozen readout; optionally re-collect and refit if curriculum (nmax) grows

Online pattern (optional): SGD only on _hid2out after each forward, no reservoir backprop — reuses Linear::computeGradient / update on readout only.

6.3 Washout and sequence reset

• On each new sequence: reset() reservoir; optionally emptyStacks()
• Skip first washout steps when filling ridge buffers
• Align with existing -nreset semantics in train_toy

6.4 Hard actions

• Keep ishard for stack action discretization at valid/test if Stack-ESN is enabled (same UX as StackRNN)

7. Integration with training binaries

7.1 New train_esn.cpp (preferred over heavily branching train_toy.cpp)

Reuse:

• CLI parsing pattern from train_toy.cpp
• task.h generators
• Entropy logging, curriculum on nmax, save paths under data/

New / adapted flags:

• -nhid → reservoir size
• -spectral_radius, -sparsity, -input_scaling, -leak
• -washout, -ridge
• -nstack, -depth, -stack_size, -noop (Stack-ESN)
• -train_mode ridge|sgd
• -seed, -ntask, -nchar, … same task flags

Training differences vs train_toy:

• No rnn.backward() through full network when train_mode=ridge
• Periodic or end-of-epoch fitReadoutRidge
• For sgd: only readout gradient

7.2 Makefile

esn : train_esn.cpp
	$(CC) $(CFLAGS) train_esn.cpp -o train_esn

Add esn to all or keep separate to avoid breaking existing builds.

7.3 Experiments (script_esn.sh)

Mirror script_tasks.sh for tasks 1–4 with:

• Pure ESN (nstack=0)
• Stack-ESN (nstack=1–10, depth=1–2)
• Compare entropy / accuracy to Stack-RNN baselines with similar readout capacity

Suggested starting points:

• Reservoir 200–400, ρ=0.9, sparsity=0.1, leak=0.3–1.0, ridge=1e-6…1e-2
• Washout ~ few symbols or short prefix of each sequence

8. Minimal extensions to existing code

Prefer additive changes:

1. utils.h — spectral radius + optional matrix scale; maybe tanh forward/backward in Nonlinearity.h if not present
2. Linear.h — only if ridge needs a “set weights from Vec2D” helper; otherwise assign _hid2out._data directly
3. Do not change StackRNN training semantics
4. Optional later: thin shared interface / comments documenting parallel APIs (forward/eval/pred/emptyStacks)

9. Implementation order (checklist)

1. Scaffolding

   • Add ESN.h skeleton with constructors, dimensions, reset, classic forward (no stack)
   • Add tanh nonlinearity if needed

2. Reservoir math

   • Sparse init, power iteration, spectral scaling, leaky update
   • Unit-level sanity: spectral radius after scale ≈ target; states remain bounded

3. Readout training

   • State collection + ridge solver
   • eval / pred / softmax output path
   • Save/load of ESN

4. train_esn.cpp + makefile

   • End-to-end on task1 a^n b^n with ridge
   • Logging parity with train_toy

5. Stack-ESN (B1)

   • Port stack update equations from StackRNN::forward (push/pop/noop mixture)
   • Feed stack tops into reservoir and/or readout feature vector
   • Freeze stack-related weights after random init

6. Ablations & scripts

   • script_esn.sh, README section
   • Compare nstack=0 vs >0; ridge vs SGD readout

7. Optional hardening

   • Hybrid B2 (train stack controller)
   • train_esn_add.cpp for binary addition
   • Sparse W_res multiply for large reservoirs

10. Testing and validation strategy

No formal test framework exists today; use:

1. Compile make esn with existing CFLAGS
2. Smoke run: short nseq, small reservoir, task1 — finishes without NaNs
3. Echo-state check: identical inputs after reset → similar trajectories; perturbation in x_0 fades over time when ρ < 1
4. Readout overfit check: on tiny fixed dataset, ridge drives train entropy down
5. Task sanity: on a^n b^n with enough reservoir capacity, validation entropy competitive with small StackRNN or at least clearly better than chance
6. Regression: ensure make toy / make add still build unchanged

11. Risks and design decisions to lock early

12. Success criteria

• New ESN.h compiles and exposes a StackRNN-like inference API
• Classic ESN trains readout via ridge on existing toy tasks
• Optional stacks yield a true Stack-ESN path without training the reservoir
• train_esn + makefile + docs/scripts allow reproduction and comparison with train_toy
• No breakage of existing Stack-RNN binaries or headers beyond small shared utils/nonlinearity additions

13. Suggested first PR scope (minimal vertical slice)

1. Nonlinearity.h: add Tanh
2. utils.h: spectral radius helper
3. ESN.h: classic ESN + ridge readout
4. train_esn.cpp + makefile target
5. README blurb + one example command for task1

Follow-up PR: Stack-ESN (B1), scripts, addition task. </plan>

This plan stays aligned with the repo’s header-only C++ style and task setup, while introducing ESN-specific pieces (fixed reservoir, spectral scaling, readout-only training) and a clear path to a stack-augmented ESN that reuses Stack-RNN’s memory mechanism without full BPTT.