Neural Tangent Knowledge Distillation for Optical Convolutional Networks

TL;DR

This paper introduces a task- and hardware-agnostic pipeline for optical neural networks that uses Neural Tangent Knowledge Distillation to improve accuracy and facilitate deployment across diverse systems and datasets.

Contribution

It proposes a novel, generalizable training and design pipeline for optical neural networks, including a new knowledge distillation method to bridge the accuracy gap.

Findings

Improved ONN performance across multiple datasets and hardware configurations.

Effective pre-fabrication accuracy estimation based on physical constraints.

Guided fine-tuning post-fabrication to compensate for implementation errors.

Abstract

Hybrid Optical Neural Networks (ONNs, typically consisting of an optical frontend and a digital backend) offer an energy-efficient alternative to fully digital deep networks for real-time, power-constrained systems. However, their adoption is limited by two main challenges: the accuracy gap compared to large-scale networks during training, and discrepancies between simulated and fabricated systems that further degrade accuracy. While previous work has proposed end-to-end optimizations for specific datasets (e.g., MNIST) and optical systems, these approaches typically lack generalization across tasks and hardware designs. To address these limitations, we propose a task-agnostic and hardware-agnostic pipeline that supports image classification and segmentation across diverse optical systems. To assist optical system design before training, we estimate achievable model accuracy based on user-specified constraints such as physical size and the dataset. For training, we introduce Neural Tangent Knowledge Distillation (NTKD), which aligns optical models with electronic teacher networks, thereby narrowing the accuracy gap. After fabrication, NTKD also guides fine-tuning of the digital backend to compensate for implementation errors. Experiments on multiple datasets (e.g., MNIST, CIFAR, Carvana Masking) and hardware configurations show that our pipeline consistently improves ONN performance and enables practical deployment in both pre-fabrication simulations and physical implementations.