ADEPT: Automatic Differentiable DEsign of Photonic Tensor Cores (2112.08703v2)

Abstract: Photonic tensor cores (PTCs) are essential building blocks for optical AI accelerators based on programmable photonic integrated circuits. PTCs can achieve ultra-fast and efficient tensor operations for neural network (NN) acceleration. Current PTC designs are either manually constructed or based on matrix decomposition theory, which lacks the adaptability to meet various hardware constraints and device specifications. To our best knowledge, automatic PTC design methodology is still unexplored. It will be promising to move beyond the manual design paradigm and "nurture" photonic neurocomputing with AI and design automation. Therefore, in this work, for the first time, we propose a fully differentiable framework, dubbed ADEPT, that can efficiently search PTC designs adaptive to various circuit footprint constraints and foundry PDKs. Extensive experiments show superior flexibility and effectiveness of the proposed ADEPT framework to explore a large PTC design space. On various NN models and benchmarks, our searched PTC topology outperforms prior manually-designed structures with competitive matrix representability, 2-30x higher footprint compactness, and better noise robustness, demonstrating a new paradigm in photonic neural chip design. The code of ADEPT is available at https://github.com/JeremieMelo/ADEPT using the https://github.com/JeremieMelo/pytorch-onn (TorchONN) library.

• The paper introduces ADEPT, a novel framework that uses gradient-based methods to automate photonic tensor core design, achieving up to 30× improvement in footprint efficiency.
• It leverages a probabilistic photonic SuperMesh and integrates foundry-specific constraints to explore vast design spaces and optimize circuit configurations.
• The methodology demonstrates superior noise robustness and adaptability across various process design kits, setting a new paradigm for photonic neurocomputing.

Automatic Differentiable Design of Photonic Tensor Cores via ADEPT

The paper introduces ADEPT, a novel framework for the automatic design of photonic tensor cores (PTCs). PTCs are critical components in photonic accelerators, which enable the acceleration of neural networks using photonic integrated circuits (PICs). The current methods for PTC design include manual construction or rely heavily on matrix decomposition theories, such as SVD, which have inherent limitations regarding adaptability to varying hardware constraints and device specifications. The paper positions itself by proposing an alternative that leverages design automation driven by AI, thereby transcending traditional manual design methodologies.

Proposed Framework: ADEPT

ADEPT is the first framework to offer a differentiable methodology for exploring and optimizing the design space of PTCs. ADEPT automates the PTC design process by efficiently searching circuit topologies adaptive to diverse foundry process design kits (PDKs) and footprint constraints. Key aspects of ADEPT include:

• Probabilistic Photonic SuperMesh: ADEPT constructs a probabilistic SuperMesh that includes a comprehensive range of potential PTC configurations. This allows the exploration of a vast search space for the optimal PTC topology in a differentiable manner.
• Differentiable Optimization Techniques: It uses gradient-based methods to optimize complex circuit configurations like waveguide connections and coupler placements through modified methods such as augmented Lagrangian and binarization-aware training.
• Footprint Constrained Design: ADEPT integrates the foundry-specific device specifications into the optimization process, ensuring that the resulting PTC designs meet specified footprint constraints while maximizing performance metrics.

Numerical Results and Evaluation

The methodology delineated in the paper is corroborated by extensive empirical validation. The authors demonstrate that their approach can effectively discover PTC designs that outperform previous designs in both footprint efficiency and computational robustness when applied to various neural network structures:

• Compactness vs. Performance: ADEPT designed PTCs demonstrate competitive matrix representability with significantly reduced circuit footprints. Notably, for configurations like 16×16 PTCs, ADEPT designs achieve up to 30× improvement in footprint compactness compared to manual designs.
• Adaptability to Foundry PDKs: The framework effectively adapts to distinct foundry PDKs, as shown in comparative studies between American Micro Foundry (AMF) and AIM Photonics PDKs. This adaptability translates into optimizing different device specifications to meet various performance and area constraints.
• Robustness: ADEPT-designed PTCs exhibit superior noise robustness, performing well under phase noise interference scenarios that were challenging for traditionally designed PTCs.

Implications and Future Research Directions

The introduction of ADEPT heralds a new paradigm in photonic circuit design by incorporating machine learning techniques to automate complex design tasks. The successful application of differentiable neural architecture search (DNAS) strategies to hardware design realms, such as photonic neurocomputing, opens many new avenues for future research and development. Notably, this could lead to more efficient and adaptable optical computing devices, broadening the horizons for their integration into more complex AI-driven applications.

Future developments could refine this approach through expanded search spaces, improved hardware-software co-design strategies, and increased interoperability with emerging photonic technologies. The scalability potential of ADEPT's methodology, combined with its adaptable design process, suggests significant applicability beyond the immediate scope of PTCs to other domains of optical and electronic systems.

In conclusion, ADEPT presents itself as a sophisticated tool in the search for efficient PTC architectures, coupling advanced design automation techniques with a keen consideration for practical foundry-specific constraints, thus setting the stage for the next iteration of photonic neurocomputing architectures.