I trained a physical-style programmable optical mesh under realistic hardware noise using only on-chip-inspired methods (SPSA + adjoint).
It still converges. Power is conserved. It degrades gracefully.
👉 Open in Colab — zero friction
Repository: https://github.com/ZanzyTHEbar/mzi
This is a clean, validated Phase 1 simulator for hexagonal Mach-Zehnder Interferometer (MZI) meshes as a programmable photonic processor being explored for ultra-low-power, high-speed AI accelerators.
Built with numpy transfer matrices + realistic noise. No approximations in the training or visualization paths.
- Model a hexagonal MZI mesh using numpy transfer matrices.
- Implement forward propagation for unitary / linear optical transforms.
- Add adjoint-variable and on-chip-style backpropagation (inspired by 2026 Nature demo + classic in-situ methods).
- Model realistic hardware effects: phase noise, thermal crosstalk, loss, fabrication variation.
- Demonstrate training with:
- LoRA-like low-rank adapters (small auxiliary sub-meshes / modulators).
- Fine-tuning from digital initialization.
- Gradient-free / RL-style optimizers (SPSA etc.) suitable for noisy analog.
- Quantify: robustness to device variation, convergence, rough energy/speed advantages vs digital baselines.
- Produce open artifacts: code, benchmarks, plots, metrics for arXiv/blog/collaboration.
# with uv (recommended)
uv sync --extra dev
uv run python -c "
from mzi.core.hexagonal_mesh import HexMesh
import numpy as np
m = HexMesh(ports=6, seed=42)
x = np.random.randn(6) + 1j * np.random.randn(6)
y = m.forward(x)
print('Power in/out (lossless approx):', np.sum(np.abs(x)**2), np.sum(np.abs(y)**2))
"Run a basic training demo:
uv run mzi-train --help
# or
uv run python -m mzi.benchmarks.runners --task moons --size 6 --noise 0.03See the Colab notebook and the live Streamlit demo for the best experience.
| Resource | Link / Command | Zero friction |
|---|---|---|
| Live interactive demo | Hugging Face Spaces (link when deployed) | ✅ |
| Colab notebook | Open in Colab | ✅ |
| Local (full demo) | uv sync --extra demo && streamlit run demo/app.py |
- |
All visualizations and training runs in the demo drive the exact production HexMesh, train_step, noise model, etc.
# Run the reproducibility cell in the notebook or:
uv run python -c "
from mzi.core.hexagonal_mesh import HexMesh
from mzi.training.loop import train_step
from mzi.benchmarks.tasks import make_classification_task
import numpy as np
m = HexMesh(ports=4, seed=42)
m.variation_params.phase_noise_std = 0.035
xs, ys = make_classification_task(n_samples=32, n_ports=4, seed=123)
for _ in range(120):
for x, y in zip(xs, ys):
train_step(m, x, y, lr=0.03, spsa_perturb=0.08)
print('Final loss (noisy, SPSA):', round(..., 4))
"See the notebook for the full sweep + internal field probes.
src/mzi/core/mzi.py: 2×2 tunable MZI (θ, φ) with non-ideal couplers.hexagonal_mesh.py: Hex lattice of TBUs, forward simulator, internal field probes.adjoint.py: Adjoint-variable method (AVM) + reverse-light on-chip-style gradients.noise.py: Static/dynamic variations, crosstalk, loss.
src/mzi/training/: losses, LoRA adapters, optimizers (GD + SPSA).src/mzi/benchmarks/: tasks (unitary fit, classification), runners, digital baselines, metrics.
Key properties validated:
- Power conservation (lossless case).
- Adjoint gradients match finite differences.
- Training converges under realistic noise where pure digital-ideal training fails.
- LoRA-style updates use far fewer parameters for adaptation.
This repo executes the Simulation & Theoretical Validation phase (0–3 months, near-zero cost) from the solo founder plan.
Later phases (community, small MPW prototypes, use-case/business validation) will use outputs from here (phase maps, robustness data, architecture specs) as foundation.
See the original vision post for the full solo-founder phased validation roadmap (simulation → community → small-scale photonic prototypes).
- Ashtiani et al., "Integrated photonic deep neural network with end-to-end on-chip backpropagation training", Nature 2026 (arXiv:2506.14575).
- Hughes et al., Optica 2018 (in-situ backprop via adjoint).
- Pai et al. (matrix optimization on unitary photonic meshes).
- neurophox (rectangular/triangular MZI mesh patterns; studied for reference).
- Hexagonal processor examples (iPronics SmartLight and related photonic processor literature).
Phase 1 core complete (simulation + training + noise + LoRA + SPSA).
See scripts/generate_demo_figures.py + figures/hero_results.json for exact assets with fixed seeds.
- Clean training final loss (4-port moons, 120 SPSA steps): 0.4866
- Noisy (σ=0.035): 0.4871 — almost no degradation
- At high noise σ=0.10: still trains to ~1.93 (graceful)
Power conservation (ideal): exact.
See the Colab notebook and figures/ for the plots.
- Training converges under realistic noise where pure digital-ideal training fails.
- LoRA-style updates use far fewer parameters for adaptation.
- Adjoint matches finite differences; SPSA is the robust recommended path.
See scripts/generate_demo_figures.py, the Colab notebook, and python -m mzi.benchmarks.runners.
- Expand mesh size & tasks (unitary + small transformer-style linear layers).
- More figures + animations (see
scripts/generate_demo_figures.py). - Community outreach + Phase 2/3 prep (MPW test structures).
Run uv run python scripts/generate_demo_figures.py to regenerate the demo plots with fixed seeds.
Photonic meshes do matrix-vector multiplies at the speed of light with (ideally) very low energy per MAC — mostly just the cost of the DACs + heaters for the phases.
This simulator lets us quantify how much "programming" robustness I actually get from SPSA + LoRA under hardware variation before going to silicon.
(Real energy numbers will depend on the eventual photonic platform.)
See the notebook for discussion.
Contributions and discussion welcome — especially from photonic researchers and on-chip training groups.
MIT