# NdLinear：下一代线性层的替代方案
这是一个可直接替换的 PyTorch 模块，由 [Ensemble AI](https://ensemblecore.ai) 提供，用于替代神经网络中的标准线性层。
**Ensemble 平台**
想要更小、更快的模型且**没有**精度损失？尝试过剪枝或量化却遇到了瓶颈？
## NdLinear
一个用于替代神经网络中标准线性层的 PyTorch 模块。
- ✅ `nn.Linear` 的即插即用替代品
- 📦 轻量级，参数高效
- 🧠 原生保留多变量结构
## Ensemble 平台（由 NdLinear 驱动）
上传任意模型 – 返回一个更小、更快的版本。
无精度损失。
- 🔁 自动替换层并调整超参数
- 📉 将模型大小（参数数量）缩减多达 8 倍
- 🛠 根据您的**硬件与微调限制**量身定制上传的模型
- 🧰 导出为 ONNX、TensorRT、SNPE 等
- 💡 旨在与其他压缩技术（剪枝、量化、蒸馏）协同工作
- 🎁 注册即赠送 **1000 万免费额度**
👉 [免费试用 Ensemble 平台→](https://app.ensemblecore.ai/signup)
📺 [或查看演示 →](https://ensemblecore.ai)
## 🧬 技术概述
NdLinear 保留了数据的多维结构，以更少的参数增强了表示能力。
它不会展平张量，而是跨结构化的向量空间集合对它们进行转换——从而捕获标准全连接层所丢失的依赖关系。
👉 [论文链接→](https://arxiv.org/pdf/2503.17353)
## 核心特性
- **结构保留：** 保留原始的数据格式和形状。
- **参数高效：** 在提升性能的同时减少参数数量。
- **极低开销：** 保持与传统线性层相同的复杂度。
- **灵活集成：** 无缝替换现有的线性层。
## 安装
要将 NdLinear 集成到您的项目中，请克隆此仓库并安装必要的依赖项：
```
git clone https://github.com/ensemble-core/NdLinear.git
cd NdLinear
pip install .
```
或者，如果已打包，请通过 pip 安装：
```
pip install ndlinear
```
或者，通过 conda 安装：
```
conda install conda-forge::ndlinear
```
## 用法
NdLinear 可用于各种神经网络架构中，例如 CNN、RNN 和 Transformer。
### 示例 1：使用 NdLinear 替换标准线性层
```
import torch
from ndlinear import NdLinear
input_tensor = torch.randn(32, 28, 28, 3) # Batch of images
ndlinear_layer = NdLinear(input_dims=(28, 28, 3), hidden_size=(64, 64, 6))
output = ndlinear_layer(input_tensor)
```
### 示例 2：Transformer
在 Transformer 架构中，您可能需要操作多维张量以实现高效的线性操作。以下是如何在 3D 输入张量中使用 `NdLinear` 的方法：
```
import torch
from ndlinear import NdLinear
input_tensor = torch.randn(32, 28, 28) # Input with shape : (batch_size, num_tokens, token_dim)
# 为线性操作重塑输入 tensor
input_tensor = input_tensor.reshape(-1, 28, 1) # New shape: (batch_size * num_tokens, token_dim, 1)
# 定义具有合适输入和隐藏维度的 NdLinear 层
ndlinear_layer = NdLinear(input_dims=(28, 1), hidden_size=(32, 1))
# 执行线性变换
output = ndlinear_layer(input_tensor)
# 处理后重塑回原始维度
output = output.reshape(32, 28, -1) # Final output shape: (32, 28, 32)
```
此示例说明了如何通过操作张量形状将 `NdLinear` 集成到 Transformer 模型中，从而保留进一步处理所需的结构，并实现高效的投影能力。
### 示例 3：多层感知机
此示例演示了如何在正向传播设置中使用 `NdLinear` 层，使得将其集成到现有的 MLP 结构中变得简单高效。
```
import torch
from ndlinear import NdLinear
input_tensor = torch.randn(32, 128)
# 为 MLP 定义第一个 NdLinear 层，其输入维度为 (128, 8)，隐藏 size 为 (64, 8)
layer1 = NdLinear(input_dims=(128, 8), hidden_size=(64, 8))
# 为 MLP 定义第二个 NdLinear 层，其输入维度为 (64, 8)，隐藏 size 为 (10, 2)
layer2 = NdLinear(input_dims=(64, 8), hidden_size=(10, 2))
x = F.relu(layer1(input_tensor))
output = layer2(x)
```
### 示例 4：边界情况
当 `input_dims` 和 `hidden_size` 为一维时，`NdLinear` 的功能等同于传统的 `nn.Linear` 层，这作为 `n=1` 时的边界情况。
```
from ndlinear import NdLinear
# 使用一维输入和隐藏 size 定义 NdLinear
layer1 = NdLinear(input_dims=(32,), hidden_size=(64,))
```
## NdLinearGated
NdLinearGated 通过引入复杂的门控机制来控制信息流，从而扩展了 NdLinear 的核心功能。这使得模型能够选择性地转换输入维度，增强了表示能力和效率。
### 核心特性
- **选择性信息流：** 控制应用哪些转换的动态门控机制
- **多种门控模式：** 支持软（连续）和硬（二值）门控方法
- **维度选择：** 将门控应用于所有维度、仅第一个维度，或自动应用于最重要的维度
### 用法
NdLinearGated 可以集成到神经网络中，并对门控行为进行细粒度控制：
```
import torch
from ndlinear import NdLinearGated
# 创建输入 tensor
input_tensor = torch.randn(32, 28, 28, 3) # Batch of images
# 初始化 NdLinearGated 并对所有维度应用 soft gating
gated_layer = NdLinearGated(
input_dims=(28, 28, 3),
hidden_size=(64, 64, 6),
gating_mode="soft",
gated_modes="all"
)
# 带 gating 的 Forward pass
output = gated_layer(input_tensor)
```
### 门控配置
NdLinearGated 提供了各种配置以适应不同的建模需求：
```
# 仅对第一个维度应用 hard gating
first_dim_gated = NdLinearGated(
input_dims=(128, 8),
hidden_size=(64, 8),
gating_mode="hard",
gated_modes="first"
)
# 对标准差最高的 top-k 维度应用 soft gating
topk_gated = NdLinearGated(
input_dims=(28, 28, 3),
hidden_size=(64, 64, 6),
gating_mode="soft",
gated_modes="topk"
)
```
### 推荐配置
基于大量的实验，我们推荐以下配置以获得最佳性能：
```
optimal_gated = NdLinearGated(
input_dims=(28, 28, 3),
hidden_size=(64, 64, 6),
gating_mode="soft",
gated_modes="topk",
gating_hidden_dim=16 # Adjust based on your model size
)
```
此配置（软门控 + top-k 模式）始终能带来：
- **更高的准确率：** 比基线 NdLinear 性能提升约 0.5-0.7%
- **计算效率：** 仅激活最有用的投影，从而减少 50-75% 的计算负载
- **训练稳定性：** 表现出稳定的训练过程和平滑的门控熵衰减
- **增强的可解释性：** 更清晰地揭示哪些维度最为重要
与硬门控相比，软门控机制提供了更好的稳定性，而 top-k 模式选择则将计算资源集中在信息量最大的维度上。可以根据您特定的模型架构和数据集要求调整 gating_hidden_dim 参数。
## 应用示例
NdLinear 用途广泛，可应用于：
- **图像分类：** 运行 `cnn_img_classification.py`。
- **时间序列预测：** 使用 `ts_forecast.py`。
- **文本分类：** 启动 `txt_classify_bert.py`。
- **Vision Transformers：** 执行 `vit_distill.py`。
## 联系方式
如有疑问或合作意向，请联系 [Alex Reneau](mailto:alex@ensemblecore.ai)。
## 许可证
本项目基于 Apache 2.0 许可证分发。查看 [LICENSE](https://github.com/ensemble-core/NdLinear/blob/main/LICENSE) 文件了解更多详情。