GreenWinters/PyTIE

# PyTIE ## PyTorch 技术推断引擎 (TIE)：训练方法与实现细节 技术推断引擎 (TIE) 是一个机器学习系统，旨在帮助网络防御者通过根据先前观察到的行为预测相关的 [MITRE ATT&CK](https://attack.mitre.org/) 技术来预判对手的下一步行动。TIE 利用了目前最大的公共数据集之一，将 CTI 报告与 ATT&CK 技术关联起来，包含 6,236 份 CTI 报告中的 43,899 条技术观察数据，覆盖了 96% 的 ATT&CK Enterprise v15.0。 ### 功能概述 `tie_model.py` 的核心功能是协调多个推荐模型的训练、评估和比较，用于技术推断。它提供了用于实验选择的 CLI，自动管理设备 (CPU/GPU) 和超参数搜索，并输出训练好的模型工件和详细指标。这支持了 TIE 的目标，即帮助防御者构建对手行动的完整图景并预测未来技术。 ### 训练方法论 #### 输入 - **数据集：** - `combined_dataset_full_frequency.json`：主要报告-技术矩阵（43,899 次观察，6,236 份报告）。 - `enterprise-attack.json`：MITRE ATT&CK 技术映射。 - **结构：** - 数据使用 `ReportTechniqueMatrixBuilder` 分为训练集、验证集和测试集。 - 所有拆分均转换为 PyTorch tensors 以实现高效计算和设备兼容性。 #### 模型训练 - **支持的模型：** - WALS (加权交替最小二乘法) - **描述：** 矩阵分解算法，使用加权最小二乘法交替优化用户和 item embeddings。专为隐式反馈数据设计，能有效处理大型稀疏矩阵。 - **优势：** 对大型数据集扩展性良好，对缺失数据鲁棒，支持 GPU 加速，embeddings 具有可解释性。 - **劣势：** 假设线性关系，可能在高度非线性或复杂交互中表现不佳，对超参数敏感。 - **预测目标：** 基于历史报告-技术关联，预测对手使用特定 ATT&CK 技术的可能性。 - **输出解释：** 分数越高表示预测的相关性越大；推荐 top-k 技术以供进一步调查。 - BPR (贝叶斯个性化排序) - **描述：** 成对排序算法，学习将观察到的（正）交互排序高于未观察到的（负）交互。优化排序质量而非绝对预测。 - **优势：** 直接优化排序指标，对推荐任务有效，能处理隐式反馈。 - **劣势：** 可能需要仔细的负采样，在超大数据集上收敛较慢，不如 WALS 易于解释。 - **预测目标：** 根据预测的相关性对每个报告/实体的 ATT&CK 技术进行排序。 - **输出解释：** 预测排名最高的技术最有可能被对手使用。 - TopItems (基于频率的基线) - **描述：** 简单的基线模型，根据技术在数据集中的整体频率进行推荐，忽略报告特定的上下文。 - **优势：** 快速，易于实现，提供比较基准。 - **劣势：** 忽略上下文和个性化，可能会漏掉罕见但重要的技术。 - **预测目标：** 最常观察到的 ATT&CK 技术。 - **输出解释：** 推荐最常见的 top-k 技术。 - Factorization (矩阵分解) - **描述：** 通过分解报告-技术矩阵来学习报告和技术的潜在 embeddings。捕获潜在的结构和关联。 - **优势：** 捕获潜在结构，embeddings 可解释，对中等规模数据集高效。 - **劣势：** 假设线性关系，对高度非线性数据效果较差，对稀疏性敏感。 - **预测目标：** 预测每个报告/实体的技术相关性。 - **输出解释：** 分数越高表示预测的关联性越强；推荐 top-k 技术。 - Implicit BPR (通过 implicit.gpu) - **描述：** 通过 `implicit.gpu.bpr` 后端实现的 GPU 加速贝叶斯个性化排序。直接在 CUDA tensors 上训练，并通过 conda-forge GPU 构建版高效运行 (`conda install -c conda-forge implicit implicit-proc=*=gpu`)。 - **优势：** 高度可扩展的排序，在 GPU 上运行，对缺失数据鲁棒。 - **劣势：** 需要仔细调优，不如标准矩阵分解易于解释，对数据分布敏感。 - **预测目标：** 对每个报告/实体的 ATT&CK 技术进行排序。 - **输出解释：** 预测排名最高的技术最有可能被对手使用。 - Implicit WALS (已弃用的 CPU 包装器) - **描述：** 仅限 CPU 的 `implicit.als` 实现的兼容性垫片。保留是为了向后兼容，但会发出弃用警告，当需要 GPU 加速时应避免使用。 - **优势：** 符合遗留行为。 - **劣势：** 仅限 CPU，速度明显较慢，无 GPU 加速，未来适应性较差。 - Implicit ALS (GPU/CPU 回退) - **描述：** 使用 `implicit.gpu.als` 的交替最小二乘法推荐器（当 CUDA 扩展可用时），并自动回退到 CPU ALS。 - **优势：** 在稀疏隐式反馈数据上高效，支持 GPU 加速，鲁棒的回退行为。 - **劣势：** 后端/库兼容性可能会影响运行时，并且 CPU 回退在大型运行中明显较慢。 - **预测目标：** 从学习到的潜在因子预测并对每个报告/实体的 ATT&CK 技术相关性进行排序。 - **输出解释：** 分数越高表示预测的相关性越强；优先考虑 top-k 技术。 - Implicit MatrixFactorizationBase (GPU/CPU 回退) - **描述：** `implicit` 矩阵分解基础模型的适配器，优先使用 GPU 实现（如果可用），并在需要时回退到 CPU。 - **优势：** 灵活的后端选择，一致的推荐器接口，对缺失 CUDA 扩展具有弹性。 - **劣势：** 性能取决于后端和环境设置；回退路径可能会增加运行时间。 - **预测目标：** 学习用户/报告和 item/技术潜在表示，以对可能的技术关联进行评分。 - **输出解释：** 分数用于排序；排名最高的技术被推荐为可能的下一步行动。 - **设备管理：** - TIE 在运行时检测 CUDA 可用性并自动选择执行设备。 - 对于基于 PyTorch 的推荐器，模型参数和 tensors 被显式移动到选定的设备。 - 对于基于 `implicit` 的推荐器，在可用且兼容时使用 GPU 后端；否则，代码回退到 CPU 实现。 - 感知设备的训练意味着相同的实验接口可以在异构环境中运行，同时保留可比的输出/指标。 - **超参数搜索：** - 自动对 embedding 维度和超参数进行搜索，并根据验证选择最佳配置。 - **训练循环：** - 对于每个配置，模型都会进行训练、验证和测试。根据性能指标选择最佳超参数。 #### 输出 - **模型工件：** - 训练好的模型保存为带时间戳文件名的 PyTorch `.pt` 文件。 - **指标：** - 性能指标保存为 JSON 文件，包括 precision、recall、NDCG 和 MSE。 - **CSV 结果：** - Embedding 维度扫描和其他实验将结果输出为 CSV 文件。 #### 测量/指标 - **Precision@k：** top-k 预测中相关技术的比例。 - **Recall@k：** 在所有可能的相关技术中检索到的相关技术的比例。 - **NDCG@k：** 评估排序质量，奖励相关技术的正确排序。 - **MSE：** 用于基于回归的模型以评估拟合质量。 - **最佳模型选择：** 模型比较并根据 NDCG@20 选择最佳模型。 ### 与原始仓库的区别 与原始 MITRE Engenuity Center for Threat-Informed Defense 仓库相比，此版本的 `tie_model.py` 包括： - **感知设备的训练：** 如果可用，所有模型和 tensors 都会移动到 GPU。 - **扩展的模型支持：** 集成了额外的推荐模型。 - **自动超参数搜索：** 多个配置并行运行，根据验证指标选择最佳。 - **鲁棒的输出管理：** 模型工件和指标保存时带有时间戳文件名。 - **全面的 CLI：** 实验选择，最佳模型返回，以及 JSON/CSV 格式的输出。 - **增强的指标和报告：** 每个实验的详细指标，所有结果保存以供分析。 - **与 RL 环境集成：** 输出工件和指标专为下游 RL 和模拟 pipeline 设计。 - **错误处理和进度反馈：** 鲁棒的错误处理，进度打印语句和设备兼容性检查。 ## Poetry 文件：功能和用途 此仓库使用由 [Poetry](https://python-poetry.org/) 管理的 `pyproject.toml` 文件进行依赖管理和打包。Poetry 简化了 Python 环境的安装、版本控制和可复现性，确保所有必需的库及其版本得到跟踪并易于安装。 为了可复现的 GPU 工作流，此项目还包括 `environment.yml` (Conda) 来定义实验使用的 CUDA 对齐基础环境。 推荐设置步骤： 1. 从 `environment.yml` 创建并激活 Conda 环境： conda env create -f technique_inference_engine/environment.yml conda activate tie_gpu 2. 使用 Poetry 安装项目依赖： poetry install 3. 确保安装了启用 GPU 的 `implicit` 工件（如果当前环境中尚不存在）： conda install -c conda-forge implicit implicit-proc=*=gpu 简而言之：`environment.yml` 建立了系统级的 Python/CUDA 基线，而 Poetry (`pyproject.toml`) 管理项目级的 Python 包版本和 lockfile 可复现性。 ## 系统参考 此仓库在具有以下配置的多 GPU 工作站上进行了更改和测试： - GPUs (4): NVIDIA RTX A6000, NVIDIA GeForce RTX 2080 Ti - OS: Ubuntu 22.04.5 LTS - CPU Model: Intel(R) Xeon(R) Silver 4114 CPU @ 2.20GHz - CPU(s): 40 此硬件和 OS 配置能够加速上述推荐模型的训练和大规模实验。 ## 许可和归属 此仓库基于 [https://github.com/center-for-threat-informed-defense/technique-inference-engine/tree/TIE-83-jupyter-instructions](https://github.com/center-for-threat-informed-defense/technique-inference-engine/tree/TIE-83-jupyter-instructions) 的原始开源版本，并包含用于研究和开发目的的修改和增强。请支持并引用原始版本及其作者。 本项目根据 Apache License, Version 2.0 的条款分发。您可以复制、修改和分发本软件的副本，前提是您遵守许可证的条件，包括保留版权声明、许可证声明和归属声明。有关详细信息，请参阅 LICENSE 文件和上方的完整许可证文本。 欲了解更多信息，请访问 [Technique Inference Engine Website](https://center-for-threat-informed-defense.github.io/technique-inference-engine/)。

标签：Apex, Cloudflare, MITRE ATT&CK, Python, PyTorch, Vectored Exception Handling, 代码示例, 凭据扫描, 威胁情报, 开发者工具, 开源, 数据分析, 无后门, 机器学习, 深度学习, 知识推理, 矩阵分解, 网络安全, 逆向工具, 隐私保护, 预测模型