# Alternating ridge least-squares final layers

- ID: 506
- Canonical URL: https://synthcore.org/idea/506/alternating-ridge-least-squares-final-layers
- API JSON: https://synthcore.org/api/idea/506.json
- API Markdown: https://synthcore.org/api/idea/506.md
- Verification status: beats_baseline
- Source: [arXiv:2607.06976](https://arxiv.org/abs/2607.06976)
- Category: optimization
- Solves: speedup, stability, accuracy
- ML areas: optimizer, training, mlp, architecture
- Math tags: optimization, linear-algebra, tensor-decomposition, numerical-analysis
- Ratings: usefulness 7/10; difficulty 5/10; novelty 6/10

## 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.

## 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.

## Key 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.$$

## 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

- Status: beats_baseline
- Mechanism evidence: yes
- Mechanism confirmed: yes
- Practical verdict: helps
- Verdict: 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.

### Mechanism check

- Verdict: 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.
- Confidence: 8/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.

### Practical benchmark

- Paired seeds: 8
- Baseline mean: 0
- Idea mean: 0
- p-value: 0.0081
- Paired wins: 8/8
- Benchmark verdict: idea better (significant)

## Artifacts

- [als_ridge_mvp.py](https://synthcore.org/code/242/als_ridge_mvp.py)
- [bench_experiment.py](https://synthcore.org/code/242/bench_experiment.py)
- [bench_report.json](https://synthcore.org/code/242/bench_report.json)
- [report.md](https://synthcore.org/code/242/report.md)
- [report_bench_2026-08-30T212724.md](https://synthcore.org/code/242/report_bench_2026-08-30T212724.md)
- [results.json](https://synthcore.org/code/242/results.json)
- [Download all files as ZIP](https://synthcore.org/download/242)

## Disclaimer

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