Transfer Learning-Assisted Inverse Modeling in Nanophotonics Based on Mixture Density Networks

Paper Info

Title	Transfer Learning-Assisted Inverse Modeling in Nanophotonics Based on Mixture Density Networks
Authors	Liang Cheng, Prashant Singh, Francesco Ferranti
Journal	IEEE Access, 2024
Paper	DOI: 10.1109/ACCESS.2024.3383790
Code	GitHub: mdn-nanophotonics-inverse-modeling

---

Motivation

Electromagnetic (EM) solvers (e.g. COMSOL) are the workhorse of nanophotonic design, but they’re computationally expensive — design optimization requires repeated simulations. Machine learning, especially deep neural networks, has emerged as a fast alternative for both forward and inverse modeling.

Inverse modeling is particularly attractive: given a desired optical response, directly predict the design parameters — no need to couple a forward model with an optimizer.

The Mixture Density Network (MDN) is a natural fit for inverse problems because it outputs multiple possible solutions as a mixture of Gaussian distributions, each with its own importance weight. This handles the well-known multi-solution problem in inverse design (different parameter sets can produce similar spectra).

However, MDNs face three key challenges:

1. The number of mixture components K is unknown — you need to try multiple values, but each requires full training
2. Joint optimization is expensive — mixing coefficients π, means μ, and standard deviations σ must all be optimized together
3. Numerical instability — optimizing everything at once often leads to NaN losses and degenerate predictions

---

Method

Overall Architecture

Input: Optical Spectrum (101 wavelength points)
    │
    ▼
┌──────────────────────────────────────┐
│  Optional Autoencoder (dimensionality  │
│  reduction → latent space)            │
└─────────────────┬────────────────────┘
                  │
                  ▼
┌──────────────────────────────────────┐
│  MDN (K Gaussian mixture components)  │
│  Hidden layers: expand → dropout →    │
│  shrink → output: π_k, μ_k, σ_k       │
└──────────────────────────────────────┘
                  │
                  ▼
         K design parameter PDFs
      (p, w, h₁, h₂, h₃)

The MDN models the conditional probability:

$p(y|x) = \sum_{k=1}^{K} \pi_k(x) \cdot \phi_k(y; \mu_k(x), \sigma_k(x))$

where φ_k is a Gaussian PDF, π_k are mixing coefficients (sum to 1), and the outputs y are the design parameters.

The loss is the negative log-likelihood (NLL), with a small epsilon (1e-5) added for numerical stability to prevent NaN.

---

Core Innovation: Transfer Learning for MDN

Key insight: MDNs with different K values share the same hidden layer structure — only the output layer differs. So we can reuse hidden layer weights when increasing K.

Training strategy: Start with K=1 (fully train it), then incrementally grow K from 2 to 10, each time initializing from the previous K-1 model.

Two Output Layer Initialization Strategies

For π (mixing coefficients), both strategies initialize to all zeros → after softmax, π_k = 1/K, giving every component an equal starting chance.

For μ and σ:

Strategy	How the K-th (new) component is initialized	Characteristic
TL1	Randomly clones one of the K-1 existing components	Encourages diversity, exploratory
TL2	Clones the component with the highest π weight	Less randomness, solutions tend to be similar

Why this works:

• TL1 introduces controlled randomness — the cloned component can then diverge during fine-tuning to discover new solution modes
• TL2 is more conservative — it doubles down on the strongest mode, useful when you expect clustered solutions
• Both avoid the cold-start problem: instead of random initialization, the new model starts from a proven-good feature extractor

---

Validation

Experimental Setup

Item	Details
Device	Grating-based multiband absorber (5-layer metal-insulator-metal periodic structure)
Design params	p (period), w (width), h₁, h₂, h₃ (thicknesses) → 5D output
Ranges	p∈[305,415]nm, w∈[45,190]nm, h₁∈[150,295]nm, h₂∈[25,200]nm, h₃∈[80,165]nm
Constraint	p − w ≥ 200 nm (inter-cell spacing)
Sampling	3848 Sobol quasi-random samples
Spectrum	Visible range [400–700] nm, 101 wavelength points
Simulation	COMSOL Multiphysics frequency-domain solver
Split	80% / 10% / 10% (train / val / test)
Platform	M1 Mac, PyTorch, 1000 epochs with early stopping

Compared Methods

Four configurations, each trained for K = 1 through 10:

1. MDN — baseline (trained from scratch for each K)
2. MDN + TL1 — transfer learning strategy 1
3. AE + MDN — autoencoder dimensionality reduction, then MDN
4. AE + MDN + TL1/TL2 — autoencoder + transfer learning

Key Results

CPU Time:

• TL1 and TL2 dramatically reduce total training time across K=1…10 (hidden layers are reused, only output layers + fine-tuning needed)
• The autoencoder has its own training cost: it helps when TL is NOT used, but with TL the benefit disappears or reverses

Accuracy (Test NLL loss):

• Transfer learning does NOT sacrifice accuracy — TL1 and TL2 achieve comparable or even better accuracy than training from scratch (better initialization → better optima)
• Autoencoder degrades overall accuracy (information loss from dimensionality reduction)
• TL1 slightly outperforms TL2 (diversity helps fit multi-modal distributions better)

Qualitative results (K=10):

• MDN+TL1 identifies multiple design solutions more accurately than baseline MDN
• The Gaussian PDFs clearly show uncertainty ranges for each design parameter
• Predicted spectra from the top solutions closely match the target spectrum

---

Key Takeaways

Dimension	Finding
Main contribution	First to apply transfer learning to MDN multi-K exploration, solving the “unknown K + expensive training + instability” trilemma
Speed	TL drastically cuts the total time for exploring multiple K values
Accuracy	TL matches or improves on from-scratch training — it’s not just a speed-up
Autoencoder	Not recommended with TL — the speed gain is negated and accuracy drops
TL1 vs TL2	TL1 (random clone) is generally better due to higher diversity
Practical value	MDN outputs PDFs (not single values), which can serve as initialization generators for further optimization

One-sentence summary: A “grow from K=1” transfer learning strategy for MDNs enables fast and accurate multi-solution inverse design of nanophotonic structures.

---

BibTeX

@article{cheng2024transfer,
  title={Transfer Learning-Assisted Inverse Modeling in Nanophotonics Based on Mixture Density Networks},
  author={Cheng, Liang and Singh, Prashant and Ferranti, Francesco},
  journal={IEEE Access},
  year={2024},
  volume={12},
  pages={55218--55227},
  doi={10.1109/ACCESS.2024.3383790}
}