Hardware-Software Co-design of Slimmed Optical Neural Networks

Transcription

1 Hardware-Sofware Co-design of Slimmed Opical Neural Neworks Zheng Zhao 1, Derong Liu 1, Meng Li 1, Zhoufeng Ying 1, Lu Zhang 2, Biying Xu 1, Bei Yu 2, Ray Chen 1, David Pan 1 The Universiy of Texas a Ausin 1 The Chinese Universiy of Hong Kong 2

2 Inroducion Emergence of dedicaed AI acceleraors Opical neural nework processor: ligh in and ligh ou» Speed-of-ligh floaing poin marix-vecor muliplicaion» >100GHz deecion rae» Ulra-low energy consumpion if configured Grea number of componens, sensiiviy o noise [Shen+, Naure Phoonics 2017]

3 Previous Opical Neural Nework (ONN) Inpu layer W Hidden layer Oupu layer in U Σ V* σ ou [Shen+, Naure Phoonics 2017] SVD decompose W = U Σ V* U and V* are uniary marices A uniary X saisfies XX* = I Implemened Mos area by expensive Mach-Zehnder inerferomeers array Σ is a diagonal marix Diagonal values are non-negaive real Implemened by opical aenuaors σ is non-linear acivaion Implemened by saurable absorber 3

4 Implemening Uniary U and V* Mach-Zehnder inerferomeers (MZI) for U and V* A single MZI implemens a 2-dim uniary in coupler ϕ coupler ou An array of n(n-1)/2 MZIs implemens an n-dim uniary in Ti,j ou i h row j h row i h col. j h col. Given an n-dim uniary, φ s can be uniquely compued 4

5 Previous ONN overview (m x n) (n x 1) W (m x 1) in U Σ V* σ ou (m x m) (m x n) (n x n) Layer size measured by # of MZIs = m(m-1)/2+n(n-1)/2 Sofware raining and hardware implemenaion Train W direcly in sofware à SVD-decomp o obain U, Σ, V* Sofware Training W SVD decomp Opical Implemenaion U Σ V*

6 Slimmed Archiecure (m x n) W (n x 1) (m x 1) T U Σ in σ ou (m x n) (n x n) (n x n) T: sparse ree nework U: uniary nework Σ: diagonal nework Use less # of MZIs = n(n-1)/2 same consrains as he previous archiecure 1 uniary marix o mainain he expressiviy An area-efficien ree nework o mach he dimension 6

7 Co-design Overview An arbirary weigh W is no TUΣ-decomposable Co-design soluion: raining and implemenaion are coupled T and Σ: Train he device parameers, consrains embedded U: Add uniary regularizaion hen approximae wih rue uniary Sofware Training Opical Implemenaion T = T Previous Train and Impl. U wih reg. Σ Approx. = U Σ Sofware Training W SVD decomp Opical Implemenaion U Σ V* 7

8 Sparse Tree Nework Sparse Tree nework (T) o mach he differen dimension Suppose in-dim > ou-dim α: linear ransfer coefficien x1 1s subree x2 xn y N x 1subree in 2nd subree ou 3rd subree 8

9 Sparse Tree Nework Implemenaion Implemened wih MZIs or direcional couplers A 2 x 1 subree x1 x2 2 x 1 subree y can be Implemened wih a single-ou MZI or direcional coupler in coupler ϕ coupler ou (energy conservaion) 9

10 Sparse Tree Nework Implemenaion Any N-inpu subree wih arbirary α s saisfying energy conservaion can be implemened i by cascading (N-1) single-ou MZIs. Energy conservaion embedded in raining Sofware Training Opical Implemenaion T U wih reg. Σ = Approx. = T U Σ

11 Uniary Nework in Training For uniary nework U saisfying UU* = I, add he regularizaion reg = UU* I F Training loss funcion Loss = Daa Loss + Regularizaion Loss leading o a near-implemenable ONN wih high accuracy Trained U ~ uniary bu only rue uniary is implemenable by MZIs 11

12 Uniary Nework in Implemenaion Approximae U by a rue uniary Ua SVD-decompose U = PSQ* à Ua = PQ* Claim. Minimize he regularizaion find he bes approximaion Min. reg Min. U - Ua F Sofware Training T = Opical Implemenaion T U wih reg. Σ Approx. = U Σ

13 Simulaion Resuls Implemened in TensorFlow for various ONN seup N1: (14 14) N4: (14 14) N7: (14 14) N2: (14 14) N5: (28 28) N8: (28 28) N3: (28 28) N6: (28 28) N9: (28 28) Tesed i on Inel Core i9-7900x CPU and an NVIDIA TianXp GPU Performed on he handwrien digi daase MNIST 13

14 Simulaion Resuls # of MZIs N1~N9: nework configuraions Our archiecure uses 15%-38% less MZIs Accuracy Similar accuracy (~0 accuracy loss) Maximum loss is Average is

15 Noise Robusness Beer resilience due o less cascaded componens Previous ONN Our ONN Accuracy Accuracy Noise Ampliude 15 Noise Ampliude

16 Training Curve Regularizaion Accuracy Epoch Epoch Accuracy Regularizaion Accuracy Regularizaion Converged in 300 epochs Balance of he accuracy and he uniary approximaion Epoch 16

17 Conribuions of This Work An new archiecure for ONN Area-efficiency ~0 accuracy loss Beer robusness o noise Hardware and sofware co-design mehodology Sofware-embedded hardware parameers Hardware consrains guaraneed by sofware 17

18 Fuure Work Beer MZI pruning mehods ~0 phase MZI à pruned + accuracy recover MZI-sparse uniary marix Design for robusness Adjus noise disribuion in raining Online raining ONN for oher neural nework archiecures CNN, RNN, ec. 18

19 19