Sequences: Recurrence and Its Limits

import numpy as np

rng = np.random.default_rng(0)

hidden_size, input_size = 4, 3
W_input = rng.normal(0, 0.5, (input_size, hidden_size))
W_hidden = rng.normal(0, 0.5, (hidden_size, hidden_size))
bias = np.zeros(hidden_size)

sequence = rng.normal(0, 1, (5, input_size))     # 5 time steps

hidden = np.zeros(hidden_size)
print('step  hidden state')
for step, x in enumerate(sequence, 1):
    hidden = np.tanh(x @ W_input + hidden @ W_hidden + bias)
    print(f'{step:4d}  {hidden.round(4)}')
print()
print(f'the same {W_hidden.size + W_input.size + bias.size} parameters were '
      f'used at all 5 steps,')
print('and would be used at all 5,000 steps of a longer sequence.')
print('That is why a recurrent network accepts any length: the')
print('length is in the loop, not in the weights.')

import numpy as np

rng = np.random.default_rng(0)
W_input = rng.normal(0, 0.5, (3, 4))
W_hidden = rng.normal(0, 0.5, (4, 4))

def final_state(sequence):
    hidden = np.zeros(4)
    for x in sequence:
        hidden = np.tanh(x @ W_input + hidden @ W_hidden)
    return hidden

sequence = rng.normal(0, 1, (5, 3))
reversed_sequence = sequence[::-1]

print('forwards ', final_state(sequence).round(4))
print('backwards', final_state(reversed_sequence).round(4))
print('same:', bool(np.allclose(final_state(sequence),
                                final_state(reversed_sequence))))
print()

# Compare with a model that averages its inputs - a "bag of words".
print('mean of the sequence, forwards ', sequence.mean(axis=0).round(4))
print('mean of the sequence, backwards', reversed_sequence.mean(axis=0).round(4))
print('same:', bool(np.allclose(sequence.mean(axis=0),
                                reversed_sequence.mean(axis=0))))
print()
print('The recurrent state depends on the order. The mean does')
print('not. Every bag-of-words model is blind to the difference')
print('between "the dog bit the man" and its reverse, and that')
print('blindness is exactly what recurrence removes.')

import numpy as np
import pandas as pd

rng = np.random.default_rng(1)

SIZE = 32


def gradient_after(steps, scale, size=SIZE, seed=0):
    generator = np.random.default_rng(seed)
    W = generator.normal(0, scale, (size, size))
    gradient = np.ones(size)
    for _ in range(steps):
        gradient = gradient @ W
        if not np.isfinite(gradient).all():
            return float('inf')
    return float(np.abs(gradient).mean())


# For a random matrix the per-step factor applied to a typical
# vector is about scale * sqrt(size). That is the number raised to
# the power of the sequence length.
rows = []
for scale in [0.10, 0.17, 0.18, 0.25]:
    row = {'init scale': scale,
           'per-step factor': round(scale * np.sqrt(SIZE), 3)}
    for steps in [10, 50, 100, 200]:
        row[f'{steps} steps'] = f'{gradient_after(steps, scale):.2e}'
    rows.append(row)
print(pd.DataFrame(rows).to_string(index=False))
print()
print('The per-step factor crosses 1 between the second and third')
print('rows. At 0.962 the gradient is 1e-45 by 200 steps. At 1.018')
print('- a difference of six percent - it is 1e+6. At 1.414 it is')
print('1e+35.')
print()
print('There is no setting that survives 200 steps, because the')
print('number being raised to the 200th power would have to be')
print('almost exactly 1. This is not the deep feed-forward')
print('problem, where each layer has its OWN matrix and the')
print('factors are independent and partly cancel. Here it is one')
print('matrix, 200 times, and the errors compound in one')
print('direction.')

import numpy as np

rng = np.random.default_rng(2)

hidden = 8
# Four gates: forget, input, output, and the candidate value.
W = {name: rng.normal(0, 0.3, (hidden + 4, hidden))
     for name in ['forget', 'input', 'output', 'candidate']}
b = {'forget': np.ones(hidden) * 1.0,      # biased OPEN at the start
     'input': np.zeros(hidden),
     'output': np.zeros(hidden),
     'candidate': np.zeros(hidden)}


def sigmoid(z):
    return 1 / (1 + np.exp(-z))


def lstm_step(x, h, c):
    joined = np.concatenate([x, h])
    f = sigmoid(joined @ W['forget'] + b['forget'])
    i = sigmoid(joined @ W['input'] + b['input'])
    o = sigmoid(joined @ W['output'] + b['output'])
    g = np.tanh(joined @ W['candidate'] + b['candidate'])
    c_next = f * c + i * g                  # ADDITION - the key line
    h_next = o * np.tanh(c_next)
    return h_next, c_next, f.mean()


sequence = rng.normal(0, 1, (12, 4))
h, c = np.zeros(hidden), np.zeros(hidden)
print('step  mean forget gate  |cell state|  |hidden state|')
for step, x in enumerate(sequence, 1):
    h, c, forget_mean = lstm_step(x, h, c)
    print(f'{step:4d}  {forget_mean:16.4f}  {np.abs(c).mean():12.4f}'
          f'  {np.abs(h).mean():14.4f}')
print()
print('The forget gate sits near 0.73, because its bias was')
print('initialised to 1.0 rather than 0 - the standard trick, and')
print('the reason every LSTM implementation has a forget_bias')
print('argument.')
print()
print('What matters is the PRODUCT of those gates, because that is')
print('the factor a gradient from the last step carries back to')
print('the first:')
for gate, label in [(0.5, 'bias 0 - the naive default'),
                    (0.73, 'bias 1 - as measured above'),
                    (0.95, 'bias 3'),
                    (0.99, 'bias 5')]:
    print(f'  gate {gate:<5} ({label:<26}) '
          f'over 12 steps {gate ** 12:8.4f}   '
          f'over 200 {gate ** 200:.2e}')
print()
print('A gate of 0.5 loses everything in a dozen steps. A gate of')
print('0.99 keeps 13% over two hundred. The arithmetic is')
print('unforgiving and it is why the forget bias is one of the few')
print('initialisation details that has a name.')

import torch
import torch.nn as nn

torch.set_num_threads(4)


def gradient_at_first_step(kind, length, forget_bias=None, hidden=32, seed=0):
    """How much of the final loss's gradient reaches input step 0."""
    torch.manual_seed(seed)
    layer = (nn.RNN if kind == 'rnn' else nn.LSTM)(1, hidden, batch_first=True)

    # PyTorch stacks the four gate biases as (input, forget, cell, output),
    # so the forget slice is the second block of `hidden` entries.
    if forget_bias is not None:
        with torch.no_grad():
            layer.bias_ih_l0[hidden:2 * hidden].fill_(forget_bias)
            layer.bias_hh_l0[hidden:2 * hidden].fill_(0.0)

    head = nn.Linear(hidden, 1)
    sequence = torch.randn(1, length, 1, requires_grad=True)
    output, _ = layer(sequence)
    head(output[:, -1]).sum().backward()
    gradients = sequence.grad[0, :, 0].abs()
    return float(gradients[0]), float(gradients[-1])


print('gradient reaching the FIRST input step:')
print()
print('length         RNN  LSTM, forget bias 0  bias 1  bias 3  bias 5')
for length in [5, 20, 50, 100, 200]:
    rnn_first, _ = gradient_at_first_step('rnn', length)
    columns = [gradient_at_first_step('lstm', length, bias)[0]
               for bias in [0.0, 1.0, 3.0, 5.0]]
    print(f'{length:6d}  {rnn_first:10.2e}  ' +
          '  '.join(f'{value:.2e}' for value in columns))
print()
last_rnn = gradient_at_first_step('rnn', 200)
print(f'for scale, the gradient at the LAST step of a 200-sequence '
      f'is {last_rnn[1]:.2e}')
print()
print('Read the first column down. By 200 steps the plain')
print('recurrence delivers a gradient of exactly zero to step 0,')
print('while the last step still gets 4.7e-02 - a ratio of')
print('infinity. The network cannot learn anything about its own')
print('early inputs, however long you train it.')
print()
print('Now read across. An LSTM at the default forget bias of 0 is')
print('no better: 3.2e-40 is zero for every practical purpose. The')
print('additive cell path does not help by existing; it helps when')
print('the forget gate is OPEN, and the bias is what opens it. At')
print('bias 5 the gradient at step 0 is 1.4e-01, larger than the')
print('plain RNN manages at its LAST step.')
print()
print('That is the honest version of "LSTMs solve the vanishing')
print('gradient". They provide a path that CAN carry a gradient.')
print('Whether it does is a matter of initialisation, and of what')
print('the gates learn.')

import numpy as np

rng = np.random.default_rng(3)
hidden, input_size, length = 6, 3, 7
sequence = rng.normal(0, 1, (length, input_size))

Wx = rng.normal(0, 0.4, (input_size, hidden))
Wh = rng.normal(0, 0.4, (hidden, hidden))


def run(sequence):
    h = np.zeros(hidden)
    states = []
    for x in sequence:
        h = np.tanh(x @ Wx + h @ Wh)
        states.append(h.copy())
    return np.array(states)


forward = run(sequence)
backward = run(sequence[::-1])[::-1]
bidirectional = np.concatenate([forward, backward], axis=1)

print('unidirectional state at each step :', forward.shape)
print('bidirectional state at each step  :', bidirectional.shape)
print()
print('step  forward sees   backward sees')
for step in range(length):
    print(f'{step:4d}  steps 0-{step:<7d} steps {step}-{length - 1}')
print()
print('A bidirectional layer doubles the state width and gives')
print('every step access to the whole sequence. It is correct for')
print('classification, translation and tagging - anything where')
print('the full input is available.')
print()
print('It is WRONG for forecasting and for generation, where the')
print('future is what you are predicting. Using a bidirectional')
print('layer there leaks the answer, and the symptom is a')
print('validation score that looks wonderful and a live model')
print('that is useless.')

import numpy as np

hidden = 512
print('encoder final state:', hidden, 'floats =', hidden * 4, 'bytes')
print()
for words in [5, 20, 60, 200]:
    print(f'{words:4d}-word sentence -> {hidden} floats '
          f'({hidden / words:.1f} floats per word)')
print()
print('A 200-word input gets 2.6 floats per word to carry its')
print('meaning, and the decoder sees nothing else. Translation')
print('quality duly collapsed with sentence length, and the fix')
print('for that - letting the decoder look back at every encoder')
print('state instead of just the last one - is attention. It was')
print('invented for this exact problem, and it turned out to make')
print('the recurrence itself unnecessary.')

import pandas as pd

rows = []
for length in [10, 100, 1000]:
    rows.append({'length': length,
                 'recurrent: steps that cannot be parallelised': length,
                 'attention: the same': 1,
                 'recurrent: hops from step 1 to step n': length,
                 'attention: the same ': 1})
print(pd.DataFrame(rows).to_string(index=False))
print()
print('A recurrent layer cannot compute step 5 before step 4. For')
print('a 1,000-token sequence that is 1,000 unavoidable sequential')
print('operations, and a GPU with thousands of cores spends almost')
print('all of them idle. Attention computes every position at')
print('once, as one large matrix multiplication, which is the')
print('single thing GPUs do best.')
print()
print('The second and fourth columns are the gradient argument:')
print('the distance information has to travel. 1,000 steps versus')
print('1. Nothing about LSTMs fixes that, it only makes the decay')
print('slower.')

import pandas as pd

rows = []
for length in [128, 1024, 8192, 65536]:
    attention_scores = length * length
    recurrent_work = length
    rows.append({'sequence length': length,
                 'attention score matrix': f'{attention_scores:,}',
                 'as float32': f'{attention_scores * 4 / 1024 ** 2:,.1f} MiB',
                 'recurrent state updates': f'{recurrent_work:,}'})
print(pd.DataFrame(rows).to_string(index=False))
print()
print('At 65,536 tokens the score matrix alone is 16 GiB per')
print('attention head per example. This is why long-context models')
print('do not use plain attention, and why state-space models -')
print('which are recurrences with a cleverer update - came back')
print('into fashion for exactly this regime.')

import numpy as np

rng = np.random.default_rng(4)
hidden = 4
Wx = rng.normal(0, 0.5, (1, hidden))
Wh = rng.normal(0, 0.5, (hidden, hidden))


def final_state(sequence, mask=None):
    h = np.zeros(hidden)
    for index, x in enumerate(sequence):
        candidate = np.tanh(x @ Wx + h @ Wh)
        if mask is None or mask[index]:
            h = candidate
    return h


real = np.array([[1.0], [-0.5], [2.0]])
padded = np.array([[1.0], [-0.5], [2.0], [0.0], [0.0], [0.0]])
mask = np.array([True, True, True, False, False, False])

print('unpadded          ', final_state(real).round(4))
print('padded, no mask   ', final_state(padded).round(4))
print('padded, with mask ', final_state(padded, mask).round(4))
print()
print('masked matches unpadded:',
      bool(np.allclose(final_state(real), final_state(padded, mask))))
print('unmasked matches unpadded:',
      bool(np.allclose(final_state(real), final_state(padded))))
print()
print('Three zero steps moved the state a long way, because')
print('tanh(0 @ Wx + h @ Wh) is not h. The unmasked version is')
print('answering a different question, and nothing in the training')
print('loop will tell you.')

import numpy as np
import torch
import torch.nn as nn

torch.manual_seed(0)

input_size, hidden_size, length = 3, 5, 4
cell = nn.RNN(input_size, hidden_size, batch_first=True, nonlinearity='tanh')

sequence = torch.randn(1, length, input_size)
output, final = cell(sequence)

# PyTorch stores the input-to-hidden weights TRANSPOSED relative to
# the (input, hidden) convention used by hand above.
W_ih = cell.weight_ih_l0.detach().numpy()      # (hidden, input)
W_hh = cell.weight_hh_l0.detach().numpy()      # (hidden, hidden)
b_ih = cell.bias_ih_l0.detach().numpy()
b_hh = cell.bias_hh_l0.detach().numpy()

h = np.zeros(hidden_size)
by_hand = []
for step in range(length):
    x = sequence[0, step].numpy()
    h = np.tanh(W_ih @ x + b_ih + W_hh @ h + b_hh)
    by_hand.append(h.copy())
by_hand = np.array(by_hand)

framework = output[0].detach().numpy()
print('framework output shape', tuple(output.shape),
      '= (batch, length, hidden)')
print('final state shape     ', tuple(final.shape),
      '= (layers, batch, hidden)')
print()
print('max difference', float(np.abs(framework - by_hand).max()))
print('agree:', bool(np.allclose(framework, by_hand, atol=1e-5)))
print()
print('Two conventions worth writing on a card. PyTorch keeps TWO')
print('biases, one for the input term and one for the hidden term,')
print('which is redundant mathematically and real in the')
print('parameter count. And batch_first=False is the DEFAULT, so')
print('the shape is (length, batch, hidden) unless you say')
print('otherwise - a transposition that trains perfectly happily')
print('and gives nonsense.')