In 2015, Ben Lynn wrote a digit-classifying neural network in Haskell that is remarkable in two ways. Firstly, he wrote the backpropagation code by hand rather than relying on a framework. Secondly, he wrote in a terse, point-free-heavy style that must have taken great mental acuity. I came across Ben’s article on Hacker News where it was suggested that, unfortunately, the style does not make it easy for the unaccustomed reader to understand. Haskell makes refactoring safe and convenient so let’s try and improve the readability of the code!
The original code appears in Ben’s article and in a Github repository. Let’s warm up by tackling something small. zLayer implements a layer of a neural network.
zLayer ::[Float] -> ([Float], [[Float]]) -> [Float])
zLayer as (bs, wvs) =
zipWith (+) bs $ sum . zipWith (*) as <$> wvsHow does it work and what can we improve? Before we go any further I want to be confident that the changes that I am making are correct. Refactoring “is the process of restructuring existing code without changing its external behavior” so to be sure that I am not changing the external behaviour I will use a test harness implemented with the excellent Hedgehog property testing library.
I’m going to make an exact duplicate of the zLayer function called zLayer_new. Then I will proceed to refactor it whilst a Hedgehog check tests that the refactored version has exactly the same behaviour as the original version. The Hedgehog code is
-- Generate a random list
floatList :: MonadGen m => m [Float]
floatList = Gen.list (Range.linear 0 10)
(Gen.float (Range.linearFrac (-10) 10))
-- Generate random inputs for zLayer
zLayerInput :: MonadGen m => m ([Float], ([Float], [[Float]]))
zLayerInput = do
l1 <- floatList
l2 <- floatList
ll <- Gen.list (Range.linear 0 10) floatList
return (l1, (l2, ll))
-- Check that my new implementation matches the original
prop_same_zLayer :: Property
prop_same_zLayer = property $ do
(l1, t) <- forAll zLayerInput
zLayer_new l1 t === zLayer l1 t
-- Run all the tests
tests :: IO Bool
tests = checkSequential $$(discover)prop_same_zLayer checks that the old and new implementation of zLayer match. After every refactoring I’m going to run my Hedgehog tests to ensure that I haven’t changed any behaviour.
Now let’s get back to looking at zLayer. I notice that zipWith (+) performs the operation of adding two vectors. The vectors are represented by lists of Float, and zipWith (+) adds the corresponding entries in two lists, which is the same thing that vector addition does. This suggests my first refactoring. I’m going to define a new operator for vector addition.
(.+) ::[Float] -> [Float] -> [Float])
(.+) = zipWith (+)
zLayer_new :: [Float] -> ([Float], [[Float]]) -> [Float]
zLayer_new as (bs, wvs) = bs .+ (sum . zipWith (*) as <$> wvs)If zipWith (+) is vector addition then what is zipWith (*)? Componentwise multiplication followed by summing the components is the definition of the dot product of two vectors so I’ll introduce a dot operation.
dot :: [Float] -> [Float] -> Float
dot v1 v2 = sum (zipWith (*) v1 v2)
zLayer_new ::[Float] -> ([Float], [[Float]]) -> [Float]
zLayer_new as (bs, wvs) = bs .+ (dot as <$> wvs)What about this <$> thing? That’s another name for fmap, and it’s mapping the dot as function across the list-of-lists called wvs. What does that mean? One way of interpreting vector-matrix multiplication is that it maps the dot operation across the columns of a matrix. Therefore I add a “vector-matrix multiplication” operation, .*.
(.*) :: [Float] -> [[Float]] -> [Float]
v .* m = dot v <$> m
zLayer_new :: [Float] -> ([Float], [[Float]]) -> [Float]
zLayer_new as (bs, wvs) = bs .+ (as .* wvs)I’m going to take advantage of the knowledge gained during refactoring to introduce some type synonyms.
type Vector = [Float]
type Matrix = [[Float]]
zLayer_new :: Vector -> (Vector, Matrix) -> Vector
zLayer_new as (bs, wvs) = bs .+ (as .* wvs)Compared to the original zLayer (reproduced below) we have made a big improvement. There is no sum, no zipWith and no non-mathematical operator. zLayer is an operation which multiplies by a matrix and then adds a vector. It’s as simple as A * B + C.
zLayer as (bs, wvs) = zipWith (+) bs $ sum . zipWith (*) as <$> wvsI’m pleased with how zLayer looks so I will go straight to the next mathematical function in the module, called feed.
feed_new :: [Float] -> [([Float], [[Float]])] -> [Float]
feed = foldl' (((relu <$>) . ) . zLayer)The type is similar to that of zLayer but I’m not going to use my Vector and Matrix synonyms until I’m sure that they are indeed what the lists of Floats are representing. (I’ll keep adding Hedgehog equivalence tests but won’t show any more of the Hedgehog code in this article because it’s completely standard testing code.)
The use of <$> and . is mysterious. I know that <$> is fmap and . is function composition but that doesn’t help me understand what they are doing. Use of function composition is popular in Haskell to avoid naming intermediate variables. Sometimes avoiding intermediate variables helps readability but here I think it hurts readability so I’m going to reintroduce them.
I’m going to expand . by replacing it with its definition. I know that the first argument of zLayer is a vector so I choose the name v for the variable of the lambda. Naturally, after each refactoring I run my Hedgehog tests to check that feed_new has the same behaviour as feed.
feed_new = foldl' (\v -> ((relu <$>) . ) (zLayer v))That’s a good start. Then I apply the . section
feed_new = foldl' (\v -> (relu <$>) . zLayer v)and I can inline the definition of . again. The second argument of zLayer is a (Vector, Matrix) so I choose the name vm for the lambda.
feed_new = foldl' (\v vm -> relu <$> zLayer v vm)This is much better! Just from knowing what a left fold is we can understand that feed starts with a vector and then successively applies zLayer and relu for each vm in a list. At this point I know it’s fine to use the type synonyms I defined earlier. I’m also going to use my domain knowledge to define a new type synonym, Layer. A neural network consists of a sequence of Layers, each layer comprising a Vector and Matrix. We started with
feed = foldl' (((relu <$>) . ) . zLayer)and we ended with
type Layer = (Vector, Matrix)
feed_new :: Vector -> [Layer] -> Vector
feed_new = foldl' (\v vm -> relu <$> zLayer vm m)That is, feed takes a vector, applies a sequence of neural network layers and returns the resulting vector.
Right, onto the next candidate, revaz.
revaz_new :: [Float]
-> [([Float], [[Float]])]
-> ([[Float]], [[Float]])
revaz_new xs =
foldl' (\(avs@(av:_), zs) (bs, wms) ->
let zs' = zLayer av (bs, wms)
in ((relu <$> zs'):avs, zs':zs)) ([xs], [])Something looks suspicious here. There is a partial pattern match on a list. What happens if the list is empty? To our relief we can see that the list starts non-empty (it is [xs]) and it only ever grows (in each iteration it has relu <$> zs' consed onto the front). This is all very well but personally I’d prefer that it was the compiler, rather than my own reasoning, that reassures me about safe program behaviour. Let’s pass the “head of the list” in a separate tuple component. We’ll have to cons the final “head of the list” onto the list after the foldl' has completed.
revaz_new xs ys =
let (av, avs, zs) =
foldl' (\(av, avs, zs) (bs, wms) ->
let zs' = zLayer av (bs, wms)
in (relu <$> zs', av:avs, zs':zs))
(xs, [], [])
ys
in (av:avs, zs)This looks a bit more messy than before but let’s keep going and see where we can get. The next thing I notice is that each iteration around the loop we stick one new element onto the front of each of avs and zs, that is, at all stages avs and zs are exactly the same length. Let’s make this property structural by building one list of pairs instead of two lists. Again, we’ll have to postprocess the result of foldl', this time unzipping one list into two lists.
revaz_new xs ys =
let (av, avs_zs) =
foldl' (\(av, avs_zs) (bs, wms) ->
let zs' = zLayer av (bs, wms)
in (relu <$> zs', (av, zs'):avs_zs))
(xs, [])
ys
(avs, zs) = unzip avs_zs
in (av:avs, zs)Now that I can see what’s going on a bit more clearly I can add type synonyms to the signature. Importantly, I can see that the [[Float]]s are not Matrixs, rather they are lists of Vectors! It’s good that I was cautious and didn’t just substitute type synonyms blindly.
revaz_new :: Vector
-> [Layer]
-> ([Vector], [Vector])
revaz_new xs ys = (av:avs, zs)
where (avs, zs) = unzip avs_zs
(av, avs_zs) =
foldl' (\(av, avs_zs) (bs, wms) ->
let zs' = zLayer av (bs, wms)
in (relu <$> zs', (av, zs'):avs_zs))
(xs, [])
ysFinally I can neaten this a little by noticing that I don’t need to unpack the (bs, wms) tuple.
revaz_new :: Vector
-> [Layer]
-> ([Vector], [Vector])
revaz_new xs ys = (av:avs, zs)
where (avs, zs) = unzip avs_zs
(av, avs_zs) =
foldl' (\(av, avs_zs) t ->
let zs' = zLayer av t
in (relu <$> zs', (av, zs'):avs_zs))
(xs, [])
ysI find this marginally clearer than what we started with but not ideal. Let’s move on and we’ll come back to revaz later.
Next up, this behemoth.
deltas xv yv layers = let
(avs@(av:_), zv:zvs) = revaz xv layers
delta0 = zipWith (*) (zipWith dCost av yv) (relu' <$> zv)
in (reverse avs, f (transpose . snd <$> reverse layers) zvs [delta0]) where
f _ [] dvs = dvs
f (wm:wms) (zv:zvs) dvs@(dv:_) = f wms zvs $ (:dvs) $
zipWith (*) [sum $ zipWith (*) row dv | row <- wm] (relu' <$> zv)The Hedgehog equivalence test requires some care. As you can see from the pattern match on the result of revaz, deltas expects both the lists in the pair to be non-empty. The first one, avs, is guaranteed to be non-empty (we can conveniently see this from our refactored version of revaz: it has av on the front). What about the second one, zvs? By inspecting revaz we can see that the length of zvs is the same as the length of the list-of-Vectors input to revaz. Therefore we need layers to be non-empty. If we tell Hedgehog to generate input data satisfying this condition then all is well.
The first thing in deltas that I want to tackle is the local function f. It looks complicated, involving zipWith, a : section and a list comprehension. Even worse, it’s a recursive function! Even though recursion is typically thought of as the bread-and-butter of functional programming, unrestricted recursion is almost as bad for code comprehensibility as GOTO is in imperative programming. We ought to seek to use recursion combinators such as maps, folds, scans, etc. in preference to direct recursion.
Let’s tackle the recursion shortly. I see something I can do straight away, which is to use the dot product operator that I defined earlier.
f (wm:wms) (zv:zvs) dvs@(dv:_) = f wms zvs $ (:dvs) $
zipWith (*) [dot row dv | row <- wm] (relu' <$> zv)dot is commutative, so I can write dot row dv as dot dv row. Then I can conveniently rewrite the list comprehension as an fmap.
f (wm:wms) (zv:zvs) dvs@(dv:_) = f wms zvs $ (:dvs) $
zipWith (*) (fmap (dot dv) wm) (relu' <$> zv)Recalling that <$> is an alias for fmap we see that the list comprehension was just doing a vector-matrix product all along!
f (wm:wms) (zv:zvs) dvs@(dv:_) = f wms zvs $ (:dvs) $
zipWith (*) (dv .* wm) (relu' <$> zv)This is looking a lot better already, but still not good. There’s a mysterious : section there. Let’s just inline it.
f (wm:wms) (zv:zvs) dvs@(dv:_) = f wms zvs $
zipWith (*) (dv .* wm) (relu' <$> zv) : dvsThat looks less mysterious and the line is now readable. The function f is still mysterious though. What’s it doing? As in revaz, it unconditionally inspects the first element, dv, of a list, so let’s pass that in as a separate argument.
in (reverse avs,
f (transpose . snd <$> reverse layers) zvs delta0 []) where
f _ [] dv dvs = dv:dvs
f (wm:wms) (zv:zvs) dv dvs = f wms zvs
(zipWith (*) (dv .* wm) (relu' <$> zv)) (dv:dvs)Next I notice that f is iterating over two lists at the same time, taking one element off each during each iteration. It may as well iterate over the zip of the two lists instead! Why is that a good thing to do? Because f is now nearly a left fold. Let’s make it closer to left fold by packaging the dv and dvs together in a tuple
in (reverse avs,
f (zip (transpose . snd <$> reverse layers) zvs) (delta0, []))
where
f [] (dv, dvs) = dv:dvs
f ((wm, zv):wms_zvs) (dv, dvs) = f wms_zvs
(zipWith (*) (dv .* wm) (relu' <$> zv), dv:dvs)and combine dv and dvs after the recursive function (now called g) has returned, rather than in the base case.
f l t = let (dv, dvs) = g l t in dv:dvs
g [] (dv, dvs) = (dv, dvs)
g ((wm, zv):wms_zvs) (dv, dvs) = g wms_zvs
(zipWith (*) (dv .* wm) (relu' <$> zv), dv:dvs)Now g is a left fold!
g l t = foldl' h t l
where h (dv, dvs) (wm, zv) =
(zipWith (*) (dv .* wm) (relu' <$> zv), dv:dvs)Shuffling some arguments around and tidying gives us
f l t = let (dv, dvs) = g t l in dv:dvs
g = foldl' (\(dv, dvs) (wm, zv) ->
(zipWith (*) (dv .* wm) (relu' <$> zv), dv:dvs))This is already looking much better than the f that we started with, but additionally, something excellent has occurred. We can see that both revaz_new and g employ the same sort of recursion pattern. They both iterate over a list with a state and push one new value onto the front of a list each iteration (treating the list as a sort of stack). We can capture this recursion pattern by abstracting the behaviour out into a recursion combinator! I’m going to call it revMapWithState, a name that describes fairly well what it does.
revMapWithState :: (state -> a -> (state, stack))
-> (state, [stack])
-> [a]
-> (state, [stack])
revMapWithState f =
foldl' (\(state, stack) item ->
let (nextState, nextStack) = f state item
in (nextState, nextStack:stack))revMapWithState captures the sort of recursion that’s happening in revaz_new and g, and we can rewrite both of them in terms of it.
revaz_new :: Vector
-> [Layer]
-> ([Vector], [Vector])
revaz_new xs ys = (av:avs, zs)
where (avs, zs) = unzip avs_zs
(av, avs_zs) =
revMapWithState (\av t ->
let zs' = zLayer av t
in (relu <$> zs', (av, zs')))
xs
ys
...
g = revMapWithState (\dv (wm, zv) ->
(zipWith (*) (dv .* wm) (relu' <$> zv), dv))The benefit of using revMapWithState is that it takes care of the result stack for us and we only have to produce the next element of the stack and the next state.
Looking at where we are now with deltas_new, I can see that expressions of the form zipWith (*) _ (relu' <$> _) appear twice. Twice isn’t quite enough on its own to make us want to pull it out into a separate function, but my domain knowledge that this is the backpropagator for relu convinces me that this is the right thing to do. Now we can tidy up, using a little domain knowledge in naming the variables, to achieve
dRelu :: Vector -> Vector -> Vector
dRelu a b = zipWith (*) a (relu' <$> b)
deltas_new :: Vector
-> Vector
-> [Layer]
-> ([Vector], [Vector])
deltas_new xv yv layers = (reverse avs, dv:dvs) where
(avs@(av:_), zv:zvs) = revaz xv layers
dv0 = dRelu (zipWith dCost av yv) zv
matrices = map snd layers
(dv, dvs) =
backProp dv0 (zip (map transpose (reverse matrices)) zvs)
backProp =
revMapWithState (\dv (wm, zv) -> (dRelu (dv .* wm) zv, dv))It’s still not easy to read, but the reason that it is hard to read has changed. It is now hard to read because it’s a hand-written backpropagation routine for a deep neural network. It used to be hard to read because it contained deeply-nested expressions, non-mathematical symbols and an ad hoc recursive function. I think that the refactored version is about as easy to read as the overall design of this module will allow.
We’ve improved the readability of the original program significantly so I’ll stop here. Whilst I think we have improved Ben’s program I’m also impressed that he was able to write it the way he did in the first place! I couldn’t have kept all the necessary details in my head. There’s more we could refactor but this article already gives us plenty to chew on. If you enjoyed this article and you’d like me to continue refactoring in a future article then please and let me know.
Haskell makes refactoring a breeze. Higher-order functions and expression-based style make it possible to slice and splice code in fine grained ways and easily capture repeated patterns. Hedgehog allows us to refactor mercilessly whilst remaining confident that we are not changing program behaviour. The examples of refactorings we saw that make code easier to understand were: using domain-specific operators and types, introducing new variable names when it helps readability, converting program properties (which can’t be checked by a type checker) into program structure (which can) and avoiding explicit recursion where possible.
If you come across Haskell code in the wild that you can’t easily read then do not be disheartened. Maybe it can be refactored using the techniques described in this article.
See also my other worked examples of refactoring in Haskell: