Physical Data Embedding for Memory Efficient AI

TL;DR

The paper introduces physical embedding, a learning paradigm where data representations are encoded directly into the coefficients of a master equation, reframing PDEs as trainable architectures. The Nonlinear Schrödinger Network (NSN) uses cascaded NLSE layers with trainable parameters $\alpha$, $\beta_2$, and $\gamma$ to perform data transformations, achieving comparable or better time-series classification accuracy with orders of magnitude fewer parameters, while preserving interpretability through physically meaningful components. An extension to the Gross-Pitaevskii Equation demonstrates the approach’s generality, and ablations quantify the contributions of dispersion and nonlinearity to performance. The discussion highlights potential analog optical implementations for ultrafast inference and acknowledges limitations in generalizability, pointing to broader applicability and hardware acceleration as key future directions, with mathematical guidance provided by NLSE and GPE formulations such as $\partial E(t,z)/\partial z = -\frac{\alpha}{2} E(t,z) - i\frac{\beta_2}{2} \frac{\partial^2 E}{\partial t^2} + i\gamma|E|^2 E$ and $i\hbar\frac{\partial\psi}{\partial t} = abla$-terms.

Abstract

Deep neural networks (DNNs) have achieved exceptional performance across various fields by learning complex, nonlinear mappings from large-scale datasets. However, they face challenges such as high memory requirements and computational costs with limited interpretability. This paper introduces an approach where master equations of physics are converted into multilayered networks that are trained via backpropagation. The resulting general-purpose model effectively encodes data in the properties of the underlying physical system. In contrast to existing methods wherein a trained neural network is used as a computationally efficient alternative for solving physical equations, our approach directly treats physics equations as trainable models. We demonstrate this physical embedding concept with the Nonlinear Schrödinger Equation (NLSE), which acts as trainable architecture for learning complex patterns including nonlinear mappings and memory effects from data. The network embeds data representation in orders of magnitude fewer parameters than conventional neural networks when tested on time series data. Notably, the trained "Nonlinear Schrödinger Network" is interpretable, with all parameters having physical meanings. This interpretability offers insight into the underlying dynamics of the system that produced the data. The proposed method of replacing traditional DNN feature learning architectures with physical equations is also extended to the Gross-Pitaevskii Equation, demonstrating the broad applicability of the framework to other master equations of physics. Among our results, an ablation study quantifies the relative importance of physical terms such as dispersion, nonlinearity, and potential energy for classification accuracy. We also outline the limitations of this approach as it relates to generalizability.

Figures (5)

• Figure 1: Comparison of Traditional Convolutional Feature Learning with Physical Embedding. Trainable steps are shown in orange. (a) Convolutional Neural Network with traditional convolutional feature learning layers followed by flattening, nonlinear activation, and linear classifier. (b) Physical Embedding Network wherein a PDE replaces the convolutional layers of a neural network. This embedding utilizes the PDE's characteristics as the parameters for a deep learning framework. The physical properties of the system are used as learnable parameters, which preserves the elegance and efficiency inherent to the underlying physics. This effectively embeds the complex patterns and relationships of the data into a handful of physically meaningful properties, such as the radius of and distance between gears, in this analogy. The transformed data can then be efficiently separated by a linear classifier.
• Figure 2: Schematic representation of the Nonlinear Schrödinger Network. The input data $x$ is treated as an input $E_{in}(t)$ propagating through a virtual medium. The network consists of $M$ cascaded Nonlinear Schrödinger Layers, each comprising a linear transformation parameterized by $\alpha$ and $\beta_2$, followed by a nonlinear transformation parameterized by $\gamma$. These layers transform the input according to the Nonlinear Schrödinger Equation (NLSE), resulting in the transformed output $E_{out}(t)$. The output is then mapped to the predicted label $\hat{y}$. The trainable parameters $\{\alpha, \beta_2, \gamma\}$ across all the layers are optimized through backpropagation and gradient descent to minimize the loss between the predicted and true labels.
• Figure 3: Model architectures of (a) baseline (linear classifier), (b) Multi-Layer Perceptron (MLP), (c) Convolutional Neural Network (CNN), (d) Long Short-Term Memory (LSTM) with Attention:.(d) Nonlinear Schrödinger Network. The evolutions of data dimensions are visualized in blue rectangles. The "embedding" layers are highlighted with the red dashed box.
• Figure 4: Learning curves of the Nonlinear Schrödinger Network trained on the Starlight dataset. The progression of loss (left) and accuracy (right) for both training and validation datasets over the training epochs are demonstrated here. The curves have been smoothed using a weighted moving average to enhance visual clarity.
• Figure 5: The convergence of $\alpha$ or attenuation, $\beta_2$ or group delay dispersion, and $\gamma$ or nonlinearity, of the physical system in the first four layers of the Nonlinear Schrödinger Network during training on the Starlight dataset. (a), (b), and (c) depict the convergence of $\alpha, \beta_2,$ and $\gamma$, respectively. All parameters start with zero and gradually converge to values that optimize the network's classification performance.