Axon v3-Pre 技术报告

摘要与核心结论

Axon v3-Pre（内部迭代代号 v2.6Exp）完成了对上一代 Axon 2.5Exp 架构的全面重构。系统从传统的“手工特征工程 + LightGBM + 规则路由”树模型体系，演进为“端到端原始字节注意力机制 + 结构化特征融合 + 级联梯度提升纠错网络 + 零依赖 Native C++ 部署”的深度学习混合架构。

核心演进与性能指标

架构范式转换：摒弃 2.5Exp 时代受限于预设特征抽取器的工作流，采用 4096 / 65536 字节原始字节流直接输入注意力网络学习隐式表示，叠加 331 维内容特征级联纠错层（Stage-2 HGB），实现特征捕获与决策矫正解耦。
显存与计算优化：针对 65536 序列长度在传统 Self-Attention 下导致单层注意力矩阵膨胀至约 2 TB（Batch Size 64）的问题，设计自研 MHDSRA2（Multi-Head Dynamic Sparse Recurrent Attention）流式注意力网络，将全序列训练峰值显存降至 2.24 GB（Chunk Size = 512, FP32）。
基底模型表现：经过 20 个 Epochs（27.2 小时）训练，在 147,796 个独立测试样本（包含 3,482 个 UPX 加壳白名单文件）上，基底深度模型实现 Test F1 0.9795、Accuracy 0.9679、AUC 0.9921。
Stage-2 纠错提升：叠加 331 维内容特征与 3-Seed HistGradientBoosting（HGB）纠错层后，Test F1 提升至 0.99341，全局分类错误数从 4,742 降至 1,511（假阳性 FP 下降 75.8%，假阴性 FN 下降 51.9%）；UPX 加壳白文件误报数由 236 降至 52。
原生部署交付：模型收敛为单一 ONNX 结构（输入包含 byte_seq[1, 4096]、pe[1, 1500]、stat[1, 49]，输出 logits[1, 2]），通过无外部依赖的动态链接库 axon_onnx_predict.dll（__cdecl 接口，18 个导出函数，96 字节 KvdConfig 结构体）交付。脱离 Python/PyTorch 运行环境，单文件平均推理延迟 ~355 ms。
现实边界约束：基于静态特征与含噪标签数据集，绝对零误报与 Recall $> 99\%$ 无法同时达成。评估表明基座模型中包含 221 个置信度 $> 90\%$ 的假阴性（FN）样本，主因为数据源标签噪声，而非模型表征缺陷。

1. 架构演进背景与遗留系统（v2.5Exp）局限性

在重构之前，Axon 2.5Exp 采用了“手工特征提取 + LightGBM 树模型 + 静态规则路由”的经典工程架构。随着恶意样本复杂度的提升与对抗技术的演进，该架构逐步显露出若干难以通过局部调优解决的技术瓶颈：

1.1 静态特征空间的表征瓶颈

特征表达能力受限：遗留系统的特征抽取过程脱离了下游任务的梯度反馈，高维二进制数据在进入模型前即被大幅度静态压缩（主要依赖轻量级哈希桶与基础 PE 结构字段）。这种缺乏上下文感知的特征设计，造成了深层结构信息的不可逆损失。
规则路由的泛化局限：原有的专家路由机制依赖硬编码的规则阈值（如加壳标志、特定节区比例）。在面对新型加壳变种或缺乏显式特征字符串的对抗样本时，硬编码规则难以自适应动态变化，容易导致路由门控退化。

1.2 树模型对字节特征的拟合瓶颈

在遗留系统的特征重要性评估中，用于刻画原始字节分布的特征字段排名普遍靠后（无一进入前 50）。这表明“人工设计预压缩规则，再由树模型拟合”的传统范式，无法有效捕获二进制文件内部的高频局部字节模式。为此，v3-Pre 决定放弃预判机制，转向端到端原始字节注意力表示学习，由神经网络直接从原始字节流中提取高阶特征。

2. v3-Pre 模型架构设计

v3-Pre 的深度神经网络采用三路异构输入 + 特征拼接（Concat）融合的拓扑结构。

                                ┌─────────────────────────────────────────────────────┐
 原始字节流 ─► byte_seq(4096/65536) │ ByteEmbedding(256→128)                              │
                                │  └─► 正弦位置编码 (max_len=65536)                   │
                                │  └─► input_proj(128→128)                            │
                                │  └─► [MultiHeadDSRA2 ×2, chunk_pool=last] ──► 128-dim│
 PE 结构特征 ─► legacy_dynamic(1500) ┼─► PEFeatureProjector(1500→256→128) ──────► 128-dim│
 字节统计特征 ─► stat(49)          └─► stat_projector(49→128) ──────────────► 128-dim│
                                └───────────────────────────────┬─────────────────────┘
                                                                ▼
                                                    Concat 融合 ──► 384-dim
                                                                ▼
                                                 LayerNorm(384) ──► Linear(64) ──► GELU ──► Linear(2)

2.1 参数量结构分解

根据模型权重文件 best_model_739k.pt 的 state_dict 实测统计，各模块参数配置如下：

模块名称	参数量	结构配置说明
ByteEmbedding	32,768	词表大小 256 $\times$ 隐层维度 128
input_proj	16,512	Linear(128 $\to$ 128) + GELU 激活
MultiHeadDSRA2 ( $\times 2$ )	196,876	单层 98,438（QKV 投影 49,152 + Out 投影 16,384 + Slot 初始化 16,384 + 门控单元 132）
PEFeatureProjector	417,664	Linear(1500 $\to$ 256) + LayerNorm + GELU + Linear(256 $\to$ 128)
stat_projector	6,656	Linear(49 $\to$ 128)
classifier	25,538	LayerNorm(384) + Linear(384 $\to$ 64) + GELU + Linear(64 $\to$ 2)
总计可训练参数	696,014	约 69.6 万可训练参数

注：权重文件中的 numel 统计值（9,183,060）包含 8,388,608 个不可训练的正弦位置编码 Buffer（ $65536 \times 128$ ）及别名引用。模型核心可训练参数仅 69.6 万，具备极高的计算效率与 CPU 运行友好度。

2.2 融合机制选择策略

在系统开发中评估了 add、gated、residual_stat_gate 以及 cross-attention 等多种融合拓扑（model.py:561-584）。实验表明，在三路输入隐层维度统一对齐至 128 的前提下，直接采用 concat（共 384 维）能够以最简单的方式实现信息无损传递，避免了复杂注意力融合带来的额外超参数引入与梯度不稳定风险。

3. MHDSRA2 流式注意力网络计算机制

为解决长字节序列在端到端训练时的显存爆炸问题，设计了 MHDSRA2（Multi-Head Dynamic Sparse Recurrent Attention）网络。

3.1 复杂度对比分析

维度	标准多头自注意力 (MHA)	MHDSRA2 流式注意力
每层注意力得分计算	$\mathcal{O}(B \cdot H \cdot T^2)$	$\mathcal{O}(B \cdot H \cdot C \cdot (K + W + \text{topk}))$ ，其中 $C=512$ (Chunk Size)
每层缓存开销	$\mathcal{O}(B \cdot H \cdot T \cdot d)$ ，随序列长度 $T$ 线性膨胀	$\mathcal{O}(B \cdot H \cdot (K + W) \cdot d)$ ，固定开销，与 $T$ 解耦
65536 序列 / Batch 64 显存	单层 Logits 占用 $\approx 65536^2 \times 64 \times 4 \times 2 \text{ Bytes} \approx \mathbf{2\text{ TB}}$	逐 Chunk 分块流式计算，实测峰值仅为 2.24 GB
跨 Chunk 状态传递	无（依赖全全局矩阵）	依赖 Slot $K/V$ 状态矩阵 $[B, H, 128, d]$ 、Local Cache ( $\le 256$ ) 及更新门控

3.2 算法处理流程

设 Chunk 大小为 $C=512$ ，槽位数量 $K=128$ ，对于输入的长字节序列，按以下三步分块处理：

Slot Read（槽位读取）：计算当前 Chunk 的 Query 与全局槽位 Key 的关联性 $Q \cdot K_{\text{slots}}^T \in \mathbb{R}^{B \times H \times 512 \times 128}$ ，取 read_topk = 8 执行稀疏 Softmax 聚合 Slot Value。
Local Attention（局部因果注意力）：在涵盖 Local Cache（ $\le 256$ Token）与当前 Chunk 的滑动窗口内执行因果 SDPA（Scaled Dot-Product Attention）计算。
Slot Write（槽位写更新）：当前 Chunk 的每 Token 按 write_topk = 4 路由至指定 Slot，利用 scatter_add 累加，并通过包含 Age/Usage/Confidence 衰减机制的 Gated Forget 单元更新全局槽位状态。

3.3 显存占用与 Chunk Size 关系测试

固定 Batch Size = 64，测试不同 Chunk 配置下的训练峰值显存：

Chunk Size ( $C$ )	训练峰值显存	运行说明
512	2.24 GB	生产与训练标准配置
1024	4.23 GB	显存膨胀约 1.9 倍
2048	8.25 GB	在 8GB 显存设备上触发 OOM 异常

其理论显存计算表达式为：

\text{Memory}_{\text{DSRA}} \approx \mathcal{O}(B \cdot H \cdot C \cdot K \cdot d_{\text{head}}) + \mathcal{O}(2 \cdot B \cdot H \cdot W \cdot d_{\text{head}})

通过用固定大小的 Slot 存储压缩记忆，成功将空间复杂度从 $\mathcal{O}(T^2)$ 压降至 $\mathcal{O}(T)$ 线性级。

3.4 精度约束与 Diversity Loss

数值稳定性约束：DSRA 模块在 FP16 混合精度训练下，由于 Chunk 内部与 Slot 状态迭代累加，极易引发前向传播数值溢出（导致 NaN，见 train_739k_full.py:234）。实验验证 BF16 虽无 NaN 溢出，但由于 Python 层面的 Chunk 循环调度无法填满 GPU 算力，相比纯 FP32 无明显吞吐提升。因此，训练过程强制采用全单精度（Pure FP32）。
Slot 坍塌正规化：为防止 128 个全局 Slot 退化为同质化表示，训练损失中引入了对 Slot Key 矩阵的重叠惩罚：

\mathcal{L}_{\text{div}} = \lambda \cdot \frac{1}{B \cdot H} \sum_{b,h} \left\Vert{} \text{Gram}\left(K_{\text{slots}}^{(b,h)}\right) - I \right\Vert{}_F^2

其中设置 $\lambda = 0.03$ 。计算上直接对 Gram 矩阵求逐元素平方和，避免了高代价的 $O(K^3)$ 矩阵求逆运算。

4. 结构化特征提取与对齐（legacy_dynamic）

在三路输入中，legacy_dynamic 模块负责处理传统的 PE 头与结构化信息（extractor.py:706-855）。

[0 : 18]         固定字段：文件头属性与安全标志
[18 : 18+3N]     动态 Section 属性：包含 exec / write / read 标志 (N 为 Section 数量)
[18+3N : 47+3N]  聚合计算特征：信息熵、导入/导出表统计、尾部 Overlay 数据、API 类别映射等
[47+3N : 1500]   零值填充（Zero-Padding）

该设计的优势在于，将 Section 维度的访问权限直接转化为动态对齐特征，保留了结构体上下文；缺点是特征索引位置依赖于 $N$ 的变化，导致静态特征名列表与物理偏移不可直接一一映射。此外，审查发现历史代码中的 idx16（has_signature）依赖未定义的 PE 库属性，导致该特征列恒为 0。在 v3-Pre 的规范化特征映射中，该无用列已被清理。

5. 训练策略与平台特定并发优化

5.1 训练超参数超细配置

参数项	配置值
损失函数	Focal Cross-Entropy ( $\gamma = 1.0, \alpha = 0.55$ ) + Label Smoothing ( $0.03$ ) + Diversity Loss ( $0.03$ )
优化器	AdamW ( $\text{lr} = 8 \times 10^{-5}, \text{weight\_decay} = 1 \times 10^{-5}, \beta = (0.9, 0.999)$ )
学习率调度	3 Epochs 线性 Warmup（ $1 \times 10^{-6}$ ） $\to$ Cosine Annealing $\to$ $\text{Min\_LR} = 1.67 \times 10^{-6}$
梯度裁剪	Gradient Norm Standard $0.75$
数值精度	Pure FP32
字节输入截断	抽取 65536 字节，训练时截断前 4096 字节

5.2 截断策略权衡

在前 4096 字节截断配置下，DSRA 迭代步数由 128 降至 8，单步训练时间由 $9.2\text{ s}$ 缩短至 $0.5\text{ s}$ （吞吐提升 18 倍）。由于 PE 文件头与入口点（Entry Point）上下文通常完整覆盖在前 4KB 空间内，深层字节信息的缺失将由 Stage-2 的全局内容特征进行弥补。

5.3 Windows 并发调度故障修复

在 Windows 平台启用 PyTorch 多进程 DataLoader（8 个 Workers）时，各子进程在初始化 OpenBLAS 运行时会默认尝试锁死所有 CPU 线程（32 线程），导致 Windows 线程栈资源耗尽而产生挂起异常。必须在导入 numpy / torch 模块之前，注入环境变量控制：

import os
os.environ.setdefault("OMP_NUM_THREADS", "1")
os.environ.setdefault("OPENBLAS_NUM_THREADS", "1")

5.4 训练收敛收尾记录

训练共执行 20 个 Epochs（耗时 27.23 小时）。基底模型获得 Test F1 0.9795、Precision 0.9723、Recall 0.9867、AUC 0.9921（包含 3,217 个假阳性 FP 与 1,525 个假阴性 FN）。整个训练过程中早停机制（Patience = 8）未被触发，模型指标保持稳步上升。

6. Stage-2 级联纠错架构（HistGradientBoosting）

6.1 纠错逻辑与特征设计

基底深度模型虽然整体校准良好（期望校准误差 ECE = 0.001），但由于训练集中良性样本有 89% 为 DLL 文件，导致模型对大型独立良性 EXE（如安装包、加载器）易产生过自信的误判。Stage-2 引入一个无需重训深度基座的轻量级级联纠错层。

Base Model (739k) ──► p_malicious ──► [p, p², |p-0.5|, log(p), log(1-p), logit(p)] (6 维)
                                                                │
原始二进制文件   ──► content_pe_v1 (100 维)                     │
                 ──► content_pe_v2 (182 维) ─────────────────────┼─► 331 维特征矩阵
                 ──► content_string (43 维)                     │
                                                                ▼
                                                3-Seed HistGradientBoosting (HGB)
                                                                ▼
                                                    概率均值 ──► 阈值 0.55 决策

基底概率衍生特征（6 维）：提取 Logit 空间与概率空间中的非线性变形特征。
content_pe_v1（100 维）：包含文件头统计、11 个数据目录属性、API 分类占比、导出表及 Overlay 组合特征。
content_pe_v2（182 维）：涵盖 32 个高频 DLL 导入依赖、16 组 API 敏感行为特征、21 项 Section/入口点特征。
content_string（43 维）：针对 ASCII/UTF-16 编码的连续游程、正则匹配（URL/IPv4/注册表/路径）及 14 类敏感语义模式进行提取。

6.2 级联纠错消融实验

在 147,796 个独立测试样本上测试 Stage-2 的性能收益：

评估配置	决策阈值	Test F1	错分样本总数	假阳性 (FP)	假阴性 (FN)
纯 Base Model Logits	0.49	0.97999	4,597	3,072	1,525
+ content_pe_v1 (100 维)	0.55	0.99282	1,647	842	805
全量 331 维特征组合	0.55	0.99341	1,511	778	733

在难度极高的 3,482 个 UPX 加壳白名单文件测试子集中，基座模型的误报数由 236 降低至 52。剔除加壳白名单干扰后，全局 Test F1 超过 0.995。

7. Native C++ 部署与端到端集成

7.1 ONNX 导出契约

生产环境模型绑定为单一 ONNX 结构（dist/axon_739k_onnx_final_20260808）：

张量名称	维度 (Shape)	数据类型 (DataType)	语义说明
`byte_seq`	`[1, 4096]`	`INT64`	截断的原始前 4KB 字节序列
`pe_features`	`[1, 1500]`	`FLOAT32`	`legacy_dynamic` 结构化特征
`stat_features`	`[1, 49]`	`FLOAT32`	字节熵与全量统计特征
`logits`	`[1, 2]`	`FLOAT32`	分类 Logits 输出

7.2 Native C++ 动态链接库设计

底层推理引擎通过 C++ 实现（axon_onnx_predict.dll），完全移除 Python、PyTorch 及 Scikit-Learn 运行时依赖。

ABI 与配置定义：采用 __cdecl 调用约定，提供 18 个 C 风格 API 函数。结构体 KvdConfig 在 x64 架构下精确定向为 96 字节，在 Rust (size_of::<KvdConfig>() == 96) 与 Node.js (koffi.sizeof == 96) 中进行 ABI 校验。
运行时环境隔离：针对 Windows 环境中 System32 下可能存在的旧版本 onnxruntime.dll（如 1.17.1）抢占符号问题，DLL 内部使用 SetDllDirectory 显式限定加载同级 bin/ 目录下的 ONNX Runtime 1.24.4 版本。
字符集编译修正：解决 MSVC 编译器在 GBK 编码下解析 UTF-8 中文注释时将续行符 \ 误判为代码换行，从而导致函数签名损坏的问题。
推理耗时预算：在预热完成后，包含文件读取、特征抽取、ONNX 节点推理的全流程平均耗时为 ~355 ms（ $< 500\text{ ms}$ ）。

8. 综合评估与数据噪声分析

针对评估指标中的基准测试差异进行说明：遗留系统 v2.5Exp 在历史评估中宣称 Acc $0.9938$ ，但其评估集中恶意样本占比仅约 $60\%$ 且完全不包含 UPX 加壳白文件。v3-Pre 测试集包含了大量对抗性强、结构复杂的样本（114,693 个恶意样本与 3,482 个 UPX 加壳良性文件）。

在全局分析中，针对错分的样本进行离线归因发现：

静态特征表达下不存在能够实现 FP = 0 且 Recall > 99% 的单点阈值。
在 Stage-2 预测的假阴性（FN）样本中，有 221 个样本在基座模型中的预测恶意概率超过 $90\%$ 。针对这部分样本进行人工逆向复核表明，大部分属于黑产样本在标注库中的误标白数据（Label Noise）。在排除误标噪声后，系统真实召回率估计可达 0.9955。

9. 架构演进中的关键技术洞察

结构化特征虚高问题：遗留系统宣称的 1500 维特征，在 C++ 提取层实际仅落地 350 维，且关键路由字段存在静默丢失，架构设计需严格以代码实现为准。
流式注意力的数值精度约束：Chunked 注意力网络由于状态跨 Chunk 累积，FP16 会直接引发数值上溢（Overflow），工业落地必须使用 Pure FP32 或针对性梯度截断。
模型参数量误区：深度模型中的位置编码矩阵（Positional Embedding Buffer）易被误计入模型可训练参数，需剔除冻结 Buffer 评估真实计算负荷。
特征对齐防线：C++ 推理端与 Python 训练端特征提取器实现必须逐位（Bit-exact）对齐。本系统中对齐误差降低至 $0.00000$ 级别。
动态拓扑的内存安全性：在 C++ 中重新实现动态 PE 特征提取时，需防止越界读取及非标准 PE 头导致的内存崩溃。
动态链接库加载防御：Windows 系统 System32 路径下的全局 DLL 易造成依赖劫持，必须在 Native 层显式锁定加载路径。
多进程并发线程竞争：Windows 平台下 PyTorch 搭配 OpenBLAS 时，必须在导入底层数学库前锁定线程数，防止线程栈耗尽。
分段互补特征设计：4096 字节截断能够换取 18 倍的训练吞吐收益，深层字节信息的缺失可由非序列化的结构与字符串特征层有效弥补。
级联树模型的部署化：将 Scikit-Learn 训练的 HistGradientBoosting 决策树参数导出为纯数组 JSON 格式，由 C++ 原生解析运行，可避免引入额外的 C++ 运行时依赖。
验证集与测试集隔离：Stage-2 的决策阈值（0.55）严格在 Validation Set 上通过网格搜索获取，Test Set 仅执行单次推理评估，确保无测试集数据泄漏。

10. 系统演进计划

良性样本扩充重训收尾：完成目前进行中的良性样本扩充重训任务（测试集包含 813,098 样本，良性样本占比提升至 29.5%），收敛后更新 Stage-2 纠错树。
Stage-2Native 集成：将导出为 JSON 数组的 Stage-2 HGB 权重接进 axon_onnx_predict.dll 的 stage2_model_json_path 接口，实现完全端到端交付。
数据集标签去噪：针对筛选出的 221 个高置信度假阴性（FN）样本建立自动化复核机制，清洗标注噪声。
长上下文检索分支验证：针对超过 65536 字节的特长文件，评估开启 PagedExactMemory 分页检索机制对分类精度的边际收益。