Memoization

This post was designed as a Jupyter Notebook. You can access the notebook here. I used the Deno kernel for JavaScript.

For the small problems of an introductory coding course, we generally don't worry too much about the efficiency of our code because it executes almost instantaneously (barring mistakes like infinite while loops). However, this isn't guaranteed, and as the size and complexity of our problems increases, we need increasingly clever solutions to keep our runtimes down. One technique that can cut down on repeated computation is memoization - saving previous results to look up later.

Let's take the Fibonacci sequence as an example. (I had originally planned to use the Collatz conjecture as my example but it didn't slow down enough to illustrate my point. It turns out the Fibonacci sequence seems to be a popular example for this use case - it was actually the example given by the Underscore.js documentation. ) The Fibonacci sequence is familiar as the sequence we get when, starting with 0 and 1, we repeatedly add the last two terms together. However, the precise mathematical definition is actually recursive:

$$F(n) = 0 \text{ for } n=0$$

$$F(n) = 1 \text{ for } n=1$$

$$F(n) = F(n-1) + F(n-2) \text{ otherwise}$$

Let's implement that in Javascript:

function fibonacci(n) {
    if (n === 0) {
        return 0;
    } else if (n === 1) {
        return 1;
    } else {
        return fibonacci(n-1) + fibonacci(n-2);
    }
}

Looks good! However, if we start calling our function for bigger and bigger numbers, we'll find that the function takes a noticeably long time to return. On my machine I start getting worried around F(40) - F(45).

console.log(fibonacci(45));

What's going on here?

Because the Fibonacci sequence is defined recursively, every time we call it, it calls itself twice. But each of those calls invoke the function twice again, and so on...

$$F(n) = F(n-1) + F(n-2)$$

$$= (F(n-2) + F(n-3)) + (F(n-3) + F(n-4))$$

$$= ((F(n-3) + F(n-4)) + (F(n-4) + F(n-5))) + ((F(n-4) + F(n-5)) + (F(n-5) + F(n-6))$$

$$= ...$$

So the number of calculations needed to compute the $n^{th}$ term of the Fibonacci sequence is about $2^n$ , which quickly gets out of hand, even for relatively small values of n. No longer can we assume our code will execute immediately.

Of course, in this case we could sidestep the issue entirely by refactoring the Fibonacci function to use a for loop instead of recursion, but not every recursive function can be rewritten, and some expensive functions aren't recursive at all. Instead, notice that our function is calling itself on the same inputs multiple times, sometimes even several times in a row. Every call to F(n) will always give the same output, so once we know that output we don't need to recompute it. Let's see if we can store the outputs somehow so we reuse them later.

let inputs = [];
let outputs = [];

function fibExternalMemo(n) {
    if (inputs.includes(n)) {
        let index = inputs.indexOf(n);
        return outputs[index];
    }

    if (n === 0) return 0;
    if (n === 1) return 1;

    let result = fibExternalMemo(n - 1) + fibExternalMemo(n - 2);
    inputs.push(n);
    outputs.push(result)
    return result;
}

console.log(fibExternalMemo(45));

Our code runs much faster now!

Every time we call our function, it checks to see if it's already seen the current input before. If so, it simply looks up the answer and returns. Only for new inputs where it doesn't already know the answer does it calculate from scratch, and the recursive calls also use the memo functionality. This is the basic premise of memoization - don't recalculate things we don't have to.

This works because the Fibonacci sequence is deterministic - the same input will always give the same output, so looking up prior inputs is a reliable way to determine the output. If our function depended on other factors, such as global variables or object states, we couldn't know for certain that the same input would always yield the same output.

Unfortunately, our reliance on external arrays breaks our function's independence, plus is a bit unwieldy. In our Underbar assignment, we created a wrapper function that uses closures to contain our inputs and outputs while keeping them out of way. We simply pass our function into the higher-order memoize function and get back a new function that we can use like the old one, but that handles all the input checking for us.

//assuming for simplicity that the input function takes only one input of a simple data type
//This avoids having to deeply compare the inputs in our example
function memoize(f) {
    let inputs = [];
    let outputs = [];

    function memoized(input) {
        for (let i = 0; i < inputs.length; i++) {
            if (inputs[i] === input) {
                return outputs[i];
            } 
        }

        let newResult = f(input);
        inputs.push(input);
        outputs.push(newResult);
        return newResult;
    }

    return memoized;
}

However if we try to feed our original fibonacci() function directly into memoize(), we'll find out that while it will almost instantly return any number in has seen before (as expected), it won't calculate new values any faster.

let wrappedFib = memoize(fibonacci);

console.log(wrappedFib(45)); //takes forever
console.log(wrappedFib(45)); //executes immediately now that 45 has been seen

This is because once it realizes it hasn't seen our input before, it calls the inner function, which recursively calls itself, never checking for seen values again. The inner function doesn't get the benefit of memoization and ends up doing everything the hard way.

We can do better than that!

I was actually quite stumped by this. Of course I could hand-craft a memoized Fibonacci function like so...

function makeMemoFib() {
    let inputs = [];
    let outputs = [];

    let fibonacci = function(n) {
        if (inputs.includes(n)) {
            let index = inputs.indexOf(n);
            return outputs[index];
        }

        if (n === 0) return 0;
        if (n === 1) return 1;

        let result = fibonacci(n - 1) + fibonacci(n - 2);
        inputs.push(n);
        outputs.push(result)
        return result;
    }

    return fibonacci;
}

let manuallyMemoized = makeMemoFib();

console.log(manuallyMemoized(45));

This accomplishes the goal of isolating the lookup arrays, but this function works exclusively for this case. If I wanted to memoize different recursive function, I'd have to write it from scratch.

Figuring I wasn't person to have this problem, I consulted the Underscore.js documentation to see if they had any tips. As it happened, the very example they provided for the _.memoize() function was the Fibonacci sequence (I mentioned it's popular), and the example code snippet gave me the (obvious in retrospect) answer:

Memoized recursive functions must be defined recursively.

let memoizedFib = memoize(n => {
    if (n === 0) {
        return 0;
    } else if (n === 1) {
        return 1;
    } else {
        return memoizedFib(n-1) + memoizedFib(n-2);
    }
});

console.log(memoizedFib(45));
console.log(memoizedFib(100));
console.log(memoizedFib(1000));

And there we have it! It takes a bit of dexterity with JavaScript's function syntax, but is much more flexible than having to write a function maker for every tricky function we want to memoize.

Of course memoization isn't a panacea. What it saves in time it costs in memory, since it has to store every input/output pair somewhere. Plus it's only useful if at least some inputs really will be duplicated - there's no point in storing every output if every input only gets used once. And as I discussed before, memoization doesn't help at all for functions that depend on exterior states.

However, for the cases where it does work, memoization really shines! Last Advent of Code I used memoization (of the global variable variety) for two problems (that I remember), and I expect knowing how to use higher-order memoizers will come in handy this year. And of course it's always a good idea to have more optimization techniques in our toolbox.