TPA

把注意力复杂度从平方降到线性，长序列不再卡顿

开源 🏷️完全免费 🔧 仅开源项目，无在线服务，需自行部署 开源模型 🍴 复刻 41 🐛 待处理 1 👀 关注 11
访问官网 ↗ GitHub ↗
GitHub 星标 ★ 464
维护状态 活跃
是否开源 开源
定价模式 🏷️完全免费 MIT 许可证开源，完全免费，可自行部署使用。

项目预览

TPA 项目界面预览图
图片来源：GitHub 官方仓库名片

项目数据

分类开源模型
开发团队tensorgi
官网地址https://tensorgi.github.io/TPA
GitHubhttps://github.com/tensorgi/TPA
定价模式🏷️完全免费
价格说明MIT 许可证开源，完全免费，可自行部署使用。
访问状态🔧 仅开源项目，无在线服务，需自行部署
是否开源是
开源协议MIT
主要语言Python
技术栈/模型deep-learning,foundation-models,large-language-models,llms,model-architectures,pretraining
GitHub 星标★ 464
上线时间2025-01-10 00:00:00
代码更新2026-11-05
仓库体积11.2 MB
本站数据更新2026-11-11
维护状态活跃
收录时间2026-08-09
浏览次数25

项目介绍

TPA（Tensor ProducT ATTenTion）是一个NeurIPS 2025 Spotlight论文的开源实现，提出了一种新型注意力机制，旨在替代标准Transformer中的多头注意力（MHA）。它通过张量积分解将注意力计算分解为低秩因子，显著降低计算和内存复杂度，同时保持或提升模型性能。TPA解决了长序列任务中注意力机制O(n²)复杂度导致的效率瓶颈，核心能力包括：线性复杂度注意力、即插即用替换现有MHA、支持多种Transformer架构适配。项目提供PyTorch实现，代码简洁，便于集成到现有NLP或视觉模型中。该工作为高效Transformer研究提供了新思路，适合对模型效率有高要求的开发者实验和部署。

Star 增长趋势

本站每日记录该项目 GitHub 星标变化，观测越久曲线越完整

数据来源：GitHub API 逐日快照 · 收录于 2026-08-09 · 当前星标 ★ 464

使用教程

ℹ 本教程由 AI 依据该项目 README 整理，个别步骤可能需要按你的环境微调。

难度：高级 约 60 分钟 部署方式：本地GPU 4 步

环境要求

  • CUDA GPU
  • PyTorch 2.4
  • 大容量磁盘

安装与启动步骤

  1. 1拉代码

    研究工程，按论文配置环境。

    git clone https://github.com/tensorgi/TPA.git

  2. 2建环境装依赖

    版本列得比较具体。

    pip install torch==2.4.0 numpy transformers datasets tiktoken wandb tqdm

  3. 3准备数据

    它给了公开数据集的预处理脚本。

    python data/openwebtext/prepare.py

  4. 4跑预训练

    按配置起训练。

    bash pretrain.sh

如何确认成功

用小规模配置能跑起预训练若干步并产出检查点，损失按预期下降

注意事项

  • 它是张量积注意力那篇论文的官方实现，属于架构研究而非应用。
  • 依赖版本列得很死，照着装基本能过。
  • 全量预训练成本很高，先用小配置验证代码正确性。
  • 数据预处理与训练是两个独立步骤，跳过前一步训练会找不到文件。

核心亮点

  • 注意力机制线性复杂度，长序列推理速度显著提升
  • 可直接替换现有Transformer的多头注意力模块，改动小
  • 论文被NeurIPS 2025接收，理论有保障，代码实现清晰

不足之处

  • 项目处于早期，文档和社区支持待完善
  • 部分场景下低秩近似可能带来精度损失，需调参验证

适用场景

  • 长文本处理（如文档摘要、代码生成）
  • 高分辨率图像或视频理解任务
  • 资源受限设备上的Transformer部署

替代项目

FlashAttention、Linformer、Performer

上一篇：RLCSD

下一篇：GRAPE

同类项目推荐

LongRecipe 开源

把长文本压缩成短序列，高效训练长上下文模型。

LongRecipe: Recipe for Efficient Long Context Generalization in Large Language Model···

★ 79 2026-08-09
llm-engine 开源

自建大模型推理与微调平台，OpenAI 接口一键迁移

Scale LLM Engine public repository

★ 841 2026-09-16
AI-Toolkit-Easy-Install 开源

Windows 上一键装好 AI 微调工具，免折腾环境直接开训

One-click Portable Windows installation of 'AI-Toolkit by Ostris'

★ 542 2026-08-09
optimate 开源

给 AI 模型瘦身提速，部署更快更省资源

A collection of libraries to optimise AI model performances

★ 8327 2026-08-09
SepLLM 开源

压缩冗余片段，让大模型推理快 2.5 倍

[ICML 2025] "SepLLM: Accelerate Large Language Models by Compressing One Segment int···

★ 573 2026-08-09
FPO 开源

用强化学习微调流匹配VLA机器人策略

Official code for "Reinforcement Fine-Tuning of Flow-Matching Policies for Vision-La···

★ 131 2026-11-11