Memoization via Representables
18 September 2022
What is the most basic container type a language can have? Some people may answer vectors, others would go with hash tables, but in this post I am arguing in favor of functions. Yes, functions. Even though they aren’t generally seem as a data structure per se, we will see that most containers are in fact a way to represent a function with a given storage layout. To illustrate this “functions are containers” idea, let’s take a look at an application that tightly couples both concepts: memoization.
By the way, in case this is the first time you hear about it: memoization is a programming technique where instead of letting a function calculate the same value whenever it is called, we instead store the already calculated value somewhere and just do a lookup instead of redoing the entire calculation.
What I like the most in this post’s code is that we’re going to delve deeply into the realm of abstraction to then emerge with a concept with clear and practical applications!
{-# LANGUAGE DeriveFunctor, TypeFamilies #-}
{-# LANGUAGE ScopedTypeVariables, RankNTypes #-}
{-# LANGUAGE TypeApplications, AllowAmbiguousTypes #-}
import Numeric.Natural
import Data.Kind (Type)
Sequences are functions that do not forget
An important fact that is normally briefly alluded in any mathematics book and immediately forgotten by (almost) every reader is that whenever you see a subindex such as x_n, what it in fact denotes is a function application x(n).1 Now consider the datatype of infinite streams as in the previous post:
data Stream a = a :> Stream a
deriving Functor
infixr 5 :>
This type models sequences of type a
, the
kind of thing a mathematician would denote as \{x_n\}_{n \in \mathbb{N}}. Oh
look, there’s a subindex there! Our previous discussion
tells us that we should be able to interpret a
Stream a
as a function
Natural -> a
. This is done by indexing: we
turn a stream xs
into the function that takes
n : Natural
to the nth element of
xs
. The definition is recursive, as one might
expect:
-- Access the nth value stored in a Stream
streamIndex :: Stream a -> (Natural -> a)
:> _) 0 = x
streamIndex (x :> xs) n = streamIndex xs (n-1) streamIndex (_
Conversely, given any f : Natural -> a
,
we can form a stream by applying f to each natural number to form
something like
[f 0, f 1, f 2, f 3,..., f n,...]
. Since
Stream
is a functor, we achieve this by mapping
f into the stream of
natural numbers:
-- Take f to [f 0, f 1, f 2, f 3, ...]
streamTabulate :: (Natural -> a) -> Stream a
= fmap f naturals where
streamTabulate f = 0 :> fmap (+1) naturals naturals
These functions are inverse to one another:
streamIndex . streamTabulate = id
streamTabulate . streamIndex = id
Meaning that we thus have a natural isomorphism
Stream a ≃ Natural -> a
. This is a strong
assertion and means that, mathematically speaking, Streams
and functions from the Naturals are essentially the same
thing. We are doing Haskell in here, however, not pure math.
And in a programming language meant to run in a real
computer, not only in the realm of ideas, we also must take
into account something more: how are those types laid out
into memory?
In the case of functions, they are compiled to chunks of
instructions that calculate some value. Specially, if you
have some rather expensive function
f : Natural -> a
and have to calculate
f n
in many parts of your program for the same
n
, all work will have to be redone each time to
get that sweet value of type a
.2 Streams, on the other
hand, are lazy infinite lists and, because of Haskell’s call-by-need
evaluation strategy, its components remain saved into
memory for reuse.3
This last paragraph is the heart of memoization in Haskell: one does not memoize functions, one turns functions into data and the language automatically memoizes the data. I don’t know about you, but I find this very cool indeed.
Memoizing with recursion
A large inspiration to this post comes from the great introduction to memoization in the Haskell wiki. Thus, we will follow their lead and explore the Fibonacci sequence as a recurring example throughout this post.
I must admit that I find illustrating a recursion concept through the Fibonacci numbers kind of a cliché… Nevertheless, clichés have their upside in that you, the reader, will have seen them so much that may even perhaps feel familiar with what we will be doing here. The Fibonacci numbers are also a well-known example where memoization can make a function go from exponential to polynomial complexity. Well, let’s start with their usual recursive definition:
fibRec :: Num a => Natural -> a
0 = 0
fibRec 1 = 1
fibRec = fibRec (n-1) + fibRec (n-2) fibRec n
Although elegant, this definition is extremely
slow. Running fibRec 100
on ghci already
took much longer than I was disposed to wait… The problem is
that the recursion has to calculate the same arguments a lot
of times, leading to an exponential complexity.
Since the problem is overlapping calculations, we can
accelerate this function using memoization. But in this
case, just turning it into a stream is not enough, because
the fibRec
will still use the slow definition
to build each of the sequence’s component. But fear nothing,
there is a salvation! It starts by writing the function in
operator form, instead of using recursion, just like we did
with the Bellman Equation in my previous post about
dynamic programming.
fibOp :: Num a => (Natural -> a) -> (Natural -> a)
0 = 0
fibOp v 1 = 1
fibOp v = v (n-1) + v (n-2) fibOp v n
You can think of fibOp
as one step of the
Fibonacci recursion, where v
is a function that
knows how to continue the process. Another way to look at
it, that is closer to the dynamic programming view, is that
if v
is an estimation of the Fibonacci values
then fibOp v
will be an improved estimation
given the available information. No matter what view you
choose, the important part to us is that the fixed point of
fibOP
is fibRec
.
= let x = f x in x
fix f
fibNaive :: Num a => Natural -> a
= fix fibOp -- same as fibRec fibNaive
Where we called it fibNaive
because it would
be rather naive to do all this refactoring in order to
arrive at the exact same thing…
Alright, with fix
we have all the necessary
building blocks to accelerate our calculations. It’s now
time to fit them together! Before fixing the operator, we
will turn it into something that “keeps a memory”. If we
compose our tabulation function with fibOp
, we
get a function turns a function v
into a
Stream, over which we can index to get back a function. In
this case, however, the same stream is shared for all
arguments. Thus, the fixed point indexes into this Stream
during the recursion process! Moreover, there is nothing
specific to the Fibonacci sequence in this process, so we
can abstract this procedure into a separate function.
streamMemoize :: ((Natural -> a) -> Natural -> a) -> Natural -> a
= fix (streamIndex . streamTabulate . f)
streamMemoize f
fibSmart :: Num a => Natural -> a
= streamMemoize fibOp fibSmart
Notice that by our previous discussion,
streamIndex . streamTabulate
equals
id
. Thus, by construction,
fibNaive
and fibMemo
are also
equal as functions. Nevertheless, their runtime behavior is
considerably different! As Orwell would put it: in terms of
execution, some equals are more equal than others.
A Call for Representation
Very well, What is a function k -> a
after all? The textbook definition says it is a rule that
for each element of type k
associates a unique
element of type a
. The previous examples have
shown us that there are data structures which, in the sense
above, behave a lot like functions. For example, we saw how
to convert between Streams and functions and even used it to
accelerate the calculation of recursive functions. In the
case of Streams, both streamIndex
and
streamTabulate
are natural transformations4, meaning that there is a
natural isomorphism between streams and functions with
Natural
as domain:
forall a. Stream a ≃ Natural -> a.
We call a functor isomorphic to a type of functions, a Representable Functor. Those have important applications in Category Theory because they are closely related to universal properties and elements. However, today we are interested in their more mundane applications, such as memoizing domains other than the Naturals.
In Haskell, we can codify being Representable as a typeclass. It must have an associated type saying to which function type the Functor is isomorphic, together with two natural transformations that witness the isomorphism.
class Functor f => Representable f where
type Key f :: Type
tabulate :: (Key f -> a) -> f a
index :: f a -> (Key f -> a)
-- Class laws:
-- index . tabulate = id
-- tabulate . index = id
As you can imagine, there is a Representable instance for Streams using what we have defined in the previous sections.
instance Representable Stream where
type Key Stream = Natural
index = streamIndex
= streamTabulate tabulate
Another interesting instance is for functions themselves! After all, the identity is, strictly speaking, a natural isomorphism.
instance Representable ((->) k) where
type Key ((->) k) = k
index = id
= id tabulate
With some type-level magic, we can write a generalized memoization procedure. It has a scarier type signure, since we’re striving for genericity, but the idea remains the same: precompose with tabulate and index before fixing. The function is essentially the same we wrote before for Streams but parameterized on our Representable Functor of choice.
-- | Memoize a recursive procedure using a Representable container of your choice.
memoize :: forall f a. Representable f => ((Key f -> a) -> (Key f -> a)) -> (Key f -> a)
= fix (index @f . tabulate . g) memoize g
We can recover our Stream-memoized Fibonacci by telling
memoize
that we choose Stream
as
the container:
fibSmart' :: Num a => Natural -> a
= memoize @Stream fibOp fibSmart'
The function above is the same as our “more manual”
fibSmart
from before. As a matter of fact, even
the naive recursion is case of these memoization schemes! By
using the Representable instance of functions, the methods
do nothing, and we get a memoization scheme that has no
storage. Well, this is equivalent to our naive approach from
before.
fibNaive' :: Num a => Natural -> a
= memoize @((->) Natural) fibOp fibNaive'
Speeding up the Fibonacci even more
With our Representable machinery all set, it would be a shame to end this post with just one example of memoization.5 So, let’s see how we can memoize the Fibonacci function using an infinite binary tree structure. This is a fun example to look at because the isomorphism is not as straightforward as with Streams and because it is much faster. We begin by defining our datatype.
data Tree a = Node a (Tree a) (Tree a)
deriving Functor
By now, you should already know how the memoization works at a high-level, so we can just define it as before.
fibTree :: Num a => Natural -> a
= memoize @Tree fibOp fibTree
Alright, how are these Tree
s representable?
In the case of Stream
s, the relation was clear:
we kept decreasing the index and advancing on the Stream
until the index was zero. And this amounted to recursing on
a unary representation of the Naturals: we advanced at a
successor and stopped at zero. The secret to translate this
idea to Tree
s is to look at a natural number as
written in binary. I personally find this relation easier to
explain with a drawing. So, while we index Streams with a
linear approach,
For Trees, we index using a breadth-first approach.
By the way, I don’t know if the figure above makes it
clear since we are starting with zero, but this arrangement
is equivalent to branching on the based on the number’s
evenness. We descend left on odd numbers and right on even
numbers. The crucial part of this representation is that we
are able to reach the n-th index in O(log n)
steps, instead of the n
steps required for the
Stream. Alright, it’s time to turn all this talking into a
proper instance!
Let’s begin with some helpers to make the evenness check more explicit.
data Eveness = Even | Odd
evenness :: Integral a => a -> Eveness
= if odd n then Odd else Even
evenness n
instance Representable Tree where
type Key Tree = Natural
We tabulate a function using the same ideas as for Streams: create a Tree of natural numbers and map the function over it. The tree is created by branching into even and odd numbers.
= fmap f nats where
tabulate f = Node 0
nats fmap (\ n -> 2*n + 1) nats)
(fmap (\ n -> 2*n + 2) nats) (
For indexing, we test for the number’s evenness and branch accordingly until we’re looking for the zeroth index.
index (Node a _ _) 0 = a
index (Node _ l r) n = case evenness n of
Odd -> index l (div n 2)
Even -> index r (div n 2 - 1)
Farewell
With this we finish our stroll through memoization-land. By the way, this post is a literate haskell file. Thus, I really recommend you to take the functions in it and try some benchmarks to see how much the memoization helps. From my own tests in ghci, the Tree Fibonacci is much faster than the other two. But compiling with optimizations, the Stream approach gets much faster, so I might need a better benchmark in there.
References
- Much of the Fibonacci example is adapted from the Haskell wiki page on memoization.
- Chapter 14 of Bartosz Milewski’s great book Category Theory for Programmers.
- The adjunctions package on Hackage.
I must comment, perhaps to the abhorrence of my Haskell readers, that I’m rather fond Fortran. One thing I like about it is how the language uses the same syntax to denote vector indexing and function application: x(n).↩︎
In fact, not all work must be necessarily redone. If we are calculating
f n
twice in the same context, the compiler may be smart enough to do some sharing and only calculate it once. Of course, this works in Haskell thanks to the magic of referential transparency.↩︎Well, at least until the next gc pass if it is not in the top-level.↩︎
In Haskell, any polymorphic function
h :: forall a. F a -> G a
, whereF
,G
are functors, is a natural transformation.↩︎Well, two if you count the naive approach.↩︎