Конвергенция в мелкой нейронной сети

Question

У меня есть один скрытый слой с 5 единицами, входной слой с 10 единицами и одна скалярная единица вывода.Я использую активацию ReLu и на выходном слое нет нелинейности, только взвешенная сумма.Вместо того чтобы использовать существующий код из сети, я подумал, что выведу уравнения.Конвергенция действительно сбивает с толку, и я уверен, что это неправильно.

import numpy as np
import matplotlib.pyplot as plt
import math

d = 10
m = 5
alp = 1e-2
W1 = np.random.randn(m,d)
W2 = np.random.randn(1,m)
a0 = np.random.randn(d,1)
b1 = np.random.randn(m,1)
b2 = np.random.randn(1,1)
y = np.random.randn(1,1)

def compute_loss(y,a2):
       return np.sum(np.power(y-a2,2))/2

def gradient_step(W1,W2,b1,b2,a1,a2,z1):
    W2 += alp*(y-a2)*a1.transpose()
    b2 += (y-a2)
    a1_deriv = np.array(reluDerivative(z1))
    b1 += (y-a2)*(np.matmul(W2,np.diagflat(a1_deriv))).transpose()
    W1 += (y-a2)*(a0.dot(W2).dot(np.diagflat(a1_deriv))).transpose()
    return W1,W2,b1,b2,a1,a2,z1

def reluDerivative(x):
     x[x<=0] = 0
     x[x>0] = 1
     return x

loss_vec = []
num_iterations = 50

for i in range(num_iterations):
    z1 = np.matmul(W1,a0)+b1
    a1 = np.maximum(0,z1)
    a2 = np.matmul(W2,a1)+b2
    loss_vec.append(compute_loss(y,a2))
    W1,W2,b1,b2,a1,a2,z1 = gradient_step(W1,W2,b1,b2,a1,a2,z1)
plt.plot(loss_vec)