Graph Neural Networks for Autonomous Multi-Scale Design of Optoelectronic Nanoelectronic Devices

Summary

This study presents an AI-driven framework for autonomously designing high-performance nanoelectronic devices. It models devices across multiple scales and optimizes them using quantum-informed reward functions. The framework performs constrained topology optimization to generate non-intuitive device configurations while maintaining compliance with specified quantum and structural constraints.

Abstract

This study presents a computational framework for autonomous multi-scale design of optoelectronic nanoelectronic devices using graph-based machine learning. The dataset consists of approximately 120,000 atomic-scale graphs, 35,000 mesoscopic-scale graphs, and 12,000 device-level graphs generated from physics-based simulations, each annotated with structural descriptors and reference energy values. Devices are represented as hierarchical graphs spanning atomic, mesoscopic, and device levels, with nodes encoding material or functional units and edges encoding physical interactions. Three scale-specific Graph Neural Networks (GNNs) are implemented using a PyTorch and PyTorch Geometric pipeline, with four message-passing layers per scale and cross-scale attention-based feature fusion. Quantum behavior is incorporated through a GNN-parameterized effective Hamiltonian, trained to approximate reference tight-binding and DFT-inspired energy spectra using a physics-regularized loss function. Device topology evolution is formulated as a constrained reinforcement learning problem in which graph modifications are treated as actions and optimized using a policy-gradient method with entropy regularization. A composite reward function balances internal quantum efficiency, structural stability, and graph complexity under physical feasibility constraints. Multi-objective optimization is performed using an augmented Lagrangian formulation to identify Pareto-consistent device configurations. Model evaluation includes energy prediction accuracy (MAE/RMSE), constraint violation rate, topology feasibility, and scalability analysis on graphs up to 10⁶ nodes. Selected designs are further validated using finite-element multiphysics simulations to verify optical and electrical consistency. This protocol provides a reproducible, multi-scale computational pipeline for physics-constrained autonomous nanoelectronic device design.

Introduction

The rapid proliferation of Internet-of-Things (IoT) systems, wearable and implantable devices, and artificial intelligence (AI) technologies has increased demand for tightly integrated electronic and photonic platforms^1,2. Optoelectronic components, including photonic integrated circuits (PICs), photodetectors, and pulse lasers, play a central role in such systems^3,4. At nanometer scales, however, device behavior is governed by quantum confinement, charge transport, and coupled optical–electrical interactions⁵. Traditional design approaches rely on first-principles simulations and empirical parameter tuning⁶. As the number of possible material combinations, geometrical layouts, and connectivity patterns grows, exhaustive simulation-based screening becomes computationally impractical.

Machine learning (ML) techniques have been introduced to accelerate modeling and inverse design tasks⁷. Early implementations used artificial neural networks and Gaussian processes as surrogate predictors⁸. More recent deep learning (DL) architectures improved representation learning for photonic and nanoelectronic systems⁹. Despite these advances, existing DL-based approaches often suffer from limited interpretability, strong data dependency, and restricted generalization across materials and structural scales¹⁰. Many frameworks also separate physics constraints from learning, treating them as external validation rather than integrated components of the optimization process¹¹.

In parallel, continued scaling of semiconductor devices has approached physical limits in conventional architectures¹². Short-channel effects and electrostatic instability become increasingly significant below 100 nm¹³, motivating exploration of alternative materials and device structures. Two-dimensional (2D) materials, particularly transition metal dichalcogenides (TMDs), offer tunable bandgaps and atomic-scale thickness compatible with CMOS processes¹⁴. Monolayer materials such as MoS₂ and WS₂ exhibit favorable carrier mobility and electrostatic control¹⁵. Junctionless architectures further demonstrate that device performance is strongly dependent on structural topology as well as material composition^16,17.

Graph Neural Networks (GNNs) provide a natural computational framework for modeling relational systems such as molecular lattices, nanostructures, and circuit topologies^18,19,20,21. By operating on graph-structured representations rather than grid-based inputs, GNNs update node embeddings through message passing among neighboring components. In simplified form, node updates can be expressed as

where ϕ_node and ψ_edge are learnable functions and denotes a permutation-invariant aggregation operator. Such models have been applied to property prediction and topology inference in nanoelectronic systems^22,23,24. Multi-scale approaches attempt to bridge atomic-scale simulations and device-level performance²⁵, yet most implementations operate at a single structural resolution or assume fixed topologies. In quantum mechanical formulations, electronic structure is governed by the eigenvalue problem

where is the Hamiltonian operator and E_k are energy eigenvalues. Density Functional Theory (DFT) provides tractable approximations but scales approximately as O(N³) with electron number, limiting direct application to large-scale topology exploration. As device complexity increases, repeated DFT evaluations become computationally prohibitive. These limitations highlight the need for differentiable approximations that retain physical structure while enabling scalable optimization^26,27,28.

To address this methodological gap, nanoelectronic devices are formulated in this work as hierarchical graphs G=(V,E) spanning atomic, mesoscopic, and device scales^{29,30,31,32,33} Scale-specific GNN encoders extract structural embeddings, and cross-scale information exchange is implemented through attention-based fusion. Instead of repeatedly solving full DFT equations during optimization, a GNN-parameterized effective Hamiltonian  is introduced to approximate reference energy behavior while preserving Hermitian structure and locality constraints^{34,35,36,37,38,39,40,41}.

Topology modification is formulated as a reinforcement learning (RL) problem⁴². For a state s_t ≡ G_t , actions modify graph connectivity under physical feasibility constraints. A composite reward function is defined as

r_t = α⋅IQE(G_t) + β⋅Stability(G_t) - γ⋅Complexity(G_t), 

where the coefficients α,β,γ regulate trade-offs among performance, robustness, and structural simplicity. Multi-objective optimization is handled through constrained formulations of the form.

Scalability to large systems is achieved using spectral graph coarsening and hierarchical aggregation, enabling efficient learning and inference on graphs containing up to millions of nodes without compromising structural fidelity. Spectral graph coarsening reduces computational complexity by clustering nodes based on eigenstructure similarity, preserving key topological and spectral properties while significantly decreasing graph size. Hierarchical aggregation further organizes the reduced graph into multi-resolution representations, allowing the model to process local and global dependencies in a staged manner. This multi-level strategy lowers memory consumption, accelerates message passing operations, and maintains predictive accuracy, thereby ensuring that the framework remains computationally tractable for large-scale nanoelectronic and optoelectronic device modeling^{43,44,45,46,47,48,49,50,51}.

The objective of this study is therefore to implement and evaluate a reproducible, multi-scale, physics-informed GNN framework for constrained nanoelectronic device design The overall workflow of the AI-driven multi-scale nanoelectronic design framework integrates dataset preparation, hierarchical graph construction, quantum-informed modeling, reinforcement learning–based topology evolution, and physics-based validation (Figure 1). The following sections describe the computational protocol, training configuration, and validation procedures.

Protocol

1. Dataset Preparation and Multi-Scale Graph Construction

1. Collect and Structure Input Data
   1. Collect atomic-scale simulation data (e.g., DFT or tight-binding outputs), mesoscopic structural data, and device-level geometries with associated performance metrics.
   2. Ensure each data record contains atomic coordinates (nm), material identifiers, energy values (eV), connectivity information, and device-level targets (e.g., internal quantum efficiency, optical absorption).
   3. Store all data in a structured format (.json or .hdf5). Verify consistent units across all records.
   4. Ensure no missing atomic coordinates or inconsistent energy units. Mixed units will corrupt normalization and training stability.
      NOTE: Use a dataset containing at least 50,000 atomic-scale graphs and at least 10,000 device-level samples
2. Clean, Normalize, and Split Data
   1. Remove incomplete entries. Normalize continuous features using z-score normalization.
   2. Encode categorical material types using one-hot encoding or learned embeddings.
   3. Split the dataset at the device level into 70% training data, 15% validation data, and 15% testing data.
   4. During training only, apply Gaussian positional noise at 1%–3% of the characteristic length scale and apply label noise up to 2%.
      NOTE: Perform dataset split before graph construction to prevent topology leakage.
3. Construct Atomic, Mesoscopic, and Device Graphs
   1. Define atomic-scale nodes as individual atoms. Define mesoscopic-scale nodes as clusters, layers, or interfaces. Define device-scale nodes as functional components such as electrodes and active regions.
   2. Define atomic-scale edges as interatomic bonds or cutoff-based interactions. Define mesoscopic-scale edges as effective coupling parameters. Define device-scale edges as electrical or optical connectivity relationships.
   3. Generate adjacency matrices.
   4. Store node and edge attributes separately for each scale.
4. Build Hierarchical Multi-Scale Representation
   1. Construct atomic graph G^(a). Apply spectral graph coarsening to produce mesoscopic graph G^(m).
   2. Aggregate mesoscopic clusters into device graph G^(d). Store cross-scale mapping indices.
5. Apply Cross-Scale Consistency Constraints
   1. Compute band alignment loss:
      L_band = ∑ ∣∣ E^(m) - A(E^(a)) ∣∣² + ∑ ∣∣ E^(d) - B(E^(m)) ∣∣² 
   2. Compute Laplacian smoothness regularization:
      
   3. Add both terms to the training objective.

2. Initialize Multi-Scale GNN Modules

1. Define node embedding dimension = 128.
   1. Define edge embedding dimension = 64.
   2. Use four message-passing layers per scale.
   3. Use ReLU activation. Initialize weights using Xavier initialization.
2. Implement Message Passing
   1. Implement edge update function as a multilayer perceptron (MLP).
   2. Implement node update function as a multilayer perceptron (MLP).
   3. Perform same-scale aggregation. Apply cross-scale attention-based fusion using Q/K/V projections.
      NOTE: Ensure consistent embedding dimensionality across scales before fusion.

3. Quantum-Informed Modeling

1. Define GNN-Parameterized Hamiltonian
   1. Replace explicit DFT Hamiltonian with:
      
   2. Enforce Hermiticity:
      f_θ(h_i,h_j) = f_θ(h_j,h_i)^* 
   3. Truncate interactions beyond predefined graph distance.
2. Define Physics-Constrained Loss
   1. Use:
      
   2. Tune λ parameters using validation set.
3. Enforce Physical Constraints
   1. Project Hamiltonian to Hermitian space each iteration.
   2. Penalize conservation law violations. Clip gradients to norm ≤ 5.

4. Reinforcement Learning for Topology Evolution

1. Define State and Action Space
   1. Represent device graph as state G_t .
   2. Define actions:
      1. Add edge
      2. Remove edge
      3. Modify connectivity.
   3. Apply action masking for illegal edits.
2. Define Reward Function
   r_t = αIQE + βStability - γComplexity 
   Tune α, β, γ via validation.
3. Train Policy Network
   1. Use policy gradient:
      ∇_θJ = E[∑∇_θlog π_θ(a_t∣s_t)A_t] 
   2. Set the discount factor γ to 0.99. Set the entropy coefficient to 0.02 and anneal during training. Use the Adam optimizer with a learning rate of 1 × 10⁻⁴ and weight decay of 1 × 10⁻⁵.
   3. Use a batch size of 32. Train for 250–300 epochs and apply early stopping based on validation convergence.

5. Multi-Objective Optimization

1. Formulate constrained objective. Apply the augmented Lagrangian.
2. Generate Pareto frontier. Retain non-dominated solutions.

6. Transfer Learning

1. Pre-train atomic GNN on energy prediction task. Save weights.
2. Fine-tune on mesoscopic and device tasks with reduced learning rate (1×10⁻⁵). Evaluate MAE and R².

7. COMSOL Multiphysics Validation

1. Import AI-Designed Geometry
   1. Export device topology as CAD file.
   2. Import into COMSOL Multiphysics.
2. Configure Simulation
   1. Use frequency-domain electromagnetic solver. Use semiconductor drift–diffusion solver.
   2. Use stationary thermal solver. Apply perfectly matched layers, ohmic contacts, and realistic excitation wavelengths.
3. Extract Metrics
   1. Measure optical absorption efficiency, internal quantum efficiency, electric-field enhancement, and peak temperature rise.
   2. Average results across excitation conditions.

8. Scalability Testing

1. Increase graph size up to 10⁶ nodes. Measure runtime and memory usage.
2. Enable spectral coarsening, sparse adjacency, mixed precision, and distributed GPU training.

9. Statistical Evaluation

1. Run three independent training seeds.
2. Report mean ± standard deviation, MAE, RMSE, and R².
3. Use identical splits and hyperparameters for baseline comparison.

10. Data and Code Archiving

1. Save trained models. Archive datasets.
2. Provide scripts for graph construction, training, and COMSOL validation.

Results

The proposed multi-scale GNN-reinforcement learning framework was evaluated on 120,000 atomic graphs, 35,000 mesoscopic graphs, and 12,000 device-level graphs. Data were split at the device level into 70% training, 15% validation, and 15% testing sets. All reported values represent mean ± standard deviation across three independent runs using different random seeds. Baseline methods included SchNet, DimeNet++, physics-informed neural networks (PINNs), generative topology optimization, and a physics-only solver baseline. ...

Discussion

This study presents a hierarchical, physics-informed graph-based framework for autonomous nanoelectronic and optoelectronic device design. The Results demonstrate that the method achieves competitive energy prediction accuracy, stable reinforcement-learning-driven topology evolution, constraint satisfaction, and physically validated improvements in device performance. The following discussion focuses on critical procedural elements, methodological considerations, limitations, contextual significance, and future extension...

Materials

List of materials used in this article

Name	Company	Catalog Number	Comments
Computational & Modeling Software
JAX / Haiku		N/A	0.3.0+
LAMMPS		N/A	-
LUMERICAL / MEEP		N/A	-
OpenAI Gym / Gymnasium		N/A	-
Python		N/A	3.8+
PyTorch / PyTorch Geometric		N/A	1.9.0+ / 2.0+
Quantum ESPRESSO		N/A	6.7+
Data & Datasets
DFT Simulation Database		N/A	Custom
Generated Nanoelectronic Graphs		N/A	~10,000 graphs
Materials Project / OQMD		N/A	Public API
Hardware & Infrastructure
Data Storage		N/A	NVMe SSDs
High-Performance Computing (HPC) Cluster	N/A	CPU Nodes
NVIDIA GPU Cluster		N/A	A100 / V100 / H100
Key Algorithms & Model Components
Adam / AdamW Optimizer		N/A	PyTorch Native
Graph Attention Network (GAT)		N/A	Custom (PyTorch)
Graph Variational Autoencoder (GVAE)		N/A	Custom (PyTorch)
Message Passing Neural Network (MPNN)	N/A	Custom (PyTorch)
Policy Gradient (e.g., PPO)		N/A	Custom (Stable-Baselines3)
Machine Learning & Optimization Libraries
NetworkX		N/A	2.6+
NumPy		N/A	1.24+
Optuna		N/A	2.0+
Scikit-learn		N/A	1.0+
Stable-Baselines3 / Ray RLLib		N/A	-
Physical & Mathematical Models
Augmented Lagrangian Method		N/A	Custom (Python)
Composite Reward Function (rt)		N/A	Custom (Python)
Parameterized Hamiltonian (Hθ)		N/A	Custom (PyTorch)
Spectral Graph Coarsening		N/A	Custom (Scikit-learn)