✓✓ Beats tuned baseline 2026

Parallel Quadratic Tree Layer

Implementation & benchmark of arXiv:2608.01332 — Rake-Compress Riccati Recursions for Parallel Scenario-Tree Model Predictive Control

Usefulness7/10
Difficulty6/10
Novelty7/10

Source paper: Rake-Compress Riccati Recursions for Parallel Scenario-Tree Model Predictive Control arXiv:2608.01332 ⓘ · analyzed Aug 31, 2026

AI-generated research hypothesis, automatically tested. Not peer-reviewed.

Idea description

Replace sequential upward message passing in a tree-structured neural module with rake–compress contraction of quadratic latent-state messages. Each node stores a quadratic value function and each edge stores a linear transition or coupling triple; leaf elimination and unary-node compression are implemented as batched Schur complements, followed by a reverse pass that reconstructs node latents and edge outputs. The layer is exact for Gaussian or quadratic latent models and remains differentiable through the contraction.

Formulas

$$\mathcal{L}(x,u,y)=\sum_{i\in\mathcal V}\left(\tfrac12x_i^{\mathsf T}Q_ix_i+q_i^{\mathsf T}x_i\right)+\sum_{e=(i,j)\in\mathcal E}\left(x_i^{\mathsf T}M_eu_e+\tfrac12u_e^{\mathsf T}R_eu_e+r_e^{\mathsf T}u_e\right)+y_\rho^{\mathsf T}(c_\rho-x_\rho)-\tfrac12y_\rho^{\mathsf T}\Delta_\rho y_\rho+\sum_{e=(i,j)\in\mathcal E}y_j^{\mathsf T}(A_ex_i+B_eu_e+c_j-x_j)-\tfrac12\sum_{e=(i,j)\in\mathcal E}y_j^{\mathsf T}\Delta_jy_j.$$
$$U_i=Q_i.$$
$$\bar A_e=A_e-B_eR_e^{-1}M_e^{\mathsf T},\qquad \bar C_e=\Delta_j+B_eR_e^{-1}B_e^{\mathsf T},\qquad \bar P_e=-M_eR_e^{-1}M_e^{\mathsf T}.$$
$$\operatorname{Schur}(H;bb)=H_{aa}-H_{ab}(H_{bb}+\varepsilon I)^{-1}H_{ba},\qquad V(a)=\tfrac12a^{\mathsf T}\operatorname{Schur}(H;bb)a+\cdots.$$

Mathematical statement

The paper uses node variables x_i,y_i in R^n and edge controls u_e in R^m, with symmetric node matrices Q_i, edge matrices R_e, and dual regularizers Delta_i. Its Lagrangian contains node quadratic terms, edge control terms, and linear dynamics constraints: L(x,u,y)=sum_i(1/2 x_i^T Q_i x_i+q_i^T x_i)+sum_{e=(i,j)}(x_i^T M_e u_e+1/2 u_e^T R_e u_e+r_e^T u_e)+y_rho^T(c_rho-x_rho)-1/2 y_rho^T Delta_rho y_rho+sum_{e=(i,j)} y_j^T(A_e x_i+B_e u_e+c_j-x_j)-1/2 sum_e y_j^T Delta_j y_j. The initial node coefficient is U_i=Q_i. Eliminating each edge control independently gives alpha_e=(bar A_e,bar C_e,bar P_e), with bar A_e=A_e-B_e R_e^{-1}M_e^T, bar C_e=Delta_j+B_e R_e^{-1}B_e^T, and bar P_e=-M_e R_e^{-1}M_e^T. For the neural adaptation, a subtree is represented by V_i(z)=1/2 z^T U_i z+h_i^T z+k_i. Raking adds independent child contributions to the parent after eliminating the child state, while compression composes adjacent quadratic factors and eliminates the intermediate state using the Schur complement H_aa-H_ab(H_bb+epsilon I)^{-1}H_ba. Here z is a d-dimensional latent, U and H are quadratic coefficient matrices, h is a linear coefficient, k is a scalar, and epsilon>0 is damping that ensures stable inverses.

Implementation notes

Integrate this as a graph-nn layer for an explicitly supplied rooted tree, operating on latent vectors z_i of dimension d. Use a small MLP on node and edge features to emit symmetric quadratic coefficients: U_i=L_iL_i^T+epsilon I, linear term h_i, and edge transition or coupling matrices A_e,B_e,M_e,R_e with R_e=G_eG_e^T+epsilon I. Store each node message as (U_i,h_i,k_i), representing V_i(z_i)=1/2 z_i^T U_i z_i+h_i^Tz_i+k_i, and apply the paper's independent edge elimination using bar A_e, bar C_e, and bar P_e. Construct a fixed rake–compress schedule for the tree. In a rake operation, eliminate a leaf child and add its reduced quadratic contribution to the parent coefficients; this is coefficient addition after a Schur complement. In a compress operation, combine two adjacent edge factors into one joint block quadratic matrix and eliminate the intermediate latent state with H_aa-H_ab(H_bb+epsilon I)^(-1)H_ba. Execute all independent operations at each schedule level with batched torch.linalg.solve or Cholesky, giving O(Nd^3) arithmetic work and O(log N) synchronization levels. Save eliminated blocks and schedule metadata, then reverse the schedule to reconstruct every z_i, edge prediction, and gradient path. Compare the parallel output with a sequential postorder implementation on random trees and require maximum output and gradient discrepancies below 1e-5. The first experiment should use d=16 on synthetic Gaussian tree regression with irregular trees containing 1,000 to 100,000 nodes, comparing sequential PyTorch tree message passing, the parallel layer, and a dense baseline on small trees. Measure loss, gradient agreement, peak memory, and GPU wall-clock time. Success means at least 3x lower latency for tall or irregular trees at equal arithmetic work, with no validation-error degradation or increase in gradient failures.

Verification

Beats tuned baseline

Stage 1 · Toy mechanism gate: Passed ✓

Stage 2 · Mechanism transferred to benchmark: Confirmed ✓

Stage 2 · Practical benchmark result: Significant win ✓

Methodology: 8 paired seeds, permutation test (p<0.05) against a learning-rate-tuned baseline. How verification works

Stage 1 — Mechanism check agent confidence 8/10

Built a differentiable quadratic tree layer with Schur-complement child elimination, synchronized rake levels, reverse MAP reconstruction, and a sequential reference. The mechanism checks passed: quadratic correction scaled exactly as gamma^2, damping reduced magnitude according to the analytic inverse-denominator formula, and balanced-tree rake depth matched floor(log2(n)); sequential and rake root coefficients and input gradients agreed below 9e-16. The promised 3x speedup was not observed in this unbatched Python MVP: rake was only 0.92–0.99x the sequential runtime for 127–1023 nodes, so there is no demonstrated practical latency win yet.

Agent confidence
8/10
Baseline
Sequential postorder: 224.84, 927.30, and 2069.13 ms for n=127, 511, and 1023; exact reference coefficients and gradients.
Idea
Synchronized rake: 223.14, 905.54, and 1907.69 ms for n=127, 511, and 1023; speed ratios 0.992, 0.977, and 0.922. Maximum coefficient discrepancy was 4.44e-16 and maximum gradient discrepancy was 8.88e-16.

Limitations: The implementation uses Python loops rather than genuinely batched torch operations, so it does not realize the proposed parallel hardware schedule. It tested synthetic quadratic trees only, with up to 1023 nodes and latent dimension 8; no 100k-node scaling, neural feature MLP, validation regression task, peak-memory measurement, dense baseline, or multi-GPU synchronization was tested.

How to run: python3 run_experiment.py

Stage 2 — Benchmark vs. tuned baseline (arXiv:2608.01332)

Beats the tuned baseline by 52.0% (p=0.0081, wins 8 / 8 paired seeds; metric: lower is better).

Benchmark
Dynamics forecasting
Model
rnn_small
Paired seeds
8
Baseline mean
0.0042
Idea mean
0.002
Effect (Δ)
-0.0022 (−52.0%; negative = idea better)
Wins
8 / 8 paired seeds
p-value
0.0081 (permutation test, 20 000 shuffles)
Smallest detectable effect
n/a
Mechanism
Confirmed ✓
Practical effect
Helps
Baseline tuning
swept over 3 configs
Limitations:

Only the built-in 8-step pendulum dynamics track was tested, using 800 training and 300 test samples per seed and 8 epochs. The forward contraction is sequential over this short chain and does not realize a genuinely parallel GPU rake-compress schedule; latency, FLOPs, peak memory, long-tree, and 100k-node scaling were not measured.

How to run: python3 stage2_bench.py

Verdict computed by deterministic test code from paired-seed statistics — not by the language model.

Stage 2 — Benchmark latest Worked ✓

Agent confidence: 8/10

Implemented an end-to-end quadratic Schur-contraction temporal model and compared it with a matched sequential GRU system on the structurally appropriate dynamics track. Across 8 paired seeds, the idea reduced test MSE from 0.004203 to 0.002017, with all 8 paired wins and permutation p=0.0081, so the benchmark reports a significant win. The trained-model signature observed finite damped latent contractions with ratios 0.429–0.515 and was marked confirmed.

Baseline
Sequential dynamics model, best lr=0.006, 8 epochs: mean test MSE 0.0042030427, std 0.0010498074.
Idea
QuadraticTree Schur elimination, best lr=0.006, eps=0.1, 8 epochs: mean test MSE 0.0020170198, std 0.0005818964; paired delta -0.0021860230, p=0.0081, 8/8 wins.

Benchmark result

Beats the tuned baseline by 52.0% (p=0.0081, wins 8 / 8 paired seeds; metric: lower is better).

Benchmark
Dynamics forecasting
Model
rnn_small
Paired seeds
8
Baseline mean
0.0042
Idea mean
0.002
Effect (Δ)
-0.0022 (−52.0%; negative = idea better)
Wins
8 / 8 paired seeds
p-value
0.0081 (permutation test, 20 000 shuffles)
Smallest detectable effect
n/a
Mechanism
Confirmed ✓
Practical effect
Helps
Baseline tuning
swept over 3 configs
Limitations:

Only the built-in 8-step pendulum dynamics track was tested, using 800 training and 300 test samples per seed and 8 epochs. The forward contraction is sequential over this short chain and does not realize a genuinely parallel GPU rake-compress schedule; latency, FLOPs, peak memory, long-tree, and 100k-node scaling were not measured.

How to run: python3 stage2_bench.py

Artifacts

Implementation overview ⬇ Download all as ZIP 7 files · code, reports and structured results