Difference equation in LSTM network on Tensorflow

Question

I'd like to use a LSTM network on Tensorflow to implement a difference equation. I searched on internet but I didn't find anything about this topic.

The equation is:

formula

in which b=[1, 2, 1] and a=[1, -1.6641, 0.8387].

My aim is to use a neural network to find the correlation between input and output. Due to that to find the output ad k-instant you have to know also the previous inputs and outputs, my idea is to implement a LSTM network (many to one structure).

If we suppose to have an input vector of 500 samples and to use a window size of 5, the input of LSTM network is a vector of shape (500,5,1) while the output is (500,1,1).

The IN%OUT of first iteration are:

[0; x(k-4), x(k-3), x(k-2), x(k-1), x(k); 1] -> [1; y(k); 1] formula

in the second iteration:

[0; x(k-3), x(k-2), x(k-1), x(k), x(k+1); 1] -> [1; y(k+1); 1] formula

So I used a LSMT network with stateful set to TRUE to allow the network to remember past states but it doesn't converge.

It seems to me that the idea is correct but I cannot see where I am going wrong. Could someone help me find the problem? I copy and paste the code below and the network is developed on Tensorflow.

# Difference equation
K = 0.0436
b = np.array([1,2,1])
a = np.array([1, -1.6641, 0.8387])
x = np.random.uniform(0, 1, 100)
y = K*(signal.lfilter(b,a,x))

# Generate Dataset
X_train     = np.random.uniform(0, 1, 100)
y_train     = K*(signal.lfilter(b,a,X_train))
X_val       = np.ones(100)
y_val       = K*(signal.lfilter(b,a,X_val))
X_test      = np.random.uniform(0.5, 0.8, 100)
y_test      = K*(signal.lfilter(b,a,X_test))

def get_x_split(data, windows_size):
    """ Return sliding window dataset. """
    x_temp = np.zeros([1,windows_size-1])
    x = np.array([])
    for i in range(0,len(data)):
        x_temp = np.append(x_temp[-windows_size+1:], data[i]).T
        x = np.append(x, x_temp, axis=0)
    x = np.reshape(x, (int(len(x)/windows_size), windows_size))
    return x

windows_size = 10
X_train     = get_x_split(X_train, windows_size)
X_val       = get_x_split(X_val, windows_size)
X_test      = get_x_split(X_test, windows_size)

X_train     = np.reshape(X_train, (X_train.shape[0], X_train.shape[1], 1))
X_val       = np.reshape(X_val, (X_val.shape[0], X_val.shape[1], 1))
X_test      = np.reshape(X_test, (X_test.shape[0], X_test.shape[1], 1))

# Model Definition
activation_function = 'tanh'
def build_model():
    input_layer = Input(shape=(X_train.shape[1],1), batch_size=1)
    HL_1 = LSTM(1, activation=activation_function, return_sequences=True, stateful = True)(input_layer)
    HL_2 = LSTM(1, activation=activation_function, return_sequences=False, stateful = True)(HL_1)
    output_layer = Dense(1, activation='relu',name='Output')(HL_2)
    model = Model(inputs=input_layer, outputs=output_layer)
    return model

model = build_model()
model.compile(optimizer=RMSprop(),
              loss={'Output': 'mse'},     #mse
              metrics={'Output': tf.keras.metrics.RootMeanSquaredError()})

# Training
history = model.fit(x=X_train,
          y=y_train,
          batch_size=1,
          validation_data=(X_val, y_val),
          epochs=5000,
          verbose=1,
          shuffle=False)

# Test
y_pred = model.predict(X_test)

pred_samples = 400   
plt.figure(dpi=1200)
plt.plot(y_test[300:pred_samples,3,0], label='true', linewidth=0.8, alpha=0.5)
plt.plot(y_pred[300:pred_samples,3,0], label='pred')
plt.legend()
plt.grid()
plt.title("Test")
plt.show()