CausalTAD解读

The code uses both VAE and confidence modeling for trajectory embedding and anomaly detection. Here's how they work together:

The Confidence Model:

1. Takes trajectory sequences and embeds them using an embedding layer
2. Encodes into latent space using mu and sigma encoders
3. Samples from the latent distribution and decodes
4. Computes a confidence score representing how "normal" each trajectory point is
5. Returns NLL loss + KL divergence as the confidence measure

The VAE Architecture:

1. Encoder:

   • Takes source trajectory and encodes into latent distribution (mu, sigma)
   • Uses RNN for sequence encoding

2. Decoder:

   • Samples from latent space
   • Decodes into trajectory reconstruction
   • Has two decoders:
     • Main decoder reconstructs full trajectory
     • SD decoder for source-destination prediction

3. Training flow:

confidence = self.confidence(src)  # Get confidence scores
src = self.road_embedding(src)  # Embed road segments
kl_loss, p_x, sd_p_x = self.vae(src, trg)  # VAE encoding/decoding
nll_loss = self.loss_fn(p_x, label, mask)  # Reconstruction loss

The combined architecture helps detect anomalies by:

• Confidence scores identify unusual patterns
• VAE reconstruction error highlights deviations from normal trajectories
• SD decoder ensures start/end points are reasonable
• Masked loss focuses on valid road segment transitions

All losses are combined to train the model to recognize normal vs anomalous trajectories.

revise

import torch
import torch.nn as nn
import math

class TransformerTrajectoryModel(nn.Module):
    def __init__(self, num_nodes, d_model=128, nhead=8, num_encoder_layers=6, 
                 num_decoder_layers=6, dim_feedforward=512, dropout=0.1):
        super().__init__()
        self.d_model = d_model
        
        # Node embeddings
        self.node_embedding = nn.Embedding(num_nodes, d_model)
        self.pos_encoder = PositionalEncoding(d_model, dropout)
        
        # Transformer
        self.transformer = nn.Transformer(
            d_model=d_model,
            nhead=nhead,
            num_encoder_layers=num_encoder_layers,
            num_decoder_layers=num_decoder_layers,
            dim_feedforward=dim_feedforward,
            dropout=dropout,
            batch_first=True
        )
        
        # Confidence module
        self.confidence_encoder = nn.Sequential(
            nn.Linear(d_model, d_model),
            nn.ReLU(),
            nn.Linear(d_model, d_model)
        )
        self.confidence_mu = nn.Linear(d_model, d_model)
        self.confidence_logvar = nn.Linear(d_model, d_model)
        
        # Trajectory VAE
        self.encoder_mu = nn.Linear(d_model, d_model)
        self.encoder_logvar = nn.Linear(d_model, d_model)
        self.decoder = nn.Linear(d_model, num_nodes)
        
        # Output projection
        self.output_projection = nn.Linear(d_model, num_nodes)
        
    def generate_mask(self, sz, device):
        mask = (torch.triu(torch.ones(sz, sz)) == 1).transpose(0, 1)
        mask = mask.float().masked_fill(mask == 0, float('-inf'))
        return mask.to(device)
        
    def encode(self, src, src_mask=None):
        src = self.node_embedding(src) * math.sqrt(self.d_model)
        src = self.pos_encoder(src)
        memory = self.transformer.encoder(src, src_mask)
        return memory
        
    def decode(self, tgt, memory, tgt_mask=None, memory_mask=None):
        tgt = self.node_embedding(tgt) * math.sqrt(self.d_model)
        tgt = self.pos_encoder(tgt)
        output = self.transformer.decoder(tgt, memory, tgt_mask, memory_mask)
        return output
        
    def compute_confidence(self, x):
        h = self.confidence_encoder(x)
        mu = self.confidence_mu(h)
        logvar = self.confidence_logvar(h)
        std = torch.exp(0.5 * logvar)
        
        # Sample using reparameterization trick
        eps = torch.randn_like(std)
        z = mu + eps * std
        
        kl_div = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp(), dim=-1)
        return z, kl_div
        
    def forward(self, src, tgt, src_mask=None, tgt_mask=None, memory_mask=None):
        # Encode trajectory
        memory = self.encode(src, src_mask)
        
        # Compute confidence
        confidence_latent, confidence_kl = self.compute_confidence(memory)
        
        # Encode to latent space
        mu = self.encoder_mu(memory)
        logvar = self.encoder_logvar(memory)
        std = torch.exp(0.5 * logvar)
        
        # Sample using reparameterization trick
        eps = torch.randn_like(std)
        z = mu + eps * std
        
        # Compute VAE KL divergence
        vae_kl = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp(), dim=-1)
        
        # Decode
        if tgt_mask is None:
            tgt_mask = self.generate_mask(tgt.size(1), tgt.device)
        output = self.decode(tgt, z, tgt_mask, memory_mask)
        
        # Project to vocabulary space
        logits = self.output_projection(output)
        
        return {
            'logits': logits,
            'confidence_kl': confidence_kl,
            'vae_kl': vae_kl
        }

class PositionalEncoding(nn.Module):
    def __init__(self, d_model, dropout=0.1, max_len=5000):
        super().__init__()
        self.dropout = nn.Dropout(p=dropout)

        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0)
        self.register_buffer('pe', pe)

    def forward(self, x):
        x = x + self.pe[:, :x.size(1)]
        return self.dropout(x)

Key changes from the original:

1. Uses Transformer encoder-decoder instead of RNN
2. Maintains same VAE + confidence structure but implements with self-attention
3. Adds positional encoding for sequence information
4. Uses masked self-attention for autoregressive decoding
5. Combines confidence and VAE latent spaces through the transformer memory

Usage remains similar but benefits from parallel processing and better long-range dependencies.

import torch
import torch.nn as nn
import math

class TransformerTrajectoryModel(nn.Module):
def init(self, num_nodes, d_model=128, nhead=8, num_encoder_layers=6,
num_decoder_layers=6, dim_feedforward=512, dropout=0.1):
super().init()
self.d_model = d_model

    # Node embeddings
    self.node_embedding = nn.Embedding(num_nodes, d_model)
    self.pos_encoder = PositionalEncoding(d_model, dropout)
    
    # Transformer
    self.transformer = nn.Transformer(
        d_model=d_model,
        nhead=nhead,
        num_encoder_layers=num_encoder_layers,
        num_decoder_layers=num_decoder_layers,
        dim_feedforward=dim_feedforward,
        dropout=dropout,
        batch_first=True
    )
    
    # Confidence module
    self.confidence_encoder = nn.Sequential(
        nn.Linear(d_model, d_model),
        nn.ReLU(),
        nn.Linear(d_model, d_model)
    )
    self.confidence_mu = nn.Linear(d_model, d_model)
    self.confidence_logvar = nn.Linear(d_model, d_model)
    
    # Trajectory VAE
    self.encoder_mu = nn.Linear(d_model, d_model)
    self.encoder_logvar = nn.Linear(d_model, d_model)
    self.decoder = nn.Linear(d_model, num_nodes)
    
    # Output projection
    self.output_projection = nn.Linear(d_model, num_nodes)
    
def generate_mask(self, sz, device):
    mask = (torch.triu(torch.ones(sz, sz)) == 1).transpose(0, 1)
    mask = mask.float().masked_fill(mask == 0, float('-inf'))
    return mask.to(device)
    
def encode(self, src, src_mask=None):
    src = self.node_embedding(src) * math.sqrt(self.d_model)
    src = self.pos_encoder(src)
    memory = self.transformer.encoder(src, src_mask)
    return memory
    
def decode(self, tgt, memory, tgt_mask=None, memory_mask=None):
    tgt = self.node_embedding(tgt) * math.sqrt(self.d_model)
    tgt = self.pos_encoder(tgt)
    output = self.transformer.decoder(tgt, memory, tgt_mask, memory_mask)
    return output
    
def compute_confidence(self, x):
    h = self.confidence_encoder(x)
    mu = self.confidence_mu(h)
    logvar = self.confidence_logvar(h)
    std = torch.exp(0.5 * logvar)
    
    # Sample using reparameterization trick
    eps = torch.randn_like(std)
    z = mu + eps * std
    
    kl_div = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp(), dim=-1)
    return z, kl_div
    
def forward(self, src, tgt, src_mask=None, tgt_mask=None, memory_mask=None):
    # Encode trajectory
    memory = self.encode(src, src_mask)
    
    # Compute confidence
    confidence_latent, confidence_kl = self.compute_confidence(memory)
    
    # Encode to latent space
    mu = self.encoder_mu(memory)
    logvar = self.encoder_logvar(memory)
    std = torch.exp(0.5 * logvar)
    
    # Sample using reparameterization trick
    eps = torch.randn_like(std)
    z = mu + eps * std
    
    # Compute VAE KL divergence
    vae_kl = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp(), dim=-1)
    
    # Decode
    if tgt_mask is None:
        tgt_mask = self.generate_mask(tgt.size(1), tgt.device)
    output = self.decode(tgt, z, tgt_mask, memory_mask)
    
    # Project to vocabulary space
    logits = self.output_projection(output)
    
    return {
        'logits': logits,
        'confidence_kl': confidence_kl,
        'vae_kl': vae_kl
    }

class PositionalEncoding(nn.Module):
def init(self, d_model, dropout=0.1, max_len=5000):
super().init()
self.dropout = nn.Dropout(p=dropout)

    pe = torch.zeros(max_len, d_model)
    position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
    div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
    pe[:, 0::2] = torch.sin(position * div_term)
    pe[:, 1::2] = torch.cos(position * div_term)
    pe = pe.unsqueeze(0)
    self.register_buffer('pe', pe)

def forward(self, x):
    x = x + self.pe[:, :x.size(1)]
    return self.dropout(x)