@@ -0,0 +1,244 @@
1+{-# LANGUAGE BangPatterns #-}
2+{-# LANGUAGE DataKinds #-}
3+{-# LANGUAGE GADTs #-}
4+{-# LANGUAGE RankNTypes #-}
5+{-# LANGUAGE RecordWildCards #-}
6+{-# LANGUAGE TypeOperators #-}
7+{-# LANGUAGE TypeFamilies #-}
8+{-# LANGUAGE MultiParamTypeClasses #-}
9+{-# LANGUAGE FlexibleContexts #-}
10+{-# LANGUAGE ViewPatterns #-}
11+module Grenade.Recurrent.Layers.LSTM (
12+ LSTM (..)
13+ , LSTMWeights (..)
14+ , randomLSTM
15+ ) where
16+
17+import Control.Monad.Random ( MonadRandom, getRandom )
18+
19+-- import Data.List ( foldl1' )
20+import Data.Singletons.TypeLits
21+
22+import Numeric.LinearAlgebra.Static
23+
24+import Grenade.Core.Network
25+import Grenade.Core.Shape
26+
27+import Grenade.Layers.Internal.Update
28+
29+import Grenade.Recurrent.Core.Network
30+
31+-- | Long Short Term Memory Recurrent unit
32+--
33+-- This is a Peephole formulation, so the recurrent shape is
34+-- just the cell state, the previous output is not held or used
35+-- at all.
36+data LSTM :: Nat -> Nat -> * where
37+ LSTM :: ( KnownNat input
38+ , KnownNat output
39+ ) => !(LSTMWeights input output) -- Weights
40+ -> !(LSTMWeights input output) -- Momentums
41+ -> LSTM input output
42+
43+data LSTMWeights :: Nat -> Nat -> * where
44+ LSTMWeights :: ( KnownNat input
45+ , KnownNat output
46+ ) => {
47+ lstmWf :: !(L output input) -- Weight Forget (W_f)
48+ , lstmUf :: !(L output output) -- Cell State Forget (U_f)
49+ , lstmBf :: !(R output) -- Bias Forget (b_f)
50+ , lstmWi :: !(L output input) -- Weight Input (W_i)
51+ , lstmUi :: !(L output output) -- Cell State Input (U_i)
52+ , lstmBi :: !(R output) -- Bias Input (b_i)
53+ , lstmWo :: !(L output input) -- Weight Output (W_o)
54+ , lstmUo :: !(L output output) -- Cell State Output (U_o)
55+ , lstmBo :: !(R output) -- Bias Output (b_o)
56+ , lstmWc :: !(L output input) -- Weight Cell (W_c)
57+ , lstmBc :: !(R output) -- Bias Cell (b_c)
58+ } -> LSTMWeights input output
59+
60+instance Show (LSTM i o) where
61+ show LSTM {} = "LSTM"
62+
63+instance (KnownNat i, KnownNat o) => UpdateLayer (LSTM i o) where
64+ -- The gradients are the same shape as the weights and momentum
65+ -- This seems to be a general pattern, maybe it should be enforced.
66+ type Gradient (LSTM i o) = (LSTMWeights i o)
67+
68+ -- Run the update function for each group matrix/vector of weights, momentums and gradients.
69+ -- Hmm, maybe the function should be used instead of passing in the learning parameters.
70+ runUpdate LearningParameters {..} (LSTM w m) g =
71+ let (wf, wf') = u lstmWf w m g
72+ (uf, uf') = u lstmUf w m g
73+ (bf, bf') = v lstmBf w m g
74+ (wi, wi') = u lstmWi w m g
75+ (ui, ui') = u lstmUi w m g
76+ (bi, bi') = v lstmBi w m g
77+ (wo, wo') = u lstmWo w m g
78+ (uo, uo') = u lstmUo w m g
79+ (bo, bo') = v lstmBo w m g
80+ (wc, wc') = u lstmWc w m g
81+ (bc, bc') = v lstmBc w m g
82+ in LSTM (LSTMWeights wf uf bf wi ui bi wo uo bo wc bc) (LSTMWeights wf' uf' bf' wi' ui' bi' wo' uo' bo' wc' bc')
83+ where
84+ -- Utility function for updating with the momentum, gradients, and weights.
85+ u :: forall x ix out. (KnownNat ix, KnownNat out) => (x -> (L out ix)) -> x -> x -> x -> ((L out ix), (L out ix))
86+ u e (e -> weights) (e -> momentum) (e -> gradient) =
87+ decendMatrix learningRate learningMomentum learningRegulariser weights gradient momentum
88+
89+ v :: forall x ix. (KnownNat ix) => (x -> (R ix)) -> x -> x -> x -> ((R ix), (R ix))
90+ v e (e -> weights) (e -> momentum) (e -> gradient) =
91+ decendVector learningRate learningMomentum learningRegulariser weights gradient momentum
92+
93+ -- There's a lot of updates here, so to try and minimise the number of data copies
94+ -- we'll create a mutable bucket for each.
95+ -- runUpdates rate lstm gs =
96+ -- let combinedGradient = foldl1' uu gs
97+ -- in runUpdate rate lstm combinedGradient
98+ -- where
99+ -- uu :: (KnownNat i, KnownNat o) => LSTMWeights i o -> LSTMWeights i o -> LSTMWeights i o
100+ -- uu a b =
101+ -- let wf = u lstmWf a b
102+ -- uf = u lstmUf a b
103+ -- bf = v lstmBf a b
104+ -- wi = u lstmWi a b
105+ -- ui = u lstmUi a b
106+ -- bi = v lstmBi a b
107+ -- wo = u lstmWo a b
108+ -- uo = u lstmUo a b
109+ -- bo = v lstmBo a b
110+ -- wc = u lstmWc a b
111+ -- bc = v lstmBc a b
112+ -- in LSTMWeights wf uf bf wi ui bi wo uo bo wc bc
113+ -- u :: forall x ix out. (KnownNat ix, KnownNat out) => (x -> (L out ix)) -> x -> x -> L out ix
114+ -- u e (e -> a) (e -> b) = tr $ tr a + tr b
115+
116+ -- v :: forall x ix. (x -> (R ix)) -> x -> x -> R ix
117+ -- v e (e -> a) (e -> b) = a + b
118+
119+ createRandom = randomLSTM
120+
121+instance (KnownNat i, KnownNat o) => RecurrentUpdateLayer (LSTM i o) where
122+ -- The recurrent shape is the same size as the output.
123+ -- It's actually the cell state however, as this is a peephole variety LSTM.
124+ type RecurrentShape (LSTM i o) = 'D1 o
125+
126+instance (KnownNat i, KnownNat o) => RecurrentLayer (LSTM i o) ('D1 i) ('D1 o) where
127+ -- Forward propagation for the LSTM layer.
128+ -- The size of the cell state is also the size of the output.
129+ runRecurrentForwards (LSTM (LSTMWeights {..}) _) (S1D cell) (S1D input) =
130+ let -- Forget state vector
131+ f_t = sigmoid $ lstmBf + lstmWf #> input + lstmUf #> cell
132+ -- Input state vector
133+ i_t = sigmoid $ lstmBi + lstmWi #> input + lstmUi #> cell
134+ -- Output state vector
135+ o_t = sigmoid $ lstmBo + lstmWo #> input + lstmUo #> cell
136+ -- Cell input state vector
137+ c_x = tanh $ lstmBc + lstmWc #> input
138+ -- Cell state
139+ c_t = f_t * cell + i_t * c_x
140+ -- Output (it's sometimes recommended to use tanh c_t)
141+ h_t = o_t * c_t
142+ in (S1D c_t, S1D h_t)
143+
144+ -- Run a backpropogation step for an LSTM layer.
145+ -- We're doing all the derivatives by hand here, so one should
146+ -- be extra careful when changing this.
147+ --
148+ -- There's a test version using the AD library without hmatrix in the test
149+ -- suite. These should match always.
150+ runRecurrentBackwards (LSTM (LSTMWeights {..}) _) (S1D cell) (S1D input) (S1D cellGrad) (S1D h_t') =
151+ -- We're not keeping the Wengert tape during the forward pass,
152+ -- so we're duplicating some work here.
153+ --
154+ -- If I was being generous, I'd call it checkpointing.
155+ --
156+ -- Maybe think about better ways to store some intermediate states.
157+ let -- Forget state vector
158+ f_s = lstmBf + lstmWf #> input + lstmUf #> cell
159+ f_t = sigmoid f_s
160+ -- Input state vector
161+ i_s = lstmBi + lstmWi #> input + lstmUi #> cell
162+ i_t = sigmoid i_s
163+ -- Output state vector
164+ o_s = lstmBo + lstmWo #> input + lstmUo #> cell
165+ o_t = sigmoid o_s
166+ -- Cell input state vector
167+ c_s = lstmBc + lstmWc #> input
168+ c_x = tanh c_s
169+ -- Cell state
170+ c_t = f_t * cell + i_t * c_x
171+
172+ -- Reverse Mode AD Derivitives
173+ c_t' = h_t' * o_t + cellGrad
174+
175+ f_t' = c_t' * cell
176+ f_s' = sigmoid' f_s * f_t'
177+
178+ o_t' = h_t' * c_t
179+ o_s' = sigmoid' o_s * o_t'
180+
181+ i_t' = c_t' * c_x
182+ i_s' = sigmoid' i_s * i_t'
183+
184+ c_x' = c_t' * i_t
185+ c_s' = tanh' c_s * c_x'
186+
187+ -- The derivatives to pass sideways (recurrent) and downwards
188+ cell' = tr lstmUf #> f_s' + tr lstmUo #> o_s' + tr lstmUi #> i_s' + c_t' * f_t
189+ input' = tr lstmWf #> f_s' + tr lstmWo #> o_s' + tr lstmWi #> i_s' + tr lstmWc #> c_s'
190+
191+ -- Calculate the gradient Matricies for the input
192+ lstmWf' = f_s' `outer` input
193+ lstmWi' = i_s' `outer` input
194+ lstmWo' = o_s' `outer` input
195+ lstmWc' = c_s' `outer` input
196+
197+ -- Calculate the gradient Matricies for the cell
198+ lstmUf' = f_s' `outer` cell
199+ lstmUi' = i_s' `outer` cell
200+ lstmUo' = o_s' `outer` cell
201+
202+ -- The biases just get the values, but we'll write it so it's obvious
203+ lstmBf' = f_s'
204+ lstmBi' = i_s'
205+ lstmBo' = o_s'
206+ lstmBc' = c_s'
207+
208+ gradients = LSTMWeights lstmWf' lstmUf' lstmBf' lstmWi' lstmUi' lstmBi' lstmWo' lstmUo' lstmBo' lstmWc' lstmBc'
209+ in (gradients, S1D cell', S1D input')
210+
211+-- | Generate an LSTM layer with random Weights
212+-- one can also just call createRandom from UpdateLayer
213+--
214+-- Has forget gate biases set to 1 to encourage early learning.
215+--
216+-- https://github.com/karpathy/char-rnn/commit/0dfeaa454e687dd0278f036552ea1e48a0a408c9
217+--
218+randomLSTM :: forall m i o. (MonadRandom m, KnownNat i, KnownNat o)
219+ => m (LSTM i o)
220+randomLSTM = do
221+ let w = (\s -> uniformSample s (-1) 1 ) <$> getRandom
222+ u = (\s -> uniformSample s (-1) 1 ) <$> getRandom
223+ v = (\s -> randomVector s Uniform * 2 - 1) <$> getRandom
224+
225+ w0 = konst 0
226+ u0 = konst 0
227+ v0 = konst 0
228+
229+ LSTM <$> (LSTMWeights <$> w <*> u <*> pure (konst 1) <*> w <*> u <*> v <*> w <*> u <*> v <*> w <*> v)
230+ <*> pure (LSTMWeights w0 u0 v0 w0 u0 v0 w0 u0 v0 w0 v0)
231+
232+-- | Maths
233+--
234+-- TODO: move to not here
235+sigmoid :: Floating a => a -> a
236+sigmoid x = 1 / (1 + exp (-x))
237+
238+sigmoid' :: Floating a => a -> a
239+sigmoid' x = logix * (1 - logix)
240+ where
241+ logix = sigmoid x
242+
243+tanh' :: (Floating a) => a -> a
244+tanh' t = 1 - s ^ (2 :: Int) where s = tanh t