AI 驱动的链上异常检测:交易图谱构建与图神经网络的风险识别方案

AI2天前发布 beixibaobao
5 0 0

AI 驱动的链上异常检测:交易图谱构建与图神经网络的风险识别方案

一、规则引擎的终点:当攻击模式无法用正则表达

2024-2026 年的链上攻击呈现出一个明显趋势:攻击不再是单一合约级别的漏洞利用,而是跨协议、多步骤、利用短期价格窗口的"闪电贷拼装"式攻击。一条典型的攻击路径可能涉及 5-8 个协议的临时交互,在 1-2 个区块内完成资金转移、清算触发和收益提取。

传统安全监控的规则引擎在面对这类攻击时几乎完全失效。规则引擎可以检测"单笔交易中这个地址调用了 flashLoan",但它理解不了"flashLoan → swap → deposit → borrow → swap → repay → withdraw"这条六步序列的整体风险语义。它只能看到孤立的函数调用,看不到交易之间形成的攻击图谱。

图神经网络(GNN)为这个问题提供了一个结构化的分析框架。核心思路是把链上交易建模为异构图(Heterogeneous Graph),其中地址是节点,交易是边。通过 GNN 的消息传递机制,模型可以从已知的恶意地址出发,沿着交易链路传播风险信号,识别出隐藏在正常交易模式中的异常行为簇。

flowchart TD
    A[链上原始数据] --> B[交易图谱构建]
    B --> C1[地址节点]
    B --> C2[交易边]
    B --> C3[合约调用子图]
    C1 --> D[节点特征工程]
    C2 --> D
    C3 --> D
    D --> E[GNN 图神经网络]
    E --> F[节点风险嵌入]
    F --> G[异常检测分类器]
    G --> H{风险分级}
    H -->|高风险| I[实时告警 + 阻断]
    H -->|中风险| J[延时监控]
    H -->|低风险| K[归档]
    subgraph 图谱构建层
        A --> B
    end
    subgraph 推理层
        E --> F --> G --> H
    end

二、交易图谱的建模与 GNN 的消息传递原理

2.1 异构图建模

交易图谱的核心是正确建模三种节点和两种边。

节点类型

  • EOA(Externally Owned Account):普通用户钱包地址。节点特征包括:交易频率、平均交易金额、合约交互数、部署合约数、活跃天数等。
  • 合约地址:已部署的智能合约。节点特征包括:合约类型(DEX/借贷/NFT/桥)、部署时间、被调用频次、关联的攻击事件历史。
  • 资产 Token:可选节点类型,如果关注代币流向。特征包括:代币类型、持有者数量分布。

边类型

  • 交易边:EOA 发起交易调用合约。边特征包括:ETH 转账额、Gas 价格、调用函数签名、交易是否成功。
  • 内部调用边:合约间通过 CALL/DELEGATECALL/STATICCALL 产生的内部消息。边特征包括:调用深度、是否转账 ETH、调用是否 revert。

2.2 GNN 消息传递

选择 GraphSAGE 作为主干模型,它在归纳式学习任务上表现稳定,且支持大图的 mini-batch 训练。

消息传递的核心操作是:每一层 GNN,每个节点都会聚合其邻居节点的特征,结合自身特征做一个非线性变换。经过 2-3 层消息传递后,一个节点的 embedding 就包含了其 k 跳邻居的图结构信息。

在异常检测场景中,这意味着一个地址的风险评分不仅依赖于它自身的行为,还依赖于它与哪些地址发生了交互、那些地址又连接到了什么已知的恶意实体。

sequenceDiagram
    participant Indexer as 链上索引器
    participant GraphDB as 图数据库(Neo4j)
    participant Train as 离线训练管线
    participant Infer as 在线推理服务
    participant Alert as 告警系统
    Indexer->>GraphDB: 增量写入新区块交易图
    loop 每 24 小时
        Train->>GraphDB: 导出全量训练数据
        Train->>Train: 特征工程 + GraphSAGE 训练
        Train->>Infer: 更新模型权重
    end
    Indexer->>Infer: 推送最新交易
    Infer->>Infer: 在已加载的图上做推理
    Infer->>Alert: 风险评分 > 阈值 → 告警
    Alert->>Alert: 通知安全团队 / 自动化阻断

三、工程实现:从图谱构建到在线推理

3.1 交易图谱构建

"""
graph_builder.py — 链上交易图谱构建
设计考量:
- 使用 NetworkX 构建内存图谱快照,Neo4j 做持久化存储
- 节点和边的 ID 使用地址哈希而非原始地址:降低内存占用,同时避免 PII 风险
- 批量写入 Neo4j 使用 UNWIND 模式:单条 CREATE 语句在大量数据下性能极差
- 图谱按时间窗口切分:全量链上数据图太大(数亿节点),训练时使用滑动窗口截取
"""
from typing import Optional
from datetime import datetime, timezone, timedelta
from dataclasses import dataclass
import hashlib
from neo4j import GraphDatabase, Driver
def _hash_address(addr: str) -> str:
    """地址哈希,生成固定长度节点 ID"""
    return hashlib.sha256(addr.lower().encode()).hexdigest()[:16]
@dataclass
class GraphNode:
    node_type: str  # 'eoa' | 'contract' | 'token'
    address: str
    features: dict
@dataclass
class GraphEdge:
    from_addr: str
    to_addr: str
    edge_type: str  # 'transaction' | 'internal_call'
    block_number: int
    features: dict
class TransactionGraphBuilder:
    def __init__(self, neo4j_uri: str, neo4j_user: str, neo4j_password: str):
        self._driver: Driver = GraphDatabase.driver(
            neo4j_uri, auth=(neo4j_user, neo4j_password),
            # 连接池配置:图写入是瓶颈操作,需要较大的连接池
            max_connection_pool_size=20,
            connection_acquisition_timeout=30,
        )
        # 内存缓存:避免对每个新区块都做 Neo4j 查询
        self._node_cache: dict[str, GraphNode] = {}
    def add_transaction(self, tx_data: dict) -> None:
        """处理单条交易,提取节点和边"""
        from_addr = tx_data.get("from", "").lower()
        to_addr = tx_data.get("to", "").lower()
        if not from_addr or not to_addr:
            return
        block = tx_data.get("blockNumber", 0)
        # 构建 EOA 节点
        if from_addr not in self._node_cache:
            self._node_cache[from_addr] = GraphNode(
                node_type="eoa",
                address=from_addr,
                features={
                    "tx_count": 1,
                    "first_seen_block": block,
                    "last_seen_block": block,
                },
            )
        else:
            node = self._node_cache[from_addr]
            node.features["tx_count"] += 1
            node.features["last_seen_block"] = block
        # 构建合约/EOA 节点
        if to_addr not in self._node_cache:
            self._node_cache[to_addr] = GraphNode(
                node_type="contract" if tx_data.get("toIsContract") else "eoa",
                address=to_addr,
                features={"tx_count": 1, "first_seen_block": block},
            )
        # 注册交易边(暂存,批量写入时统一处理)
        self._pending_edges.append(GraphEdge(
            from_addr=from_addr,
            to_addr=to_addr,
            edge_type="transaction",
            block_number=block,
            features={
                "value_wei": int(tx_data.get("value", "0"), 16),
                "gas_price": int(tx_data.get("gasPrice", "0"), 16),
                "success": tx_data.get("status") == "0x1",
            },
        ))
    _pending_edges: list[GraphEdge] = []
    def flush_to_neo4j(self) -> int:
        """批量写入 Neo4j,返回写入的节点+边总数"""
        count = 0
        with self._driver.session() as session:
            # 批量写入节点
            if self._node_cache:
                nodes_batch = [
                    {
                        "id": _hash_address(addr),
                        "type": n.node_type,
                        "addr": n.address,
                        "tx_count": n.features.get("tx_count", 0),
                    }
                    for addr, n in self._node_cache.items()
                ]
                session.run(
                    """
                    UNWIND $batch AS node
                    MERGE (n:Address {id: node.id})
                    SET n.type = node.type,
                        n.address = node.addr,
                        n.tx_count = coalesce(n.tx_count, 0) + node.tx_count
                    """,
                    batch=nodes_batch,
                )
                count += len(nodes_batch)
            # 批量写入边
            if self._pending_edges:
                edges_batch = [
                    {
                        "from_id": _hash_address(e.from_addr),
                        "to_id": _hash_address(e.to_addr),
                        "block": e.block_number,
                        "value": e.features.get("value_wei", 0),
                    }
                    for e in self._pending_edges
                ]
                session.run(
                    """
                    UNWIND $batch AS edge
                    MATCH (f:Address {id: edge.from_id})
                    MATCH (t:Address {id: edge.to_id})
                    MERGE (f)-[r:TRANSACTED]->(t)
                    SET r.block = edge.block,
                        r.value = edge.value
                    """,
                    batch=edges_batch,
                )
                count += len(edges_batch)
            self._node_cache.clear()
            self._pending_edges.clear()
        return count
    def close(self):
        self._driver.close()

3.2 图神经网络训练

"""
gnn_trainer.py — 图神经网络训练管线
设计考量:
- 使用 PyTorch Geometric 代替原始 DGL:生态更成熟,与 PyTorch 训练管线集成更好
- 训练集正负样本比例 1:10:链上数据天然严重不平衡,恶意地址占比极低
- 加权损失函数:正样本权重是负样本的 5 倍,迫使模型关注异常模式而非盲目预测"正常"
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch_geometric.nn import SAGEConv
from torch_geometric.data import Data, DataLoader
from sklearn.metrics import precision_recall_fscore_support
class AddressRiskGNN(nn.Module):
    """
    地址风险评分 GNN 模型
    架构:2层 GraphSAGE + 1层全连接分类头
    为什么选 GraphSAGE 而非 GAT:
    - GraphSAGE 的归纳式学习能力更适合增量图场景(每天新增数万地址)
    - GAT 的注意力机制在超大图上计算成本过高,且容易过拟合到训练集的小邻域
    """
    def __init__(self, in_channels: int, hidden_channels: int = 128):
        super().__init__()
        self.conv1 = SAGEConv(in_channels, hidden_channels, aggr="mean")
        self.conv2 = SAGEConv(hidden_channels, hidden_channels, aggr="mean")
        self.classifier = nn.Sequential(
            nn.Linear(hidden_channels, 64),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(64, 2),  # 二分类:正常 / 异常
        )
    def forward(self, x, edge_index):
        x = F.relu(self.conv1(x, edge_index))
        x = F.dropout(x, p=0.2, training=self.training)
        x = F.relu(self.conv2(x, edge_index))
        return self.classifier(x)
def train_epoch(
    model: AddressRiskGNN,
    loader: DataLoader,
    optimizer: torch.optim.Optimizer,
    device: torch.device,
) -> float:
    model.train()
    total_loss = 0.0
    # 正样本权重设为 5,负样本 1:缓解天然类别不平衡
    pos_weight = torch.tensor([5.0], device=device)
    criterion = nn.CrossEntropyLoss(weight=torch.tensor([1.0, 5.0], device=device))
    for batch in loader:
        batch = batch.to(device)
        optimizer.zero_grad()
        out = model(batch.x, batch.edge_index)
        loss = criterion(out, batch.y)
        loss.backward()
        # 梯度裁剪:防止恶意构造的图谱数据导致的梯度爆炸
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
        optimizer.step()
        total_loss += loss.item()
    return total_loss / len(loader)
def evaluate(model: AddressRiskGNN, data: Data, device: torch.device) -> dict:
    model.eval()
    with torch.no_grad():
        out = model(data.x.to(device), data.edge_index.to(device))
        pred = out.argmax(dim=1).cpu()
        precision, recall, f1, _ = precision_recall_fscore_support(
            data.y.cpu(), pred, average="binary", pos_label=1
        )
    return {"precision": precision, "recall": recall, "f1": f1}
# 训练主循环
# 数据准备从 Neo4j 导出为 PyG Data 对象
# 特征包括:节点度数、交易频率、交易金额分布、合约交互多样性等 16 维特征
# 标签来自已知的安全事件标注(如被标记的钓鱼地址、攻击合约等)

3.3 在线推理服务

"""
inference_server.py — 在线风险推理服务
设计考量:
- 模型加载为单例:避免每个请求重新加载模型权重
- 推理使用 torch.inference_mode() 而非 torch.no_grad():关闭 autograd 引擎降低内存
- 图更新使用增量策略:新区块交易追加到内存子图中,无需重建全图
- 风险评分阈值分层:>0.8 即时告警,0.5-0.8 加入观察列表
"""
import torch
from torch_geometric.data import Data
import asyncio
from collections import defaultdict
class RiskInferenceServer:
    HIGH_RISK_THRESHOLD = 0.8
    MEDIUM_RISK_THRESHOLD = 0.5
    def __init__(self, model_path: str, device: str = "cpu"):
        self._device = torch.device(device)
        self._model: AddressRiskGNN | None = None
        self._model_path = model_path
        # 内存子图:不存储全链数据,仅维护活跃窗口(最近 7 天)
        self._active_graph: dict = {
            "nodes": {},        # node_id -> feature_vec
            "edges": [],        # [(src, dst)]
            "node_to_idx": {},  # address -> contiguous index
        }
        self._lock = asyncio.Lock()  # 图更新与推理之间的并发控制
    async def load_model(self):
        self._model = AddressRiskGNN(in_channels=16)
        self._model.load_state_dict(torch.load(self._model_path, map_location=self._device))
        self._model.to(self._device)
        self._model.eval()
    async def add_transactions(self, txs: list[dict]):
        """增量更新内存交易图"""
        async with self._lock:
            for tx in txs:
                from_addr = tx["from"]
                to_addr = tx["to"]
                for addr in (from_addr, to_addr):
                    if addr not in self._active_graph["node_to_idx"]:
                        # 初始化节点特征向量
                        idx = len(self._active_graph["node_to_idx"])
                        self._active_graph["node_to_idx"][addr] = idx
                        self._active_graph["nodes"][idx] = [0.0] * 16
                src = self._active_graph["node_to_idx"][from_addr]
                dst = self._active_graph["node_to_idx"][to_addr]
                self._active_graph["edges"].append((src, dst))
    async def predict_risks(self, addresses: list[str]) -> dict[str, float]:
        """对指定地址列表做风险评分"""
        if self._model is None:
            await self.load_model()
        async with self._lock:
            # 构建 PyG Data 对象
            num_nodes = len(self._active_graph["node_to_idx"])
            x = torch.zeros((num_nodes, 16), device=self._device)
            for idx, feat in self._active_graph["nodes"].items():
                x[idx] = torch.tensor(feat, device=self._device)
            edge_index = torch.tensor(
                self._active_graph["edges"], device=self._device
            ).t().contiguous()
            data = Data(x=x, edge_index=edge_index)
        with torch.inference_mode():
            logits = self._model(data.x, data.edge_index)
            probs = torch.softmax(logits, dim=1)
            results = {}
            for addr in addresses:
                idx = self._active_graph["node_to_idx"].get(addr)
                if idx is not None:
                    # prob[1] 是异常类的概率
                    results[addr] = float(probs[idx][1])
                else:
                    results[addr] = -1.0  # 未知地址
        return results

四、边界分析

数据标注是最大的瓶颈。GNN 的异常检测效果高度依赖标注数据的质量。链上恶意地址的标签来源分散:Etherscan 标签、慢雾攻击库、OpenChain 的 AML 数据、以及安全团队的内部情报。这些标签覆盖面和更新速度有限,模型对新型攻击模式的泛化能力取决于标注库的及时扩展。

图的时效性问题。链上交易图是高度动态的。训练时使用的图结构和推理时的图结构存在时间差,新出现的合约交互模式在训练图中不存在,GNN 消息传递无法利用这些新连接。解决方案是采用频繁重训练策略(如每 6 小时增量训练),但这需要持续的计算资源投入。

特征工程的领域依赖。节点特征的构造直接影响 GNN 的表达能力。如果将"合约交互多样性"这类高级特征写成硬编码的统计指标,模型会错过未知的特征组合模式。更好的做法是使用 Node2Vec 或 Graph Autoencoder 做图结构的无监督预训练,让模型自动学习节点在交易图中的拓扑角色。

不适用场景

  • 只有几笔交易的全新地址:没有足够图谱结构供 GNN 推理,退化为随机猜测
  • 非交易型链上行为:如 NFT Mint 的签名验证、零知识证明验证等不产生交易图结构的操作
  • 需要逐笔交易实时拦截的极高吞吐场景:GNN 推理延迟(毫秒级起步)不适合交易所级别的吞吐

五、总结

维度 要点
核心思路 将链上交易建模为异构图,通过 GNN 消息传递传播风险信号
图建模 EOA / 合约 / Token 三类节点,交易边 + 内部调用边,特征含交易频率和金额分布
模型选型 GraphSAGE 适合大规模增量图,2层卷积 + MLP 分类头
工程管线 Neo4j 持久化 + PyG 训练 + 增量内存图推理
两大瓶颈 标注数据覆盖面和时效性、模型对新交互模式的泛化能力
适用边界 适合批量异常检测和大规模地址风险评分,不适合单笔实时拦截和全新地址评估
© 版权声明

相关文章