KFLSTMModel

Core EKF-LSTM model components. Source: src/KFLSTMModel.py.

---

FeedForwardDecoder

FeedForwardDecoder(input_size, hidden_size, output_size)

Bases: Module

Two-layer feedforward network used as the nonlinear observation function C(h). Maps the EKF state h to predicted sensor measurements y_h = C(h).

Architecture: Linear(input_size, hidden_size) -> ReLU -> Linear(hidden_size, output_size).

Source code in src/KFLSTMModel.py

def __init__(self, input_size, hidden_size, output_size):
    super(FeedForwardDecoder, self).__init__()
    self.fc1 = nn.Linear(input_size, hidden_size)
    self.relu = nn.ReLU()
    self.fc2 = nn.Linear(hidden_size, output_size)

---

LSTMModel

LSTMModel(input_size, hidden_size, output_size, num_layers)

Bases: Module

Single-layer (or multi-layer) LSTM with a linear output head and ReLU activation. Used as a sub-model within KFLSTMModel for state transition (model_f), process noise covariance (model_q), and measurement noise covariance (model_r).

Source code in src/KFLSTMModel.py

def __init__(self, input_size, hidden_size, output_size, num_layers):
    super(LSTMModel, self).__init__()
    self.lstm = nn.LSTM(input_size, hidden_size, num_layers, batch_first=True)
    self.relu = nn.ReLU()
    self.fc = nn.Linear(hidden_size, int(output_size))

---

KFLSTMModel

KFLSTMModel(is_f, is_q, is_r, os_f, os_q, os_r, hidden_size, num_layers, ff_hidden_size, lr, la_mse, la_nll_r, la_nll_h, la_CPCT_cov, la_K)

Bases: Module

LSTM-based Extended Kalman Filter (EKF) model for ICS time series anomaly detection.

Learns to predict the EKF residual r_t and innovation covariance S_t such that r_t ~ N(0, S_t). The model comprises four learned sub-models: - model_f: LSTM state transition function f(h) -> h_next - model_C: FeedForward observation function C(h) -> y_h - model_q: LSTM predicting lower-triangular process noise covariance Q_t - model_r: LSTM predicting lower-triangular measurement noise covariance R_t

Jacobians of model_f and model_C are computed via torch.autograd to linearize the nonlinear dynamics for the EKF covariance update equations.

Parameters:

Name	Type	Description	Default
is_f	int	State space dimension (SD).	required
is_q	int	Input dimension for the process noise LSTM (= is_f).	required
is_r	int	Measurement space dimension (MD).	required
os_f	int	Output dimension for the state transition LSTM (= is_f).	required
os_q	int	Number of lower-triangular elements in Q_t (= is_f*(is_f+1)//2).	required
os_r	int	Number of lower-triangular elements in R_t (= is_r*(is_r+1)//2).	required
hidden_size	int	LSTM hidden state size.	required
num_layers	int	Number of LSTM layers.	required
ff_hidden_size	int	Hidden size of the FeedForwardDecoder.	required
lr	float	Learning rate (stored for reference).	required
la_mse, la_nll_r, la_nll_h, la_CPCT_cov, la_K	float	Loss amplifier weights.	required

Source code in src/KFLSTMModel.py

def __init__(self, is_f, is_q, is_r, os_f, os_q, os_r, hidden_size, num_layers, ff_hidden_size,
             lr, la_mse, la_nll_r, la_nll_h, la_CPCT_cov, la_K):
    super(KFLSTMModel, self).__init__()

    # Set batch_size to 1 for ONNX compatibility
    self.batch_size = 1  # Force batch_size=1
    self.is_f = is_f

    # Initialize matrices
    P = torch.tril(torch.randn(is_f, is_f))
    self.P = torch.matmul(P, P.T)

    CPCT_cov = torch.tril(torch.randn(is_r, is_r))
    self.CPCT_cov = torch.matmul(CPCT_cov, CPCT_cov.T)

    CPT = torch.tril(torch.randn(is_f, is_f))
    self.CPT = torch.matmul(CPT, CPT.T)

    self.h_out = torch.randn(is_f, 1)
    self.K = torch.ones(is_f, is_r)
    self.KC = torch.ones(is_f, is_f)

    # Store dimensions
    self.is_q = is_q
    self.is_r = is_r
    self.os_f = os_f
    self.os_q = os_q
    self.os_r = os_r
    self.hidden_size = hidden_size
    self.num_layers = num_layers
    self.total_loss = 0
    self.eps = 1e-2
    self.la_mse = la_mse
    self.la_nll_r = la_nll_r
    self.la_nll_h = la_nll_h
    self.la_CPCT_cov = la_CPCT_cov
    self.la_K = la_K

    self.sq_cpc_penalty = False

    # Initialize sub-models
    self.model_C = FeedForwardDecoder(is_f, ff_hidden_size, is_r)
    self.model_f = LSTMModel(is_f, hidden_size, os_f, num_layers)
    self.model_w = LSTMModel(is_f, hidden_size, os_f, num_layers)
    self.model_q = LSTMModel(is_q, hidden_size, os_q, num_layers)
    self.model_r = LSTMModel(is_r, hidden_size, os_r, num_layers)

    # Register buffers for matrix reconstruction
    self.register_buffer("h_tril_indices", torch.tril_indices(4,4, offset=0))
    self.register_buffer("r_tril_indices", torch.tril_indices(5,5, offset=0))

    self.kalman_recon_check = KalmanReconstructionCheck()
    self.If = torch.eye(self.is_f)

computeKalmanCond

computeKalmanCond(Kt, bsize, compute_stat=False)

Reshapes the Kalman gain to (SD x MD). Optionally computes and prints condition number diagnostics (singular values, spectral norms) for debugging.

Parameters:

Name	Type	Description	Default
Kt	Tensor	Kalman gain K^T.	required
bsize	int	Batch size.	required
compute_stat	bool	If True, prints SVD-based condition number stats.	False

Returns:

Type	Description
	torch.Tensor: Kalman gain reshaped as (SD x MD).

Source code in src/KFLSTMModel.py

def computeKalmanCond(self,Kt,bsize,compute_stat=False):
    """
    Reshapes the Kalman gain to (SD x MD). Optionally computes and prints
    condition number diagnostics (singular values, spectral norms) for debugging.

    Args:
        Kt (torch.Tensor): Kalman gain K^T.
        bsize (int): Batch size.
        compute_stat (bool): If True, prints SVD-based condition number stats.

    Returns:
        torch.Tensor: Kalman gain reshaped as (SD x MD).
    """
    K2 = Kt.reshape(self.is_f, self.is_r)
    if compute_stat:
        K2 = Kt #torch.bmm(torch.transpose(Kt, 1, 2), Kt).reshape(self.is_r, self.is_r)
        K2_s1 = torch.linalg.norm(K2, ord=2)
        K2_sn = torch.linalg.norm(K2, ord=-2)
        K2_cond = K2_s1 / K2_sn
        U, S, Vh = torch.linalg.svd(K2)
        print(S, K2_s1, K2_sn, K2_cond)
    return K2

forward

forward(hin, yin, jac_fh, jac_Ch, P, Kt)

Primary EKF forward pass used in the zkSNARK prediction circuit.

Takes pre-computed Jacobians and Kalman gain as explicit circuit inputs so they can be committed to via Poseidon hashes. Runs one EKF update step and returns the updated state and all intermediate quantities.

Parameters:

Name	Type	Description	Default
hin	Tensor	Current state estimate h_t (1 x 1 x SD).	required
yin	Tensor	Current sensor measurement y_t (1 x 1 x MD).	required
jac_fh	Tensor	Pre-computed Jacobian of f w.r.t. h (1 x SD x SD).	required
jac_Ch	Tensor	Pre-computed Jacobian of C w.r.t. h (1 x MD x SD).	required
P	Tensor	Current state error covariance (1 x SD x SD).	required
Kt	Tensor	Pre-computed Kalman gain K^T (1 x MD x SD).	required

Returns:

Name	Type	Description
Tuple		(h_out, out_h, dy, cov_h, cov_r, CPCT_cov, P_updated, kalman_check_status, kalman_diff, kalman_threshold)

Source code in src/KFLSTMModel.py

def forward(self,hin, yin, jac_fh, jac_Ch, P, Kt,):
    """
    Primary EKF forward pass used in the zkSNARK prediction circuit.

    Takes pre-computed Jacobians and Kalman gain as explicit circuit inputs
    so they can be committed to via Poseidon hashes. Runs one EKF update step
    and returns the updated state and all intermediate quantities.

    Args:
        hin (torch.Tensor): Current state estimate h_t (1 x 1 x SD).
        yin (torch.Tensor): Current sensor measurement y_t (1 x 1 x MD).
        jac_fh (torch.Tensor): Pre-computed Jacobian of f w.r.t. h (1 x SD x SD).
        jac_Ch (torch.Tensor): Pre-computed Jacobian of C w.r.t. h (1 x MD x SD).
        P (torch.Tensor): Current state error covariance (1 x SD x SD).
        Kt (torch.Tensor): Pre-computed Kalman gain K^T (1 x MD x SD).

    Returns:
        Tuple: (h_out, out_h, dy, cov_h, cov_r, CPCT_cov, P_updated,
                kalman_check_status, kalman_diff, kalman_threshold)
    """
    out_h = self.model_f(hin)
    y_h = self.model_C(out_h)
    cov_h = self.model_q(hin)
    cov_r = self.model_r(yin)

    cov_h_raw = torch.exp(cov_h).reshape(self.os_q, )
    cov_r_raw = torch.exp(cov_r).reshape(self.os_r, )

    cov_h = torch.zeros(self.is_f, self.is_f)
    h_tril_idx = torch.tril_indices(self.is_f, self.is_f, offset=0)
    cov_h[h_tril_idx[0], h_tril_idx[1]] = cov_h_raw

    cov_r = torch.zeros((self.is_r, self.is_r))
    r_tril_idx = torch.tril_indices(self.is_r, self.is_r, offset=0)
    cov_r[r_tril_idx[0], r_tril_idx[1]] = cov_r_raw

    cov_h = cov_h.reshape(1, self.is_f, self.is_f)
    cov_r = cov_r.reshape(1, self.is_r, self.is_r)

    cov_h = torch.bmm(torch.transpose(cov_h, 1, 2), cov_h)
    cov_r = torch.bmm(torch.transpose(cov_r, 1, 2), cov_r)

    jac_fh_T = jac_fh.transpose(1, 2)
    P = torch.bmm(torch.bmm(jac_fh, P), jac_fh_T)
    P = torch.add(P,cov_h)

    jac_Ch_T = jac_Ch.transpose(1, 2)
    CPCT_cov = torch.bmm(torch.bmm(jac_Ch, P), jac_Ch_T) + cov_r

    # print("zkforward:CPCT_cov:",torch.linalg.norm(CPCT_cov))
    # print("zkforward:cov_r:", torch.linalg.norm(cov_r))
    # print("zkforward:jac_Ch:", torch.linalg.norm(jac_Ch))
    # print("zkforward:P:", torch.linalg.norm(P))

    kalman_check_status,kalman_diff, kalman_threshold = self.kalman_recon_check.forward(CPCT_cov,jac_Ch,P,Kt)

    K = Kt.transpose(1, 2)

    KC = torch.bmm(K, jac_Ch)

    P = torch.bmm(self.If - KC, P)

    dy = yin - y_h

    dh_out = torch.bmm(K, dy.transpose(1, 2))

    out_h = out_h.reshape(1,self.is_f,1)
    #print(out_h.shape,h_out.shape)

    h_out = dh_out + out_h
    #print(dh_out,K,Kt,h_out,dy)


    return h_out, out_h, dy, cov_h, cov_r, CPCT_cov, P, kalman_check_status, kalman_diff,kalman_threshold

forward_one_step

forward_one_step(hin, yin)

Single EKF step with internally computed Jacobians. Used during non-ZKP inference and testing. Iterates over the batch dimension.

Parameters:

Name	Type	Description	Default
hin	Tensor	Batch of state inputs (batch x 1 x SD).	required
yin	Tensor	Batch of measurement inputs (batch x 1 x MD).	required

Returns:

Name	Type	Description
Tuple		(h_out, delta_h, delta_y, cov_h, cov_r, CPCT_cov, K) all reshaped to (batch, ...).

Source code in src/KFLSTMModel.py

def forward_one_step(self, hin, yin):
    """
    Single EKF step with internally computed Jacobians. Used during
    non-ZKP inference and testing. Iterates over the batch dimension.

    Args:
        hin (torch.Tensor): Batch of state inputs (batch x 1 x SD).
        yin (torch.Tensor): Batch of measurement inputs (batch x 1 x MD).

    Returns:
        Tuple: (h_out, delta_h, delta_y, cov_h, cov_r, CPCT_cov, K)
            all reshaped to (batch, ...).
    """
    h_out_list = []
    delta_y = []
    delta_h = []
    cov_h_list = []
    cov_r_list = []
    CPCT_cov_list = []
    K_list = []

    bsize = hin.shape[0]
    for b in range(bsize):
      h=hin[b,:,:]

      y=yin[b,:,:]

      out_h = self.model_f(h)

      y_h = self.model_C(out_h)

      jac_fh = torch.autograd.functional.jacobian(self.output_function_f,h,create_graph=True)
      jac_fh = jac_fh.reshape(self.is_f, self.is_f)

      jac_Ch = torch.autograd.functional.jacobian(self.output_function_C,out_h,create_graph=True)
      jac_Ch = jac_Ch.reshape(self.is_r, self.is_f)

      cov_h = self.model_q(out_h)
      cov_r = self.model_r(y)

      cov_h_raw = torch.exp(cov_h).reshape(self.os_q,)
      cov_r_raw = torch.exp(cov_r).reshape(self.os_r,)


      cov_h = torch.zeros(self.is_f,self.is_f)
      h_tril_indices = torch.tril_indices(self.is_f, self.is_f,offset=0)
      cov_h[h_tril_indices.tolist()] = cov_h_raw

      cov_r = torch.zeros((self.is_r,self.is_r))
      r_tril_indices = torch.tril_indices(self.is_r, self.is_r)
      cov_r[r_tril_indices.tolist()] = cov_r_raw


      self.P = jac_fh @ self.P @ jac_fh.T + cov_h.T@cov_h

      CPCT_cov = jac_Ch @ self.P @ jac_Ch.T + cov_r.T@cov_r

      CPT = jac_Ch @ self.P.T

      Kt = torch.linalg.solve(CPCT_cov.T, CPT)

      self.K = Kt.T

      self.KC = self.K@jac_Ch

      I = torch.eye(self.is_f)
      self.P = (I - self.KC) @ self.P
      dy = (y-y_h.reshape_as(y))

      h_out = out_h.T + self.K @ dy.T

      self.h_out = h_out.clone().detach().reshape(1,1,self.is_f)

      delta_y.append(dy)

      h_out_list.append(h_out)
      delta_h.append(out_h.T)

      cov_h_list.append(cov_h)
      cov_r_list.append(cov_r)
      CPCT_cov_list.append(CPCT_cov)

      K2 = self.computeKalmanCond(Kt,bsize,)

      K_list.append(K2)

    return (torch.stack(h_out_list).reshape(bsize,hin.shape[2]),
            torch.stack(delta_h).reshape(bsize,hin.shape[2]),
            torch.stack(delta_y).reshape(bsize,yin.shape[2]),
            torch.stack(cov_h_list).reshape(bsize,hin.shape[2],hin.shape[2]),
            torch.stack(cov_r_list).reshape(bsize,yin.shape[2],yin.shape[2]),
            torch.stack(CPCT_cov_list).reshape(bsize,yin.shape[2],yin.shape[2]),
            torch.stack(K_list).reshape(bsize, hin.shape[2], yin.shape[2])
            )

forward_train

forward_train(hin, yin)

Batched EKF forward pass for training. Computes Jacobians in batch (detached for efficiency) and returns all quantities needed for loss_func().

Parameters:

Name	Type	Description	Default
hin	Tensor	Batch of state inputs (batch x 1 x SD).	required
yin	Tensor	Batch of measurement inputs (batch x 1 x MD).	required

Returns:

Name	Type	Description
Tuple		(h_out, out_h, dy, cov_h, cov_r, CPCT_cov, K2) all shaped (batch, ...).

Source code in src/KFLSTMModel.py

def forward_train(self, hin, yin):
    """
    Batched EKF forward pass for training. Computes Jacobians in batch
    (detached for efficiency) and returns all quantities needed for loss_func().

    Args:
        hin (torch.Tensor): Batch of state inputs (batch x 1 x SD).
        yin (torch.Tensor): Batch of measurement inputs (batch x 1 x MD).

    Returns:
        Tuple: (h_out, out_h, dy, cov_h, cov_r, CPCT_cov, K2)
            all shaped (batch, ...).
    """
    bsize = hin.shape[0]
    out_h = self.model_f(hin)
    y_h = self.model_C(out_h)
    cov_h = self.model_q(hin)
    cov_r = self.model_r(yin)
    P = self.P.reshape(bsize, self.is_f, self.is_f)

    cov_h_raw = torch.exp(cov_h).reshape(self.os_q, )
    cov_r_raw = torch.exp(cov_r).reshape(self.os_r, )

    cov_h = torch.zeros(self.is_f, self.is_f)
    h_tril_indices = torch.tril_indices(self.is_f, self.is_f, offset=0)
    cov_h[h_tril_indices.tolist()] = cov_h_raw

    cov_r = torch.zeros((self.is_r, self.is_r))
    r_tril_indices = torch.tril_indices(self.is_r, self.is_r)
    cov_r[r_tril_indices.tolist()] = cov_r_raw

    cov_h = cov_h.reshape(bsize, self.is_f, self.is_f)
    cov_r = cov_r.reshape(bsize, self.is_r, self.is_r)

    cov_h = torch.bmm(torch.transpose(cov_h, 1, 2), cov_h)
    cov_r = torch.bmm(torch.transpose(cov_r, 1, 2), cov_r)

    self.out_h = Variable(out_h.clone(), requires_grad=True)
    self.y_h = Variable(y_h.clone(), requires_grad=True)
    self.cov_h = Variable(cov_h.clone(), requires_grad=True)
    self.cov_r = Variable(cov_r.clone(), requires_grad=True)
    self.P = Variable(P.clone(), requires_grad=True)

    jac_fh = torch.autograd.functional.jacobian(self.output_function_f, hin, create_graph=True).detach()
    jac_fh = jac_fh.reshape(bsize, self.is_f, self.is_f)

    jac_Ch = torch.autograd.functional.jacobian(self.output_function_C, hin, create_graph=True).detach()
    jac_Ch = jac_Ch.reshape(bsize, self.is_r, self.is_f)

    jac_fh_transposed = torch.transpose(jac_fh, 1, 2)
    P = torch.bmm(torch.bmm(jac_fh, P), jac_fh_transposed) + cov_h.reshape(bsize, self.is_f, self.is_f)


    jac_Ch_transposed = torch.transpose(jac_Ch, 1, 2)
    CPCT_cov = torch.bmm(torch.bmm(jac_Ch, P), jac_Ch_transposed) + cov_r.reshape(bsize, self.is_r,self.is_r)
    self.CPCT_cov = Variable(CPCT_cov.clone(), requires_grad=True)

    # print(":CPCT_cov:", torch.linalg.norm(CPCT_cov))
    # print(":cov_r:", torch.linalg.norm(cov_r))
    # print(":jac_Ch:", torch.linalg.norm(jac_Ch))
    # print(":P:", torch.linalg.norm(P))


    CPT = torch.bmm(jac_Ch, torch.transpose(P, 1, 2))
    self.CPT = Variable(CPT.clone(), requires_grad=True)

    Kt = torch.linalg.solve(torch.transpose(CPCT_cov, 1, 2), CPT).detach()
    self.K = torch.transpose(Kt, 1, 2).detach()


    KC = torch.bmm(self.K, jac_Ch)
    self.KC = Variable(KC.clone(), requires_grad=True)

    I = torch.eye(self.is_f)
    self.P = torch.bmm(I - KC, P)
    P = self.P.reshape(bsize, self.is_f, self.is_f).detach()

    dy = (yin - y_h.reshape(bsize, 1, self.is_r))

    h_out = torch.bmm(self.K, torch.transpose(dy, 1, 2))

    out_h = out_h.reshape_as(h_out)
    #print(h_out.shape,self.is_r,self.is_f)
    h_out = h_out + out_h
    h_out = h_out.reshape(1,1,self.is_f)

    # h_out = h_out.reshape(bsize, self.is_f)
    # out_h = out_h.reshape(bsize, self.is_f)
    # dy = dy.reshape(bsize, self.is_r)
    # K2 = Kt.reshape(self.is_f, self.is_r)
    # print(P.shape)
    #print(torch.linalg.norm(hin),torch.linalg.norm(Kt), torch.linalg.norm(dy),  torch.linalg.norm(h_out), torch.linalg.norm(out_h))
    return h_out, out_h, dy, cov_h, cov_r, CPCT_cov, K2

forward_yield_jacobians

forward_yield_jacobians(hin, yin)

Full EKF step that additionally returns the Jacobians and Kalman gain.

Called before zkSNARK circuit export to produce the jac_fh, jac_Ch, and Kt inputs that will be committed to in the prediction circuit. These are then passed as explicit inputs to forward() during witness generation.

Parameters:

Name	Type	Description	Default
hin	Tensor	Current state estimate (1 x 1 x SD).	required
yin	Tensor	Current measurement (1 x 1 x MD).	required

Returns:

Name	Type	Description
Tuple		(h_out, jac_fh, jac_Ch, P_updated, Kt)

Source code in src/KFLSTMModel.py

def forward_yield_jacobians(self, hin, yin):
    """
    Full EKF step that additionally returns the Jacobians and Kalman gain.

    Called before zkSNARK circuit export to produce the jac_fh, jac_Ch, and
    Kt inputs that will be committed to in the prediction circuit. These are
    then passed as explicit inputs to forward() during witness generation.

    Args:
        hin (torch.Tensor): Current state estimate (1 x 1 x SD).
        yin (torch.Tensor): Current measurement (1 x 1 x MD).

    Returns:
        Tuple: (h_out, jac_fh, jac_Ch, P_updated, Kt)
    """
    bsize = hin.shape[0]
    out_h = self.model_f(hin)
    y_h = self.model_C(out_h)
    cov_h = self.model_q(hin)
    cov_r = self.model_r(yin)
    P = self.P.reshape(bsize, self.is_f, self.is_f)

    cov_h_raw = torch.exp(cov_h).reshape(self.os_q, )
    cov_r_raw = torch.exp(cov_r).reshape(self.os_r, )

    cov_h = torch.zeros(self.is_f, self.is_f)
    h_tril_indices = torch.tril_indices(self.is_f, self.is_f, offset=0)
    cov_h[h_tril_indices.tolist()] = cov_h_raw

    cov_r = torch.zeros((self.is_r, self.is_r))
    r_tril_indices = torch.tril_indices(self.is_r, self.is_r)
    cov_r[r_tril_indices.tolist()] = cov_r_raw

    cov_h = cov_h.reshape(bsize, self.is_f, self.is_f)
    cov_r = cov_r.reshape(bsize, self.is_r, self.is_r)

    cov_h = torch.bmm(torch.transpose(cov_h, 1, 2), cov_h)
    cov_r = torch.bmm(torch.transpose(cov_r, 1, 2), cov_r)

    self.out_h = Variable(out_h.clone(), requires_grad=True)
    self.y_h = Variable(y_h.clone(), requires_grad=True)
    self.cov_h = Variable(cov_h.clone(), requires_grad=True)
    self.cov_r = Variable(cov_r.clone(), requires_grad=True)
    self.P = Variable(P.clone(), requires_grad=True)

    jac_fh = torch.autograd.functional.jacobian(self.output_function_f, hin, create_graph=True).detach()
    jac_fh = jac_fh.reshape(bsize, self.is_f, self.is_f)

    jac_Ch = torch.autograd.functional.jacobian(self.output_function_C, hin, create_graph=True).detach()
    jac_Ch = jac_Ch.reshape(bsize, self.is_r, self.is_f)

    jac_fh_transposed = torch.transpose(jac_fh, 1, 2)
    P = torch.bmm(torch.bmm(jac_fh, P), jac_fh_transposed) + cov_h.reshape(bsize, self.is_f, self.is_f)


    jac_Ch_transposed = torch.transpose(jac_Ch, 1, 2)
    CPCT_cov = torch.bmm(torch.bmm(jac_Ch, P), jac_Ch_transposed) + cov_r.reshape(bsize, self.is_r,self.is_r)
    self.CPCT_cov = Variable(CPCT_cov.clone(), requires_grad=True)

    # print(":CPCT_cov:", torch.linalg.norm(CPCT_cov))
    # print(":cov_r:", torch.linalg.norm(cov_r))
    # print(":jac_Ch:", torch.linalg.norm(jac_Ch))
    # print(":P:", torch.linalg.norm(P))


    CPT = torch.bmm(jac_Ch, torch.transpose(P, 1, 2))
    self.CPT = Variable(CPT.clone(), requires_grad=True)

    Kt = torch.linalg.solve(torch.transpose(CPCT_cov, 1, 2), CPT).detach()
    self.K = torch.transpose(Kt, 1, 2).detach()


    KC = torch.bmm(self.K, jac_Ch)
    self.KC = Variable(KC.clone(), requires_grad=True)

    I = torch.eye(self.is_f)
    self.P = torch.bmm(I - KC, P)
    P = self.P.reshape(bsize, self.is_f, self.is_f).detach()

    dy = (yin - y_h.reshape(bsize, 1, self.is_r))

    h_out = torch.bmm(self.K, torch.transpose(dy, 1, 2))

    out_h = out_h.reshape_as(h_out)
    #print(h_out.shape,self.is_r,self.is_f)
    h_out = h_out + out_h
    h_out = h_out.reshape(1,1,self.is_f)

    # h_out = h_out.reshape(bsize, self.is_f)
    # out_h = out_h.reshape(bsize, self.is_f)
    # dy = dy.reshape(bsize, self.is_r)
    # K2 = Kt.reshape(self.is_f, self.is_r)
    # print(P.shape)
    #print(torch.linalg.norm(hin),torch.linalg.norm(Kt), torch.linalg.norm(dy),  torch.linalg.norm(h_out), torch.linalg.norm(out_h))
    return h_out, jac_fh, jac_Ch, P, Kt, #h_out, out_h, dy, cov_h, cov_r, CPCT_cov, K2

loss_func

loss_func(outputs, targets, delta_h, delta_y, cov_h_list, cov_r_list, CPCT_cov, K)

Composite training loss combining state prediction MSE with condition number penalties on the learned covariance matrices.

Loss = la_mse * MSE(h_pred, h_real) + la_nll_r * SSE(cond(R_t), 1) + la_nll_h * SSE(cond(Q_t), 1) + la_CPCT_cov * SSE(cond(CPCT_cov), 1) + la_K * SSE(cond(K_t), 1)

The condition number penalties encourage well-conditioned covariance matrices, improving numerical stability in both the EKF and the zkSNARK fixed-point arithmetic.

Returns:

Type	Description
	Tuple[torch.Tensor, ...]: (total_loss, mse_loss, nll_loss_r, nll_loss_h)

Source code in src/KFLSTMModel.py

def loss_func(self, outputs, targets, delta_h, delta_y, cov_h_list, cov_r_list, CPCT_cov, K):
    """
    Composite training loss combining state prediction MSE with condition number
    penalties on the learned covariance matrices.

    Loss = la_mse * MSE(h_pred, h_real)
         + la_nll_r * SSE(cond(R_t), 1)
         + la_nll_h * SSE(cond(Q_t), 1)
         + la_CPCT_cov * SSE(cond(CPCT_cov), 1)
         + la_K * SSE(cond(K_t), 1)

    The condition number penalties encourage well-conditioned covariance matrices,
    improving numerical stability in both the EKF and the zkSNARK fixed-point arithmetic.

    Returns:
        Tuple[torch.Tensor, ...]: (total_loss, mse_loss, nll_loss_r, nll_loss_h)
    """
    if torch.cuda.is_available():
        outputs = outputs.to("cuda")
        targets = targets.to("cuda")
        delta_h = delta_h.to("cuda")
        delta_y = delta_y.to("cuda")
        cov_h_list = cov_h_list.to("cuda")
        cov_r_list = cov_r_list.to("cuda")

    real_dh = outputs - targets

    # mse_loss = F.binary_cross_entropy(outputs, targets, reduction="mean")
    lf = nn.MSELoss(reduction="mean")

    mse_loss = lf(outputs.reshape_as(targets), targets)

    mse_loss = mse_loss + self.la_mse * lf(delta_h.reshape_as(targets), targets)

    cov_r_ltril = cov_r_list
    cov_h_ltril = cov_h_list

    nll_loss_h = 0
    nll_loss_r = 0
    nll_loss_cpct = 0
    nll_loss_K = 0

    if self.la_nll_h:
        if self.sq_cpc_penalty:
            cov_h = torch.matmul(torch.transpose(cov_h_ltril, dim0=0, dim1=1), cov_h_ltril)
        else:
            cov_h_ltril = cov_h_ltril.reshape(self.is_f, self.is_f)
            cov_h = cov_h_ltril.T @ cov_h_ltril

        cov_h_s1 = torch.linalg.norm(cov_h, ord=2)
        cov_h_sn = torch.linalg.norm(cov_h, ord=-2)
        h_cond = cov_h_s1 / cov_h_sn
        target_h_cond = torch.ones_like(h_cond)

        nll_loss_h = self.la_nll_h * lf(h_cond, target_h_cond)

    if self.la_nll_r:
        if self.sq_cpc_penalty:
            cov_r = torch.matmul(torch.transpose(cov_r_ltril, dim0=0, dim1=1), cov_r_ltril)
        else:
            cov_r_ltril = cov_r_ltril.reshape(self.is_r, self.is_r)
            cov_r = cov_r_ltril.T @ cov_r_ltril #################################

        cov_r_s1 = torch.linalg.norm(cov_r, ord=2)
        cov_r_sn = torch.linalg.norm(cov_r, ord=-2)
        r_cond = cov_r_s1 / cov_r_sn
        target_r_cond = torch.ones_like(r_cond)
        nll_loss_r = self.la_nll_r * lf(r_cond, target_r_cond)

    if self.la_CPCT_cov:
        if self.sq_cpc_penalty:
            CPCT_cov = torch.matmul(torch.transpose(CPCT_cov, dim0=1, dim1=2), CPCT_cov)
        else:
            CPCT_cov = CPCT_cov.reshape(self.is_r, self.is_r)

        CPCT_cov_s1 = torch.linalg.norm(CPCT_cov, ord=2)
        CPCT_cov_sn = torch.linalg.norm(CPCT_cov, ord=-2)
        CPCT_cov_cond = CPCT_cov_s1 / CPCT_cov_sn
        target_CPCT_cov_cond = torch.ones_like(CPCT_cov_cond)
        nll_loss_cpct = self.la_CPCT_cov * lf(CPCT_cov_cond, target_CPCT_cov_cond)

    if self.la_K:
        if self.sq_cpc_penalty:
            K = torch.matmul(torch.transpose(K, dim0=1, dim1=2), K)
        else:
            K = K.reshape(self.is_f, self.is_r)

        K_s1 = torch.linalg.norm(K, ord=2)
        K_sn = torch.linalg.norm(K, ord=-2)
        K_cond = K_s1 / K_sn
        target_K_cond = torch.ones_like(K_cond)
        nll_loss_K = self.la_k * lf(K_cond, target_K_cond)

    total_loss = mse_loss + (nll_loss_r + nll_loss_h + nll_loss_cpct + nll_loss_K)
    self.total_loss = total_loss

    return total_loss, mse_loss, nll_loss_r, nll_loss_h

squared_exponential_kernel

squared_exponential_kernel(x1, x2, lengthscale, variance)

Computes the squared exponential kernel function between two sets of points.

Parameters:

Name	Type	Description	Default
x1	Tensor	First set of points (shape: [N, D])	required
x2	Tensor	Second set of points (shape: [M, D])	required
lengthscale	Tensor	Lengthscale parameter (shape: [1] or [D])	required
variance	Tensor	Variance parameter (shape: [1])	required

Returns:

Type	Description
	torch.Tensor: Kernel matrix (shape: [N, M])

Source code in src/KFLSTMModel.py

def squared_exponential_kernel(self, x1, x2, lengthscale, variance):
    """
    Computes the squared exponential kernel function between two sets of points.

    Args:
        x1 (torch.Tensor): First set of points (shape: [N, D])
        x2 (torch.Tensor): Second set of points (shape: [M, D])
        lengthscale (torch.Tensor): Lengthscale parameter (shape: [1] or [D])
        variance (torch.Tensor): Variance parameter (shape: [1])

    Returns:
        torch.Tensor: Kernel matrix (shape: [N, M])
    """
    # Compute the pairwise squared Euclidean distances
    pairwise_sq_dists = torch.cdist(x1.reshape(1,x1.shape[1],1) / lengthscale, x2.reshape(1,x2.shape[1],1) / lengthscale, p=2).pow(2)
    # Compute the kernel matrix
    kernel_matrix = variance * torch.exp(-0.5 * pairwise_sq_dists)

    #kernel_matrix = kernel_matrix + 1e-3*torch.eye(x1.shape[0],x1.shape[0])

    return kernel_matrix.reshape(x1.shape[1],x2.shape[1])