RNN：順序資訊怎麼進入模型？

a_t = W_ih x_t + b_ih + W_hh h_(t−1) + b_hh
h_t = tanh(a_t)

W_ih = [[ 0.4, −0.2],     W_hh = [[ 0.5, 0.15],
        [ 0.1,  0.3]]             [−0.25, 0.4]]

b_ih = [0.05, −0.10]
b_hh = [−0.02, 0.03]
h_0  = [0.10, −0.20]

input  : [batch=1, time=3, input_dim=2]
output : [batch=1, time=3, hidden_dim=2]
h_0/h_n: [layers×directions=1, batch=1, hidden_dim=2]

L = 0.5 × Σ_t Σ_k (h_t[k] − target_t[k])²

δ_t = (h_t − target_t + W_hhᵀ δ_(t+1)) ⊙ (1 − h_t²)
∂L/∂W_hh = Σ_t δ_t h_(t−1)ᵀ

python3 -m venv .venv-rnn
source .venv-rnn/bin/activate
python -m pip install torch==2.3.1 --index-url https://download.pytorch.org/whl/cpu
python rnn_mechanics.py

"""CPU/float64 RNN mechanics, with no corpus, fitting or benchmark.

PyTorch stores W_ih/W_hh as (hidden,input)/(hidden,hidden):
h_t = tanh(W_ih x_t + b_ih + W_hh h_prev + b_hh), in column notation.
The plain-Python calculation below uses this orientation explicitly.
"""
import json
import math
import platform
import torch


def plain_recurrence(xs, h0, wi, wh, bi, bh):
    """Independent Python float/math.tanh forward, not a torch rewrite."""
    h = list(h0)
    history, preactivations = [], []
    for x in xs:
        a = [sum(wi[i][j] * x[j] for j in range(2)) + bi[i]
             + sum(wh[i][j] * h[j] for j in range(2)) + bh[i]
             for i in range(2)]
        h = [math.tanh(v) for v in a]
        preactivations.append(a)
        history.append(h)
    return history, preactivations


def main():
    torch.manual_seed(42)  # Initialization is overwritten; not a sampled experiment.
    dtype = torch.float64
    device = torch.device('cpu')
    wi = [[0.4, -0.2], [0.1, 0.3]]
    wh = [[0.5, 0.15], [-0.25, 0.4]]
    bi, bh = [0.05, -0.1], [-0.02, 0.03]
    xs = [[1.0, 0.0], [0.0, 1.0], [1.0, 1.0]]
    initial = [0.1, -0.2]  # Nonzero on purpose: W_hh is active at all three steps.
    targets = [[0.2, -0.1], [-0.2, 0.3], [0.1, 0.2]]
    rnn = torch.nn.RNN(2, 2, num_layers=1, nonlinearity='tanh',
                      bias=True, batch_first=True, dropout=0.0,
                      bidirectional=False, device=device, dtype=dtype)
    with torch.no_grad():
        for name, value in [('weight_ih_l0', wi), ('weight_hh_l0', wh),
                            ('bias_ih_l0', bi), ('bias_hh_l0', bh)]:
            getattr(rnn, name).copy_(torch.tensor(value, dtype=dtype))
    x = torch.tensor([xs], dtype=dtype, device=device)  # batch1,time3,input2
    h0 = torch.tensor([[initial]], dtype=dtype, device=device)  # layer1,batch1,hidden2
    target = torch.tensor([targets], dtype=dtype, device=device)
    python_h, preact = plain_recurrence(xs, initial, wi, wh, bi, bh)
    output, hn = rnn(x, h0)
    manual = torch.tensor([python_h], dtype=dtype)
    torch.testing.assert_close(output, manual, rtol=0, atol=1e-14)
    torch.testing.assert_close(hn[0], output[:, -1], rtol=0, atol=1e-14)

    # Same multiset of tokens, different order, unchanged weights and initial state.
    order = [2, 0, 1]
    with torch.no_grad():
        reordered, reordered_hn = rnn(x[:, order], h0)
        order_difference = (reordered_hn - hn).abs().max().item()
        assert order_difference > 1e-3
        first, state = rnn(x[:, :2], h0)
        second, carried_hn = rnn(x[:, 2:], state)
        carried = torch.cat([first, second], dim=1)
        torch.testing.assert_close(carried, output, rtol=0, atol=1e-14)
        torch.testing.assert_close(carried_hn, hn, rtol=0, atol=1e-14)
        reset_second, reset_hn = rnn(x[:, 2:], torch.zeros_like(state))
        reset_difference = (reset_hn - hn).abs().max().item()
        assert reset_difference > 1e-3

    # Artificial hidden-state squared-error objective, not next-token CE or accuracy.
    # L = 0.5 * sum over all times and both hidden components (h_t - target_t)^2.
    step_losses = 0.5 * (output - target).square().sum(dim=2).squeeze(0)
    per_loss_grad = [torch.autograd.grad(v, rnn.weight_hh_l0, retain_graph=True)[0]
                     for v in step_losses]
    loss = step_losses.sum()
    loss.backward()
    autograd = rnn.weight_hh_l0.grad.detach().clone()
    torch.testing.assert_close(sum(per_loss_grad), autograd, rtol=0, atol=1e-14)

    # Analytic BPTT, pure Python: delta_t = (h_t-y_t + W_hh^T delta_next)
    # * (1-h_t^2). Each use of shared W_hh contributes outer(delta_t,h_prev).
    next_delta = [0.0, 0.0]
    deltas, contributions = [None] * 3, [None] * 3
    for t in range(2, -1, -1):
        delta = [(python_h[t][i] - targets[t][i]
                  + sum(wh[k][i] * next_delta[k] for k in range(2)))
                 * (1 - python_h[t][i] ** 2) for i in range(2)]
        previous = initial if t == 0 else python_h[t-1]
        contributions[t] = [[delta[i] * previous[j] for j in range(2)] for i in range(2)]
        deltas[t], next_delta = delta, delta
    analytic = torch.tensor([[sum(c[i][j] for c in contributions)
                              for j in range(2)] for i in range(2)], dtype=dtype)
    torch.testing.assert_close(analytic, autograd, rtol=0, atol=1e-14)
    assert all(abs(c[0][0]) > 1e-8 for c in contributions)

    # Central finite difference uses the actual torch.nn.RNN forward, not autograd.
    eps = 1e-6
    with torch.no_grad():
        original = rnn.weight_hh_l0[0, 0].item()
        values = []
        for sign in (1, -1):
            rnn.weight_hh_l0[0, 0] = original + sign * eps
            perturbed, _ = rnn(x, h0)
            values.append((0.5 * (perturbed - target).square().sum()).item())
        rnn.weight_hh_l0[0, 0] = original  # Restore the exact original parameter.
    finite_difference = (values[0] - values[1]) / (2 * eps)
    finite_error = abs(finite_difference - autograd[0, 0].item())
    assert finite_error < 1e-8
    assert rnn.weight_hh_l0[0, 0].item() == wh[0][0]

    result = {
        'scope': 'Hand-set CPU mechanics only; no corpus, optimizer, trained task, accuracy or benchmark.',
        'environment': {'python': platform.python_version(), 'torch': torch.__version__,
                        'device': str(output.device), 'dtype': str(output.dtype)},
        'parameters_pytorch_orientation': {'W_ih': wi, 'W_hh': wh, 'b_ih': bi, 'b_hh': bh},
        'inputs': {'token_names': ['A', 'B', 'C'], 'values': xs, 'initial_hidden': initial},
        'shapes': {'input_batch_first': list(x.shape), 'initial_hidden': list(h0.shape),
                   'output_all_hidden': list(output.shape), 'final_hidden': list(hn.shape)},
        'forward': {'preactivations_python': preact, 'hidden_python': python_h,
                    'hidden_torch': output.detach().squeeze(0).tolist(),
                    'max_abs_error': (output - manual).abs().max().item()},
        'order': {'permutation': order, 'token_names': ['C', 'A', 'B'],
                  'reordered_final_hidden': reordered_hn.squeeze().tolist(),
                  'original_final_hidden': hn.detach().squeeze().tolist(),
                  'max_abs_final_difference': order_difference,
                  'interpretation': 'This parameter/input example is order-sensitive; not proof all RNN settings distinguish all permutations.'},
        'chunks': {'split_lengths': [2, 1], 'carry_max_abs_error': (carried-output).abs().max().item(),
                   'reset_to_zero_final_hidden': reset_hn.squeeze().tolist(),
                   'reset_max_abs_final_difference': reset_difference,
                   'interpretation': 'Forward state carry matches full sequence. This no_grad example does not test truncated-BPTT gradients.'},
        'gradient': {'objective': '0.5*sum_t,sum_hidden (h_t-target_t)^2, no averaging',
                     'artificial_targets': targets, 'step_losses': step_losses.detach().tolist(),
                     'loss': loss.item(), 'autograd_shared_W_hh': autograd.tolist(),
                     'analytic_bptt_shared_W_hh': analytic.tolist(),
                     'analytic_deltas_time_order': deltas,
                     'parameter_reuse_contributions_time_order': contributions,
                     'per_time_loss_total_gradients': [g.tolist() for g in per_loss_grad],
                     'partition_warning': 'Reuse contributions include downstream losses; per-time-loss gradients include earlier parameter uses. These are distinct decompositions of the same total.',
                     'finite_difference_W_hh_0_0': finite_difference, 'finite_difference_epsilon': eps,
                     'finite_difference_abs_error': finite_error,
                     'all_three_reuse_contributions_0_0_nonzero': True},
        'checks': 'PASS: Python/nn.RNN forward, all output/h_n shapes, order difference, carry/reset, per-loss sum, analytic BPTT, finite difference.'}
    print(json.dumps(result, ensure_ascii=False, indent=2, allow_nan=False))


if __name__ == '__main__':
    main()