Memoization in Neural networks

Memoization is an interesting computer science technique, mostly used in dynamic programming problems like Fibonacci or shortest-path algorithms. Using this technique the results of repetitive function calls or computations are calculated once and stored to avoid re-computation of same function calls. This eliminates redundancy and improves efficiency of the algorithm by improving time complexity.

Before we see how it works in neural networks, let’s understand what memoization actually does and what problem it solves using an example:

1. Without Memoization:

let’s understand and solve Fibonacci sequence problem:

In the Fibonacci sequence, each number is the sum of the previous two numbers.

F(n)=F(n−1)+F(n−2)

so given a number n , we have to calculate f(n) by calculating f(n-1) and f(n-2)

def fib(n):
    if n <= 1:
        return n
    return fib(n - 1) + fib(n - 2)

using the above code we’ll try to solve an example

for fib(5) , the below is the call stack

In the above example if we notice clearly, the functions are re-calculated repeatedly

• fib(3) repeats.
• fib(2) repeats many times.
• fib(1) ,fib(0) also repeat many times.

The function keeps doing the same work again and again.

This creates unnecessary computation, slower execution, exponential growth in function calls.

2. With Memoization:

Instead of recomputing the same function calls repeatedly, memoization stores previously computed results in memory and reuses them whenever needed(basically caching).

mem = {}
def fib(n):
    if n in mem:
        return mem[n]
    if n <= 1:
        return n
    mem[n] = fib(n - 1) + fib(n - 2)
    return mem[n]

In the above code , we first check whether the result is previously computed or not and calculate only if it isn’t cached/stored . Using this, every Fibonacci sequence is only calculated once. This drastically improves efficiency of the algorithm and reduces complexity. By using memoization we are not changing the logic or approach, but only how it handles repetitive computations .

Now that we understand memoization in the context of dynamic programming, let’s see how it improves neural networks.

1. The problem with chain rule:

Neural networks learn through a repeated cycle of four steps:

1. Input pass — the input data is fed into the network through the input layer
2. Forward pass — for each layer, the weighted sum W·X + b is computed and passed through an activation function
3. Loss computation — the network’s prediction is compared to the actual value using a loss function, such as Mean Squared Error: L = (yᵢ − ŷᵢ)²
4. Backpropagation — gradients of the loss are computed with respect to each weight using the chain rule, and weights are updated to minimize the loss

Backpropagation is an important optimization step and this is where the computation becomes a real problem.

Get Dineshbolledula’s stories in your inbox

Join Medium for free to get updates from this writer.

Remember me for faster sign in

To understand this more intuitively, we’ll take a look at an example neural network with 1 input layer , 2 hidden layers and an output layer using a dataset with 4 input features.

In our example network, the weight matrices are:

• W¹ — shape (3×4), connecting input layer to hidden layer 1
• W² — shape (2×3), connecting hidden layer 1 to hidden layer 2
• W³ — shape (1×2), connecting hidden layer 2 to output layer

During backpropagation, we update weights from right to left — starting at W³, then W², then W¹.Each weight requires its own gradient and we’ll see, many of these gradients share common sub-expressions that get recomputed from scratch every time.

1. W³ Gradient Computation:

W³ is a (1×2) matrix with two weights — W³₁₁ and W³₁₂. We need the gradient of loss L with respect to each:

we can clearly notice that ∂L/∂o₃₁is calculated twice .It is computed once for W³₁₁ and then recomputed again for W³₁₂. At this layer redundancy may not be a big problem . But as we move to previous layers ,the redundancy increases .

2. W² Gradient Computation:

W² is a (2×3) matrix with 6 weights — W²₁₁, W²₁₂, W²₁₃, W²₂₁, W²₂₂ , W²₂₃. We need the gradient for each:

If we observe the above gradients clearly , we could see the following:

• ∂L/∂o₃₁ is calculated 6 times.
• ∂o₃₁/∂o₂₁ is calculated 3 times.
• ∂o₃₁/∂o₂₂ is calculated 3 times.

We can observe the computation of gradients is compounding and also redundant.

3. W¹ Gradient Computation:

W¹ is a (3×4) matrix with 12 weights — W ¹₁₁, W ¹₁₂, W ¹₁₃, W ¹₁₄, W¹₂₁, W¹₂₂ , W¹₂₃ , W¹₂₄, W¹ ₃₁, W¹ ₃₂, W¹ ₃₃, W ¹₃₄.

Gradient calculation for the W¹ matrix is a bit tricky. Let’s first calculate the gradient for the weight W ¹₁₁ and see.

The weight W ¹₁₁ affects/impacts next outputs using 2 different paths i.e., one is through node f₂₁ and other through f₂₂. This is a special case for chain rule.

Using this, the gradient for W ¹₁₁ is :

now similarly a part of this chain will be recomputed for 12 different times.

In our small 3-layer network, the same intermediate gradients are recomputed 20 times in total. Now imagine a network with 100 layers and millions of weights. This redundancy doesn’t grow linearly — it explodes and increases the complexity of the model.

This is exactly the same problem we saw with Fibonacci. And the fix is exactly the same — cache the repeated computations.

2. Using Memoization to fix the issue:

Previously we saw how multiple gradients are recomputed repeatedly across weight updates. This is where memoization comes handy. Instead of computing gradients each time , compute each intermediate gradient once, store it and reuse it.

Instead of recomputing ∂L/∂o₃₁ for every weight in W³, W², and W¹ — we compute it once during the backward pass and cache it. Same for other gradients.

In the above example we can observe how each intermediate gradient is computed and cached. Instead of 20+ redundant computations, each delta/gradient is computed exactly once. The total number of gradient computations drops from O(n²) to O(n) — where n is the number of layers.

In our small 3-layer network, caching reduced redundant computations from 20+ to 3. In a 100-layer network with millions of weights, the difference between cached and uncached backpropagation is the difference between a model that trains in hours and one that never finishes. Frameworks like PyTorch and TensorFlow have this built in — every time you call loss.backward() , memoization is quietly doing the heavy lifting. The next time you train a model, you'll know exactly what's happening under the hood.This is why memoization isn’t just an optimization trick in neural networks — it’s what makes training deep networks computationally feasible in the first place. Without it, backpropagation at scale would be practically impossible.