DynamicSVD: fit_on_signal + >2D inputs, and actually use the mu scaffold in example 05

[cweniger]

Extracted from plan/gaussianized-flow-matching (21cee8a, 79b8b53), plus the follow-through needed to make the feature real in a shipped example.

The bug this uncovers

examples/05_linear_regression declares scaffolds: [mu] and passes _input_: [x, mu] to DynamicSVD. Its forward model is deliberately split into SignalSimulator → NoiseSimulator for the sole purpose of exposing the noiseless mu.

DynamicSVD then throws it away. On main, signal is touched in exactly one place:

if self.whitener is not None and signal is not None:
    self.whitener.update((x - signal).detach())

No whitener is attached in that config, so mu is plumbed through the entire graph and silently discarded. The example pays for a scaffold node it never uses.

What changes

fit_on_signal fits the eigenbasis on the clean signal stream while projections in forward() stay on x. Two reasons that matters:

• Basis selection. Fitting on noisy x lets noise eigenvalues — spread up to the Marchenko–Pastur edge (1+√(D/N))²σ² — masquerade as structure and displace real components from the top-k.
• The Wiener step. The whitener normalizes noise to unit variance, so fitting on x gives λ ≈ λ_s + 1 and the gain λ/(λ+1) becomes (λ_s+1)/(λ_s+2), bounded below by ½. A pure-noise direction reaches the downstream network at half amplitude instead of being suppressed. Fitting on the signal gives λ = λ_s and the textbook λ_s/(λ_s+1).

Net effect on output scale, in whitened units:

	λ used	output std at λ_s = 0 (pure noise)	λ_s → ∞
fit on x	λ_s + 1	0.50	→ 1
fit on signal	λ_s	0	→ 1

>2D inputs are flattened to (batch_size, D) internally, so image-shaped data can be fed in directly.

Commits

1. fit_on_signal flag — cherry-pick of 21cee8a. Default False, and with the flag off the expression reduces to the previous line exactly, so it is a no-op for existing callers.
2. flatten >2D inputs — cherry-pick of 79b8b53, covering update() and forward().
3. flatten in reconstruct() too — the one entry point commit 2 missed, which would still reject the inputs the other two now accept.
4. docstring — it was stale under every setting: it described signal as whitener-only, and advertised step 2 as a "Wiener filter" without noting that this holds only when λ is signal power in noise units. Now states all three configurations, including that with no whitener the + 1 sits in raw data units and the threshold is arbitrary. No behaviour change.
5. example 05 — see below.

Example 05: the whitener is not optional here

fit_on_signal alone would not have been correct in this example. With no whitener the eigenvalues are in raw data units, so λ/(λ+1) compares signal power against 1.0 when the noise variance is σ² = 0.01. The two settings only make sense together:

embedding:
  _target_: falcon.embeddings.DynamicSVD
  _input_: [x, mu]
  n_components: 32
  momentum: 0.1
  fit_on_signal: true
  whitener:
    _target_: falcon.embeddings.DiagonalWhitener
    dim: 1000   # n_bins

The whitener trains on noise = x - mu, which is zero-mean with σ = 0.1, so it drives std → σ and the eigenvalues become genuine signal-to-noise ratios.

The nested _target_ resolves correctly — builder.py:336 runs every embedding kwarg through _instantiate_sub_targets, whose docstring calls out this exact case (a raw dict would make DynamicSVD call dict.update() instead of the whitener's .update()).

⚠️ This changes the example's behaviour: different basis, differently scaled coefficients. 05 has an analytic posterior, so it is directly checkable — please re-run it and confirm it still recovers Sigma_post / mu_post before merging.

Default left at False

fit_on_signal defaults to False and is set explicitly in example 05. Flipping the constructor default remains open — worth doing once 05 has demonstrated the setting end to end, since at that point every in-repo caller already opts in. Say the word and it is a one-line follow-up.

Coverage note

Before this PR, none of the affected paths were reachable from any committed config: nothing set fit_on_signal, nothing fed >2D input, and nothing passed a whitener and a signal together — so even main's existing signal handling was dead code. Commit 5 makes the first and third live. The >2D flatten path still has no in-repo caller.

Testing

Not run — import torch fails in my environment (libcudnn.so.9: cannot open shared object file), which breaks pytest tests collection on main too. Verification here is static review, py_compile, YAML parse of the edited config, and a check that dim: 1000 matches the n_bins: 1000 on the mu simulator.

The re-run of example 05 is the real test and I cannot do it.

🤖 Generated with Claude Code

[codecov]

Codecov Report

❌ Patch coverage is 0% with 26 lines in your changes missing coverage. Please review.
✅ Project coverage is 10.21%. Comparing base (9db772f) to head (950503d).

Files with missing lines	Patch %	Lines
src/falcon/embeddings/svd.py	0.00%	26 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #119      +/-   ##
==========================================
- Coverage   10.26%   10.21%   -0.05%     
==========================================
  Files          30       30              
  Lines        3955     3973      +18     
==========================================
  Hits          406      406              
- Misses       3549     3567      +18     

Flag	Coverage Δ
unit	10.21% <0.00%> (-0.05%)	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[cweniger]

Pushed two more commits implementing the on-the-fly noise estimate — you were right that requiring a whitener was the wrong shape for this.

68ce127 — estimate σ² when no whitener is attached

The Wiener step only means something if λ and the denominator share units. A whitener guarantees that by normalizing the noise; without one the + 1 is arbitrary. Now, when a signal is available but a whitener is not, a single scalar noise variance is estimated from x - signal under a white-noise assumption (EMA, reusing the SVD blend momentum), and the gain becomes λ/(λ+σ²).

Precedence:

situation	σ²
whitener attached	1 — noise already normalized per feature
no whitener, signal given	running scalar estimate
neither	1 — nothing to estimate from

Only the middle row changes behaviour, and no config in the repo was in it.

Steps 2–3 of forward() are now one factor. Substituting λ → λ/σ² throughout and simplifying:

old:  Λ/(Λ+1)/√Λ
new:  √Λ/(Λ+σ²)

These are algebraically identical at σ²=1. Checked numerically across λ ∈ [1e-12, 1e6]: max relative difference 2.2e-16, i.e. float rounding only. The new form is also better conditioned — the old one divides by √Λ after clamping Λ at 1e-12, so by 1e-6. reconstruct() gets the same denominator.

_noise_var persists through get/set_extra_state, read back with .get() so pre-existing checkpoints still load and fall back to the old constant.

7258748 — example 05 comment

The previous comment claimed the whitener was required there. That is no longer true: 05's noise is homoscedastic (σ = 0.1 in every bin), so the scalar estimate covers it exactly and dropping the block would also be correct. I kept the whitener — it demonstrates the per-feature case you need once noise varies across bins — but the comment now says so rather than overstating it.

Consequence for the re-run

The validation ask is unchanged but now covers a bit more. Since 05 keeps its whitener it takes the σ²=1 path, so the new estimator is not exercised by that run. If you want the scalar path covered too, commenting out the whitener: block is a second cheap run against the same analytic posterior — it should land in the same place, and if it doesn't, the estimator is wrong.

[cweniger]

Pushed 950503d: shrinkage → wiener.

Renaming turned out to be free — nothing outside the class ever set the flag (no config, no example, no other module), so there is no compatibility surface and no shim is needed. grep -rn shrinkage over the repo now returns nothing.

Reasoning, for the record:

• It gated two operations and named one. In forward() the flag controls the Wiener gain and the 1/√λ rescaling. shrinkage=False therefore doesn't give unshrunk-but-still-sane coefficients — it gives raw eigenbasis projections with scale √(λ+σ²), varying by orders of magnitude across components. That second effect is the one that matters to a downstream network, and the old name said nothing about it.
• It meant different things in different methods. Wiener + rescaling in forward(), Wiener alone in reconstruct() — the latter correctly, since the rescaling would be wrong heading back to data space. Now documented explicitly.
• "Shrinkage" is already taken in this neighbourhood: in covariance estimation it means Ledoit–Wolf-style eigenvalue regularization, which a reader would go looking for in _svd_update.

I did not split it into separate wiener / normalize flags. The combination wiener=False, normalize=True would divide by √λ with no noise suppression, amplifying exactly the directions the filter exists to kill — better not to make that expressible.

The commit also fixes a docstring claim that was never accurate. Step 3 was described as "normalize to ~unit variance", but the output standard deviation is √(λ/(λ+σ²)) = √(SNR/(SNR+1)). Unit scale is only the λ ≫ σ² asymptote; suppressing the rest toward zero is the point of the step, not a side effect.

No behaviour change in this commit — pure rename plus docs.