提出单步邻居聚合方法，显著提升知识图谱更新效率与内存利用率。
FuseSampleAgg: One-Pass Neighborhood Estimation for Budgeted Knowledge-Graph Refresh and Validation
- 融合采样与平均聚合操作，避免中间数据结构构建。
- 相比基线提速2.24至3.48倍，显存峰值降低最高160倍。
- 支持可复现采样，适合严格预算下的图谱刷新与验证场景。
网络与网络安全领域中，运营级知识图谱（KG）管道需在严格的时间、内存和审计预算下刷新嵌入表示，尤其在人工筛选数据流和大模型辅助抽取加速更新的背景下。当前小批量训练中的关键开销在于邻居上下文估计：无放回均匀采样后进行均值聚合。主流框架通过采样子图构建和中间特征收集实现该过程，导致内核启动频繁、内存分配压力大及瞬时显存激增。本文提出单步邻居估计（One-Pass Neighborhood Estimation），一种融合的PyTorch CUDA算子，直接采样邻居并输出其均值，无需显式构造子图块，同时保持与GraphSAGE均值语义一致。该方法支持种子控制采样及可选的索引重播，便于可复现验证与回归测试。在大规模图的小批量工作负载中，其FP32端到端步骤延迟相比调优后的DGL基线提升2.24x–3.48x，瞬时显存消耗减少最多160倍。在OGB KG补全基准（如WikiKG2和BioKG）上，该方法在保持排名质量（在种子波动范围内）的同时，显著降低步骤耗时与峰值显存，提升受限预算下知识图谱刷新的时效性。
原文摘要 · Abstract (English)
Operational knowledge-graph (KG) pipelines in networking and cybersecurity increasingly need to refresh embeddings under strict time, memory, and audit budgets, especially as curated feeds and LLM-assisted extraction accelerate KG updates. A recurring per-step cost in mini-batch KG learning is neighborhood-context estimation: uniform neighbor sampling without replacement followed by mean aggregation. Common frameworks implement this estimator through sampled-subgraph materialization and intermediate feature gathers, adding kernel launches, allocator pressure, and transient memory spikes. We present One-Pass Neighborhood Estimation, a fused PyTorch CUDA operator that samples neighbors and directly emits the sampled-neighborhood mean, avoiding explicit block construction while preserving GraphSAGE-mean semantics for the same sampled neighbor IDs. It supports seed-controlled sampling and optional saved-index replay for reproducible validation and regression testing. Across large-graph mini-batch workloads, it improves FP32 end-to-end step latency by 2.24x-3.48x over tuned DGL baselines and reduces transient GPU memory by up to 160x in our measurements. On OGB KG completion benchmarks such as WikiKG2 and BioKG, it reduces step time and peak VRAM while matching ranking quality within seed variability, improving time-to-quality for budgeted KG refresh.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。