In-network optical inference
In-network Optical Inference (IOI) provides low-latency machine learning inference by leveraging programmable switches and optical matrix multiplication. IOI uses a transceiver module, called a Neuro Transceiver, with an optical processor to perform linear operations, such as matrix multiplication, in the optical domain. IOI's transceiver modules can be plugged into programmable packet switches, which are programmed to perform non-linear activations in the electronic domain and to respond to inference queries. Processing inference queries at the programmable packet switches inside the network, without sending them to cloud or edge inference servers, significantly reduces end-to-end inference latency experienced by users.
1 . A method of inference processing, the method comprising:
transmitting a first data packet containing an input vector to a layer of neurons in an artificial neural network from a client to a programmable packet switch comprising a programmable packet processing chip and an optical computing unit;
removing source and/or destination metadata from the first data packet;
after removing the source and/or destination metadata from the first data packet, concatenating, by the programmable packet processing chip, a weight vector corresponding to the layer of neurons with the first data packet;
transmitting the first data packet and the weight vector from the programmable packet processing chip to the optical computing unit;
computing a product of the input vector and the weight vector with the optical computing unit;
transmitting a second data packet containing the product of the input vector and the weight vector from the optical computing unit to the programmable packet processing chip; and
applying a nonlinear activation function to the product with the programmable packet processing chip.
2 . The method of claim 1 , further comprising:
storing the destination metadata in a memory while computing the product of the input vector and the weight vector and applying the nonlinear activation function;
adding the destination metadata to an output packet containing an output of the artificial neural network; and
transmitting the output packet from the programmable packet switch to a destination specified in the destination metadata.
3 . The method of claim 1 , wherein applying the nonlinear activation function comprises retrieving information from a match-action table programmed into a memory of the programmable packet processing chip.
4 . The method of claim 1 , wherein the layer of neurons is a first layer of neurons of the artificial neural network, the weight vector is a first weight vector, and further comprising:
concatenating, by the programmable packet processing chip, a second weight vector corresponding to a second layer of neurons in the artificial neural network to the second data packet; and
transmitting the second data packet and the second weight vector from the programmable packet processing chip to the optical computing unit.
5 . The method of claim 1 , further comprising, before transmitting the second data packet:
digitizing the product of the input vector and the weight vector; and
storing the product in a memory.
6 . A method of inference processing, the method comprising:
transmitting a first data packet containing an input vector to a layer of neurons in an artificial neural network from a client to a programmable packet switch comprising a programmable packet processing chip and an optical computing unit;
concatenating, by the programmable packet processing chip, a weight vector corresponding to the layer of neurons with the first data packet;
transmitting the first data packet and the weight vector from the programmable packet processing chip to the optical computing unit;
computing a product of the input vector and the weight vector with the optical computing unit;
transmitting a second data packet containing the product of the input vector and the weight vector from the optical computing unit to the programmable packet processing chip; and
applying a nonlinear activation function to the product with the programmable packet processing chip,
wherein computing the product of the input vector and the weight vector comprises:
modulating, with a first modulator, an optical pulse with a waveform proportional to the input vector;
modulating, with a second modulator in optical communication with the first modulator, the optical pulse with a waveform proportional to the weight vector; and
detecting the optical pulse with a photodetector in optical communication with the second modulator.
7 . The method of claim 6 , further comprising:
delaying the optical pulse between the first modulator and the second modulator by a delay equal to a duration of the waveform proportional to the input vector.
8 . A programmable packet switch comprising:
a programmable packet processing chip to concatenate a first data packet representing an input vector to a first layer of neurons of an artificial neural network with a first weight vector corresponding to the first layer of neurons; and
an optical processor, operably coupled to the programmable packet processing chip, to receive the first data packet and the first weight vector from the programmable packet processing chip, to compute a product of the input vector and the first weight vector, and to transmit the product of the input vector and the first weight vector to the programmable packet processing chip as a second data packet,
wherein the programmable packet processing chip is configured to remove source and/or destination metadata from the first data packet before transmitting the first data packet and the first weight vector to the optical processor.
9 . The programmable packet switch of claim 8 , wherein the programmable packet switch further comprises:
a memory to store the destination metadata while the optical processor computes the product of the input vector and the first weight vector,
wherein the programmable packet processing chip is configured to add the destination metadata stored in the memory to an output packet containing an output of the artificial neural network.
10 . The programmable packet switch of claim 8 , wherein the programmable packet processing chip is configured to perform a nonlinear activation on the product of the input vector and the first weight vector, thereby producing an output of the first layer of neurons.
11 . The programmable packet switch of claim 10 , wherein the programmable packet processing chip comprises a memory to store a match-action table representing the nonlinear activation.
12 . The programmable packet switch of claim 10 , wherein the programmable packet processing chip is further configured to concatenate the output of the first layer of neurons with a second weight vector corresponding to a second layer of neurons of the artificial neural network and to transmit the output of the first layer of neurons and the second weight vector to the optical processor.
13 . A programmable packet switch of claim 8 comprising:
a programmable packet processing chip to concatenate a first data packet representing an input vector to a first layer of neurons of an artificial neural network with a first weight vector corresponding to the first layer of neurons; and
an optical processor, operably coupled to the programmable packet processing chip, to receive the first data packet and the first weight vector from the programmable packet processing chip, to compute a product of the input vector and the first weight vector, and to transmit the product of the input vector and the first weight vector to the programmable packet processing chip as a second data packet,
wherein the optical processor comprises:
a first modulator to modulate an optical beam with a first analog waveform representing the input vector;
a second modulator, in optical communication with the first modulator, to modulate the optical beam with a second analog waveform representing the first weight vector; and
a photodetector, in optical communication with the second modulator, to detect the optical beam.
14 . The programmable packet switch of claim 13 , further comprising:
at least one digital-to-analog converter, operably coupled to the first modulator and to the second modulator, to convert the input vector into the first analog waveform and to convert the first weight vector into the second analog waveform.
15 . The programmable packet switch of claim 13 , further comprising:
an analog-to-digital converter, operably coupled to the photodetector, to convert a photocurrent generated by the photodetector in response to detecting the optical beam into a digital signal representing the product of the input vector and the first weight vector.
16 . The programmable packet switch of claim 13 , wherein the optical processor further comprises:
an optical delay line, coupling the first modulator to the second modulator, to delay the optical beam by a delay equal to a duration of the first analog waveform.
17 . A method of inference processing, the method comprising:
receiving a packet with a header comprising source/destination metadata and a payload comprising an input to a deep neural network (DNN);
removing the source/destination metadata from the header;
adding a weight vector corresponding to a first layer of the DNN to the header;
transmitting the packet to an optical processor comprising a first modulator, a second modulator in series with the first modulator, and a photodetector;
converting the input to the DNN into a first analog waveform;
converting the weight vector into a second analog waveform;
modulating, by the first modulator, an amplitude of an optical beam with the first analog waveform;
modulating, by the second modulator, the amplitude of the optical beam with the second analog waveform;
transducing, by the photodetector, the optical beam into an electrical signal with an amplitude representing a product of the input to the DNN and the weight vector; and
performing a nonlinear activation on the electrical signal to produce an output of the first layer of the DNN.
18 . The method of claim 17 , further comprising:
transmitting a packet containing the output of the first layer of the DNN and a weight vector corresponding to a second layer of the DNN to the header to the optical processor.