CoderDoge1108/ContraBin

# ContraBin 对比使用单纯形插值（simplex interpolation）进行二进制代码表示的对比预训练。 [](https://github.com/CoderDoge1108/ContraBin/actions/workflows/ci.yml) [](https://www.python.org/downloads/) [](LICENSE) [](https://arxiv.org/abs/2210.05102) ContraBin 通过对比学习 对齐同一程序的三个视角——**源代码**、**二进制 / LLVM IR** 以及简短的**自然语言注释**——实现可迁移的二进制代码表示， 采用单纯形插值的两阶段课程进行训练。 本仓库是该工作的开源、研究级 ## 亮点 - **模块化架构**——编码器、投影头、单纯形插值与对比损失均为独立模块，可自由组合。 - **离线友好**——每个组件均提供轻量测试主干（`contrabin-tiny`），完整训练流程可在笔记本电脑 CPU 上运行。 - **改进的下游任务**——四个设计更合理、更具研究价值的新任务（详见[下游任务](#downstream-tasks)）。 - **纯 CLI 工作流**——`contrabin build-triplets | pretrain | embed | task ...`。 - **现代 Python 打包**——基于 `pyproject.toml`、`pydantic` 配置与 `typer` CLI，集成 `ruff` + `mypy` + `pytest`，包含 30+ 个测试。 ## 架构 ``` ┌──────────┐ ┌──────────┐ ┌──────────┐ │ Source │ frozen │ Binary │ trainable │ Comment │ frozen │ encoder │──┐ │ encoder │──┐ │ encoder │──┐ └──────────┘ │ └──────────┘ │ └──────────┘ │ ▼ │ ▼ ┌────────────┐ ┌───────────┐ ┌────────────┐ │ src head │ │ bin head │ │ cmt head │ └─────┬──────┘ └─────┬─────┘ └─────┬──────┘ │ │ │ ├────────── CLIP loss ──┤── CLIP loss ─────────┤ │ │ └──── Simplex interpolation Γ(src, cmt; λ) ────┘ │ ▼ InfoNCE loss against binary (curriculum: naive -> linear -> nonlinear) ``` - 两个**锚定**（冻结）编码器共享权重，分别生成源编码与注释编码。 - 一个独立的**可训练**编码器生成二进制 / IR 编码。 - 训练时，两个锚定编码通过**单纯形插值模块**混合，生成*中间视角*，再与二进制编码通过 InfoNCE 目标对齐。 - 标量课程控制插值过程：从*朴素*（无混合）到*线性*（标量 λ）再到*非线性*（逐元素 λ）。 详见 [docs/method.md](docs/method.md) 与 [docs/training.md](docs/training.md) 获取完整的数学公式及论文算法 1。 ## 快速开始 ``` # 1. 安装（可编辑模式 + 开发依赖）。 pip install -e '.[dev]' # 2. 在合成数据上对整个堆栈进行冒烟测试（无需网络 / 无需 Clang）。 pytest -q # 3. 使用对 CPU 友好的配置在合成数据上驱动端到端训练。 contrabin make-synthetic --output data/processed/train.jsonl -n 64 contrabin make-synthetic --output data/processed/val.jsonl -n 16 --seed 1 contrabin pretrain --config configs/smoke.yaml # 4. 从 C 文件目录构建真实三元组（需要 Clang）。 contrabin build-triplets \ --input data/raw/anghabench \ --output data/processed/triplets.jsonl \ --comment-generator heuristic # 5. 使用完整课程进行预训练。 contrabin pretrain --config configs/pretrain.yaml # 6. 运行四个下游任务中的任意一个。 contrabin task retrieve --config configs/pretrain.yaml --gallery data/processed/gallery.jsonl --checkpoint outputs/contrabin/final.pt contrabin task name-recovery --config configs/pretrain.yaml --train data/.../train.jsonl --val data/.../val.jsonl contrabin task summarize --config configs/pretrain.yaml --train data/.../train.jsonl --val data/.../val.jsonl contrabin task provenance --config configs/pretrain.yaml --train data/.../train.jsonl --val data/.../val.jsonl ``` ## 下游任务 论文评估了四个任务，其中两个（函数名恢复与二进制摘要）任务设计良好并予以保留； 另外两个任务被重新设计为更易复现、更直接地衡量*表示质量*而非*微调头容量*： | 任务 | 论文设定 | 本框架设定 | 指标 | |------|----------|------------|------| | 二进制功能相似性 | POJ-104 104 分类 | **基于 POJ-104 标签的检索**（冻结编码，可选线性探针） | mAP、MRR、Recall@{1,5,10} | | 函数名恢复 | 多标签子词分类 | 同上 | 精确匹配、子词 F1 | | 二进制摘要 | 二进制 → 自然语言序列 | 同上 | BLEU、ROUGE-L | | 逆向工程 | 二进制 → C 源代码 | **编译器溯源恢复**（编译器 / 优化级别 / 语言） | 每头与联合准确率 | 完整说明请参见 [docs/tasks.md](docs/tasks.md)。 ## 仓库结构 ``` contrabin/ ├── contrabin/ # Library code │ ├── config.py # Pydantic configs + YAML loading │ ├── data/ # Triplet builder, clang frontend, datasets │ ├── models/ # Encoders, heads, simplex interpolation │ ├── losses/ # CLIP-style + InfoNCE objectives │ ├── training/ # Trainer, curriculum scheduler, callbacks │ ├── tasks/ # Four downstream tasks │ ├── evaluation/ # Metrics (mAP, MRR, BLEU, ROUGE, F1, ...) │ ├── utils/ # Logging, IO, seeding, visualization │ └── cli.py # Typer CLI (contrabin ...) ├── configs/ # YAML configs (pretrain / smoke / finetune) ├── tests/ # pytest suite (CPU-only, offline) ├── docs/ # Method, data, training, tasks, FAQ ├── scripts/ # Shell wrappers (build_triplets.sh, ...) ├── examples/ # Minimal runnable usage examples ├── notebooks/ # Jupyter quickstart + analysis notebooks ├── pyproject.toml ├── CHANGELOG.md └── README.md ``` ## 引用 若在研究中使用 ContraBin，请引用： ``` @article{zhang2025contrabin, title = {Pre-Training Representations of Binary Code Using Contrastive Learning}, author = {Zhang, Yifan and Huang, Chengzhi and Zhang, Yichi and Shao, Haoran and Leach, Kevin and Huang, Yu}, journal = {Transactions on Machine Learning Research}, year = {2025}, url = {https://arxiv.org/abs/2210.05102} } ``` ## 许可证 MIT。详见 [LICENSE](LICENSE)。

标签：Apex, CPU训练, LLVM IR, PyPI, Python, PyTorch, SIMD插值, TensorFlow, TMLR, TMLR 2025, TypeScript, 下游任务, 二进制, 二进制代码, 云资产清单, 代码克隆, 代码向量, 代码嵌入, 代码搜索, 代码检索, 代码理解, 代码相似性, 代码表示, 凭据扫描, 单纯形插值, 可移植表示, 多视图学习, 安全插件, 安全规则引擎, 对比学习, 开源, 文档结构分析, 无后门, 机器学习, 模块化架构, 测试骨干, 源代码, 研究级实现, 离线训练, 自然语言注释, 迁移学习, 逆向工具, 逆向工程, 预训练