[Python] Layer-wise Relevance Propagation (LRP) not conserving total relevance in my BiLSTM

I'm implementing Layer-wise Relevance Propagation manually for a BiLSTM model, following Arras et al. implementation.

def lrp_bilstm(R_seq_3d, manual_bilstm_results, bilstm_weights, epsilon=1e-6):

# R_seq_3d.shape= (1, 100, 256)

forward_states = manual_bilstm_results["forward"]
backward_states = manual_bilstm_results["backward"]

batch_size, seq_len, total_hidden = R_seq_3d.shape
hidden_dim = total_hidden // 2

R_forward = R_seq_3d[:, :, :hidden_dim]
R_backward = R_seq_3d[:, :, hidden_dim:]

print(f"Total input relevance sum before lrp in forward branch(between merge and bilstm): {np.sum(R_forward):.6f}")
print(f"Total input relevance sum before lrp in backward branch (between merge and bilstm): {np.sum(R_backward):.6f}")

R_input_forward, R_reste_forward = lrp_lstm(R_forward, forward_states, bilstm_weights["forward"], epsilon )

R_input_backward_temp, R_reste_backward = lrp_lstm(R_backward, backward_states, bilstm_weights["backward"], epsilon )

R_input_backward = R_input_backward_temp[:, ::-1, :]

print(f"Résidu du lrp bilstm forward: {R_reste_forward:.6f}")
print(f"Résidu du lrp bilstm backward: {R_reste_backward:.6f}")

# concatener forward and backward relevances
R_input_total = np.concatenate([R_input_forward, R_input_backward], axis=-1)

print(f"Forward input relevance sum after lrp bilstm: {np.sum(R_input_forward):.6f}")
print(f"Backward input relevance sum after lrp bilstm: {np.sum(R_input_backward):.6f}")
print(f"Total input relevance sum after lrp bilstm: {np.sum(R_input_total):.6f}")

print(R_input_total.shape)

return R_input_total


def lrp_lstm(R_h_input, lstm_states, weights, epsilon=1e-7):
batch_size, seq_len, hidden_dim = R_h_input.shape
input_dim = lstm_states["x_t"].shape[2]

W = weights["W"]
U = weights["U"]
b = weights["b"]

# Attention: ordre Keras = [i, f, g, o]
W_i, W_f, W_g, W_o = np.split(W, 4, axis=1)
U_i, U_f, U_g, U_o = np.split(U, 4, axis=1)
b_i, b_f, b_g, b_o = np.split(b, 4)

R_input = np.zeros_like(lstm_states["x_t"])
R_c = np.zeros((batch_size, seq_len + 1, hidden_dim))
R_h = np.zeros((batch_size, seq_len + 1, hidden_dim))
R_g = np.zeros((batch_size, seq_len, hidden_dim))

R_h[:, seq_len - 1, :] = R_h_input[:, seq_len - 1, :]

for t in reversed(range(seq_len)):
x_t = lstm_states["x_t"][:, t, :]
h_prev = lstm_states["h"][:, t - 1, :] if t > 0 else np.zeros_like(lstm_states["h"][:, 0, :])
c_prev = lstm_states["c"][:, t - 1, :] if t > 0 else np.zeros_like(lstm_states["c"][:, 0, :])

i_t = lstm_states["i"][:, t, :]
f_t = lstm_states["f"][:, t, :]
g_t = lstm_states["g"][:, t, :]
o_t = lstm_states["o"][:, t, :]
c_t = lstm_states["c"][:, t, :]

R_c[:, t, :] += R_h[:, t, :]

Rc_from_f = lrp_linear(
f_t * c_prev,
np.identity(hidden_dim),
np.zeros(hidden_dim),
c_t,
R_c[:, t, :],
bias_nb_units=2 * hidden_dim,
eps=epsilon,
bias_factor=0.0,
debug= False
).T

Rc_from_g = lrp_linear(
i_t * g_t,
np.identity(hidden_dim),
np.zeros(hidden_dim),
c_t,
R_c[:, t, :],
bias_nb_units=2 * hidden_dim,
eps=epsilon,
bias_factor=0.0,
debug= False
).T

R_c[:, t - 1, :] += Rc_from_f
R_g[:, t, :] = Rc_from_g

g_pre = np.dot(x_t, W_g) + np.dot(h_prev, U_g) + b_g

Rx = lrp_linear(
x_t,
W_g,
b_g,
g_pre,
R_g[:, t, :],
bias_nb_units=input_dim + hidden_dim,
eps=epsilon,
bias_factor=0.0,
debug= False
).T

Rh = lrp_linear(
h_prev,
U_g,
b_g,
g_pre,
R_g[:, t, :],
bias_nb_units=input_dim + hidden_dim,
eps=epsilon,
bias_factor=0.0,
debug= False
).T

R_input[:, t, :] = Rx
R_h[:, t - 1, :] += Rh

total_in = np.sum(R_h_input)
total_out = np.sum(R_input) + np.sum(R_c[:, 0, :]) + np.sum(R_h[:, 0, :])
print("Différence de pertinence :", total_out - total_in)

R_reste = total_in - total_out
return R_input, R_reste

def lrp_linear(hin, w, b, hout, Rout, bias_nb_units, eps, bias_factor=0.0, debug=0):

# hin.shape=(1, 128)
# hout.shape=(1, 128)
# w.shape =(128, 128)
# b.shape =(128,)
# Rout.shape= (1, 128)
# bias_nb_units= 256

sign_out = np.where(hout >= 0, 1., -1.)
#sign_out.shape= (1, 128)

numer = (w * hin) + (bias_factor * (b+ eps * sign_out) / bias_nb_units)
# numer.shape= (128, 128)

denom = hout + (eps * sign_out)
# denom.shape=(1, 128)

message = (numer / denom) * Rout
# message.shape= (128, 128)

Rin = message.sum(axis=0)

# Rin.shape=(128,)

if debug:
print("local diff: ", Rout.sum() - Rin.sum())

return Rin



But the total input relevance doesn't match the total output relevance, so relevance is not preserved, leading to large discrepancies.

I tried printing the message.sum() and Rout.sum() who should match, but they show differences, and I don't get it where this difference comes from.

Continue reading...