✓✓ Beats tuned baseline 2026

Alternating ridge least-squares final layers

Implementation & benchmark of arXiv:2607.06976 — Hybrid Least Squares/Gradient Descent Methods for MIONets

Usefulness7/10
Difficulty5/10
Novelty6/10

Source paper: Hybrid Least Squares/Gradient Descent Methods for MIONets arXiv:2607.06976 ⓘ · analyzed Aug 30, 2026

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

Idea description

Replace gradient updates for one branch's final linear layer at a time with an exact ridge least-squares solve while holding the other branches, trunk, and hidden layers fixed. The method applies to any model whose output is a sum of products of branch factors and a trunk factor, including MIONets and tensorized neural networks.

Formulas

$$G(u_1,\ldots,u_N)(y)\approx\left\langle (\mathbf b_1\circ\phi^1_{M_1})(u_1)\odot\cdots\odot(\mathbf b_N\circ\phi^N_{M_N})(u_N),\mathbf t(y)\right\rangle.$$
$$\hat y_j=\sum_{i=1}^{I}t_i(y_j)\,b_{m,i}(u_{m,j})\prod_{r\neq m}b_{r,i}(u_{r,j})=\big(c_j^\top\otimes h_m(u_{m,j})^\top\big)\operatorname{vec}(W_m),\quad c_{j,i}=t_i(y_j)\prod_{r\neq m}b_{r,i}(u_{r,j}).$$
$$w_m^{+}=\arg\min_{w}\sum_{s=1}^{S}\left\|A_sD_mw-z_s\right\|_2^2+\lambda_m\|w\|_2^2,$$
$$\left(\sum_{s=1}^{S}D_m^\top A_s^\top A_sD_m+\lambda_m I\right)w_m^{+}=\sum_{s=1}^{S}D_m^\top A_s^\top z_s.$$

Mathematical statement

The paper represents a multiple-input operator as \(G(u_1,\ldots,u_N)(y)\approx\left\langle (\mathbf b_1\circ\phi^1_{M_1})(u_1)\odot\cdots\odot(\mathbf b_N\circ\phi^N_{M_N})(u_N),\mathbf t(y)\right\rangle\), where \(u_m\in X_m\) are inputs, \(\phi^m_{M_m}\) extracts the first \(M_m\) Schauder coefficients, \(\mathbf b_m\in\mathbb R^I\) is branch \(m\)'s rank-\(I\) feature vector, \(\mathbf t\in\mathbb R^I\) is the trunk feature vector, and \(\odot\) is the Hadamard product. Let the trainable last layer of branch \(m\) be \(W_m\in\mathbb R^{H_m\times I}\), with hidden feature \(h_m(u_m)\in\mathbb R^{H_m}\) and \(\mathbf b_m=h_m(u_m)^\top W_m\). For sample-location pair \(j\), define \(c_{j,i}=t_i(y_j)\prod_{r\neq m}b_{r,i}(u_{r,j})\). The prediction is linear in \(w_m=\operatorname{vec}(W_m)\): \(\hat y_j=d_{m,j}w_m\), where \(d_{m,j}=c_j^\top\otimes h_m(u_{m,j})^\top\) and \(c_j=(c_{j,1},\ldots,c_{j,I})\). For target vectors \(z_s\), linear observation operators \(A_s\), and ridge coefficient \(\lambda_m\ge 0\), solve \(\min_{w_m}\sum_s\|A_sD_mw_m-z_s\|_2^2+\lambda_m\|w_m\|_2^2\), giving the stated ridge normal equations. Here \(D_m\) stacks rows \(d_{m,j}\), and the identity matrix has dimension \(H_mI\). The benefit is exact minimization of each convex block subproblem and ridge-controlled conditioning, although the joint objective remains nonconvex.

Implementation notes

Integrate this into the optimizer for a MIONet or any tensor-factorized MLP with output \(\hat y=\sum_i t_i\prod_m b_{m,i}\). Partition each branch into a nonlinear hidden network producing \(h_m\in\mathbb R^{H_m}\) and a final matrix \(W_m\in\mathbb R^{H_m\times I}\). During each hybrid training cycle: (1) run a minibatch forward pass and cache all \(h_m\), \(b_m=h_m^\top W_m\), and trunk vectors \(t\); (2) for \(m=1,\ldots,N\), recompute \(c_{j,i}=t_{j,i}\prod_{r\ne m}b_{r,j,i}\), forming the implicit row operator \(D_mw\) without materializing \(D_m\); (3) solve the ridge normal equations with 5--20 conjugate-gradient iterations. A matrix-vector product with \(D_m\) reshapes \(w\) into \(W_m\), computes \(h_mW_m\), multiplies elementwise by \(c_j\), and contracts over rank \(i\); the transpose product is implemented by reverse accumulation. (4) Take one or several AdamW steps on all hidden-layer and trunk parameters; (5) repeat the alternating branch solves and gradient steps. Use \(\lambda_m=10^{-4}\) to \(10^{-2}\) times the mean diagonal of the empirical Gram operator, and reject or damp a solve if its residual increases. The paper supplies the multilinear representation and block least-squares principle; the implementation estimates the Gram action from the current minibatch and approximates the solve with CG. First test on MNIST or CIFAR-10 with a three-branch tensorized MLP, rank \(I=32\), and hidden width 128. Compare AdamW with AdamW plus one alternating solve per branch every 10 gradient steps. Track training loss versus optimizer steps, wall-clock time, CG residual, gradient norm, and validation accuracy. Success means faster loss descent at equal forward FLOPs, lower seed variance, or higher accuracy at the same training budget. Also test larger batch sizes to determine whether the least-squares update is harmed by minibatch noise.

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 an alternating ridge least-squares MVP for a tensor-factorized multilinear model, including explicit and implicit design/transposed operators and alternating exact branch solves. The mechanism checks succeeded: implicit-vs-explicit errors were 2.66e-15 forward and 1.78e-15 transpose, ridge stationarity norms stayed below 1.5e-14 across lambda=0 to 10, and every exact block update decreased the regularized objective. Conditioning improved from 20.13 at lambda=0 to 7.28 at lambda=10; with the same initialization, ALS reached MSE 3.60e-4 versus Adam's 10.48 after 75 branch updates, but ALS was slower wall-clock.

Agent confidence
8/10
Baseline
Adam, 75 branch updates, final training MSE 10.4801, 0.0707 s
Idea
Exact alternating ridge solves, 25 sweeps/75 branch updates, final training MSE 0.0003605, 0.8602 s; condition number 20.13 at lambda=0 versus 7.28 at lambda=10

Limitations: This is a small full-batch NumPy toy regression, not MNIST/CIFAR or a full neural MIONet. It uses explicit dense normal-equation solves rather than the planned implicit conjugate-gradient implementation, tests one seed, and does not measure validation accuracy, minibatch noise, FLOPs, or multi-seed variance.

How to run: python3 als_ridge_mvp.py

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

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

Benchmark
Tabular regression (Friedman)
Model
tensorized_mlp
Paired seeds
8
Baseline mean
8.5742
Idea mean
4.3161
Effect (Δ)
-4.2581 (−49.7%; negative = idea better)
Wins
8 / 8 paired seeds
p-value
0.0081 (permutation test, 20 000 shuffles)
Smallest detectable effect
±19.2%
Mechanism
Confirmed ✓
Practical effect
Helps
Baseline tuning
swept over 3 configs
Limitations:

The experiment used a local tensorized MLP on the registered tabular optimizer track rather than stock mlp_tiny because the intervention requires multiplicative branch-factor structure. The implementation used explicit design matrices and direct solves rather than implicit conjugate gradients; only 18 epochs, 400 training samples, one batch size, and one ridge fraction were tested, with no FLOP or wall-clock efficiency comparison.

How to run: python3 bench_experiment.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

The registered built-in tabular track was evaluated with 8 paired seeds, a tuned baseline sweep, and a shared learning-rate union. Alternating ridge final-layer solves achieved significantly lower test MSE than matched Adam: 4.3161 versus 8.5742, delta=-4.2581, permutation p=0.0081, winning on all 8 seeds. The trained-model mechanism signature confirmed non-increasing regularized block objectives.

Baseline
Tuned Adam, lr=0.01, mean test MSE 8.57417118549347
Delta mean
-4.258101403713226
Idea
Alternating ridge final-layer solves, lr=0.001, mean test MSE 4.316069781780243
P value
0.0081

Benchmark result

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

Benchmark
Tabular regression (Friedman)
Model
tensorized_mlp
Paired seeds
8
Baseline mean
8.5742
Idea mean
4.3161
Effect (Δ)
-4.2581 (−49.7%; negative = idea better)
Wins
8 / 8 paired seeds
p-value
0.0081 (permutation test, 20 000 shuffles)
Smallest detectable effect
±19.2%
Mechanism
Confirmed ✓
Practical effect
Helps
Baseline tuning
swept over 3 configs
Limitations:

The experiment used a local tensorized MLP on the registered tabular optimizer track rather than stock mlp_tiny because the intervention requires multiplicative branch-factor structure. The implementation used explicit design matrices and direct solves rather than implicit conjugate gradients; only 18 epochs, 400 training samples, one batch size, and one ridge fraction were tested, with no FLOP or wall-clock efficiency comparison.

How to run: python3 bench_experiment.py

Artifacts

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