122 lines
4.2 KiB
Python
122 lines
4.2 KiB
Python
import torch
|
||
import torch.nn as nn
|
||
import torch.nn.functional as F
|
||
|
||
class ChebSpectralReconstructor(nn.Module):
|
||
"""
|
||
Spectral reconstruction using Chebyshev polynomial approximation.
|
||
Input:
|
||
X: Tensor of shape [N, T, d] (node features over time)
|
||
A: Adjacency matrix [N, N]
|
||
Output:
|
||
X_tilde: Tensor [N, T, d*(K+2)] for downstream model
|
||
"""
|
||
def __init__(self, K=3, in_dim=16, hidden_dim=32, lr=1e-3, device='cpu'):
|
||
super().__init__()
|
||
self.K = K
|
||
self.in_dim = in_dim
|
||
self.hidden_dim = hidden_dim
|
||
self.device = device
|
||
|
||
# Spectral coefficients α_k (learnable)
|
||
self.alpha = nn.Parameter(torch.randn(K + 1, 1))
|
||
|
||
# Propagation weights W_k (learnable)
|
||
self.W = nn.ParameterList([
|
||
nn.Parameter(torch.randn(in_dim, hidden_dim)) for _ in range(K + 1)
|
||
])
|
||
|
||
self.optimizer = torch.optim.Adam(self.parameters(), lr=lr)
|
||
self.to(device)
|
||
|
||
def compute_normalized_laplacian(self, A):
|
||
"""Compute normalized Laplacian L = I - D^{-1/2} A D^{-1/2}"""
|
||
I = torch.eye(A.size(0), device=A.device)
|
||
D = torch.diag(torch.sum(A, dim=1))
|
||
D_inv_sqrt = torch.pow(D, -0.5)
|
||
D_inv_sqrt[torch.isinf(D_inv_sqrt)] = 0.0
|
||
L = I - D_inv_sqrt @ A @ D_inv_sqrt
|
||
return L
|
||
|
||
def chebyshev_polynomials(self, L_tilde, X_t):
|
||
"""Compute [T_0(L_tilde)X_t, ..., T_K(L_tilde)X_t] recursively"""
|
||
T_k_list = [X_t]
|
||
if self.K >= 1:
|
||
T_k_list.append(L_tilde @ X_t)
|
||
for k in range(2, self.K + 1):
|
||
T_k_list.append(2 * L_tilde @ T_k_list[-1] - T_k_list[-2])
|
||
return T_k_list
|
||
|
||
def forward_one_step(self, X_t, L):
|
||
"""Compute one-step reconstruction of X_{t+1}"""
|
||
lambda_max = torch.linalg.eigvalsh(L).max().real
|
||
L_tilde = (2.0 / lambda_max) * L - torch.eye(L.size(0), device=L.device)
|
||
T_k_list = self.chebyshev_polynomials(L_tilde, X_t)
|
||
|
||
H_list = []
|
||
for k in range(self.K + 1):
|
||
H_k = T_k_list[k] @ self.W[k]
|
||
H_list.append(H_k)
|
||
H_stack = torch.stack(H_list, dim=-1) # [N, hidden_dim, K+1]
|
||
H_weighted = torch.sum(H_stack * self.alpha.view(1, 1, -1), dim=-1)
|
||
X_hat_next = torch.tanh(H_weighted)
|
||
return X_hat_next, H_list
|
||
|
||
def forward_sequence(self, X, A):
|
||
"""
|
||
Compute decoupled multi-scale representation over all time steps.
|
||
X: [N, T, d], A: [N, N]
|
||
"""
|
||
L = self.compute_normalized_laplacian(A)
|
||
N, T, d = X.shape
|
||
multi_scale_list = [X[:, 0, :]] # include X_1
|
||
|
||
for t in range(T - 1):
|
||
X_hat_next, H_list = self.forward_one_step(X[:, t, :], L)
|
||
X_tilde_next = torch.cat(H_list + [X_hat_next], dim=-1)
|
||
multi_scale_list.append(X_tilde_next)
|
||
|
||
X_tilde = torch.stack(multi_scale_list, dim=1)
|
||
# Shape: [N, T, hidden_dim*(K+2)]
|
||
return X_tilde
|
||
|
||
def loss_fn(self, X_true, X_pred):
|
||
return F.mse_loss(X_true, X_pred)
|
||
|
||
def fit(self, X, A, num_epochs=200, verbose=True):
|
||
"""Optimization process to learn α_k and W_k"""
|
||
L = self.compute_normalized_laplacian(A)
|
||
for epoch in range(num_epochs):
|
||
total_loss = 0
|
||
for t in range(X.size(1) - 1):
|
||
X_t = X[:, t, :]
|
||
X_true_next = X[:, t + 1, :]
|
||
X_pred_next, _ = self.forward_one_step(X_t, L)
|
||
loss = self.loss_fn(X_true_next, X_pred_next)
|
||
|
||
self.optimizer.zero_grad()
|
||
loss.backward()
|
||
self.optimizer.step()
|
||
total_loss += loss.item()
|
||
|
||
if verbose and (epoch + 1) % 10 == 0:
|
||
print(f"Epoch [{epoch+1}/{num_epochs}] | Loss: {total_loss:.6f}")
|
||
|
||
print("Training complete.")
|
||
|
||
# ---------------------------
|
||
# Example usage
|
||
# ---------------------------
|
||
if __name__ == "__main__":
|
||
N, T, d = 20, 10, 16
|
||
A = torch.rand(N, N)
|
||
A = (A + A.T) / 2 # symmetrize
|
||
X = torch.randn(N, T, d)
|
||
|
||
model = ChebSpectralReconstructor(K=3, in_dim=d, hidden_dim=16, lr=1e-3, device='cpu')
|
||
model.fit(X, A, num_epochs=50)
|
||
|
||
X_tilde = model.forward_sequence(X, A)
|
||
print("Final embedding shape:", X_tilde.shape) # [N, T, hidden_dim*(K+2)]
|
||
|