arlenecc/potion-multilingual-128M-mlx

potion-multilingual-128M — MLX 版（INT8 / INT4）

给 Apple Silicon 用的多语言静态嵌入模型量化权重包。 源模型 minishlab/potion-multilingual-128M （vocab 500353 × dim 256，fp32 约 512 MB）。

目录	权重文件	大小	相比 fp32	量化方案
int8/	weights.npz	136.10 MB	3.765x	对称 INT8，group_size=32
int4/	weights.npz	80.06 MB	6.40x	非对称 INT4，group_size=32，2 nibbles/byte

每个变体目录自包含：weights.npz + config.json + tokenizer.json + example.py。

potion-multilingual-128M-mlx/
├── int8/{weights.npz, config.json, tokenizer.json, example.py}
├── int4/{...}
├── convert_mlx.py                 # 转换脚本（纯 numpy，不需要 mlx）
├── verify_mlx.py                  # 验证脚本（含 numpy 版 mlx.core 替身）
├── _example_template.py           # 各变体 example.py 的来源，改这里后重跑 convert_mlx.py 同步
├── eval_data_multilingual_bitext.json   # 2124 条英↔法/德/西/日/阿 平行句对（验证语料）
├── conversion_report_mlx.json
├── verification_report_mlx.json
└── requirements.txt

1. 运行环境

pip install mlx tokenizers numpy safetensors     # Apple Silicon (M 系列) + macOS
python int8/example.py                            # 跑 demo

⚠️ MLX 只支持 Apple Silicon（Intel Mac / Linux / Windows 不行）。转换过程不需要 MLX， 只有加载推理时需要。

2. 量化布局

	INT8（对称）	INT4（非对称）
embeddings	int8 [V, D]	uint8 [V, D/2]：低 4 位=偶数维，高 4 位=奇数维
scales	fp16 [V, D/group]	fp16 [V, D/group]
biases	—	fp16 [V, D/group]
还原公式	x = q * scale	x = q * scale + bias
每组	`scale = max	x

group_size 默认 32，由 ONNX 版的 sweep.py 在 1800 句多语言真实语料上扫描确定（同为一阶最优点）。

**为什么不用 nn.QuantizedEmbedding**：MLX 原生量化层内部用 uint32 打包 + 特定的权重排列， 不同 mlx 版本细节可能变化。这里改用「显式量化的裸张量 + 只用基础算子写完 forward」， 任何 mlx 版本都能加载；同时也便于用 numpy 逐算子复刻来验证（见第 4 节）。 想用原生量化层的话，等价写法是 nn.QuantizedEmbedding.from_embedding(emb, group_size=32, bits=4)。

3. 使用

from example import StaticEmbeddingMLX

model = StaticEmbeddingMLX.from_dir("int8")     # 或 "int4"
vec = model.encode("Bonjour le monde")          # -> mx.array, shape (256,)
emb = model.encode(["Hello world", "你好世界", "こんにちは世界"])
sim = (emb @ emb.T)                             # mx.array，跨语言可直接比较

命令行：

python int8/example.py
python int4/example.py "What time does the train leave?" "¿A qué hora sale el tren?"

必须对齐的 3 个细节（已封装进 example.py 的 Model2VecTokenizer）： add_special_tokens=False；丢弃 unk token（id = 1）；超长文本截断到 512 token （model2vec 还会先按字符截到 512 × 7 = 3584 个字符）。输出向量已 L2 归一化，点积即余弦相似度。

4. 验证方式与结果

本机是 Intel Mac，pip install mlx 没有对应分发，因此无法真机跑 MLX。为了让验证仍然可信， verify_mlx.py 实现了 numpy 版的 mlx.core 替身（只实现用到的那十来个算子），注入 sys.modules 后真实 import 并执行 example.py 里的 forward——验证的是用户最终会跑的那份代码本身， 而不是另一套独立的理论公式。真机上 pip install mlx 后执行同样的脚本即可复验。

4248 句真实多语言语料（2124 条平行句对，英↔法/德/西/日/阿，来自 sentence-transformers/parallel-sentences-news-commentary，已随目录附带）：

检查项	结果
权重包完整性（npz 可加载、shape/dtype 与 config 一致、无 NaN/Inf）	✅ int8 / int4 均通过
分词与 model2vec 官方一致	✅ 4248/4248 句 token ids 完全一致
forward 输出	✅ 无 NaN，‖v‖ 偏差 ≤ 1e-15（单位向量）
边界输入（空串/单 token/CJK/emoji/超长/混合 padding batch）	✅ 全部 OK

4.1 精度（对照官方 model2vec fp32 实现）

变体	cosine mean	cosine min	cosine p1	max|Δ|
int8	0.999991	0.999980	0.999984	1.97e-03
int4	0.998085	0.995826	0.996736	2.80e-02

4.2 跨语言检索（query=英文句，候选=目标语言译文集合）

模型	R@1 均值	MRR@10	en-fr	en-de	en-es	en-ja	en-ar
model2vec 官方（参考）	0.7582	0.8124	0.836	0.822	0.858	0.594	0.681
MLX int8	0.7589	0.8126	0.836	0.822	0.858	0.597	0.681
MLX int4	0.7563	0.8113	0.830	0.820	0.856	0.601	0.675

排序差异（5 个语言对汇总）：

变体	名次完全一致率	Top-1 一致率	R@1 掉点	MRR@10 掉点
int8	0.9671	0.9993	-0.0007（实际略优）	-0.0002
int4	0.8805	0.9891	+0.0019	+0.0011

→ 量化带来的业务指标变化都在 ±0.002 以内（噪声量级），int8 基本无损，int4 损失可控。

4.3 与 ONNX 版的交叉一致性

指标	ONNX int8	MLX int8	ONNX int4	MLX int4
cosine mean	0.999991	0.999991（一致）	0.998050	0.998085
max|Δ|	1.965e-03	1.965e-03（一致）	2.59e-02	2.80e-02
R@1	0.7587	0.7589	0.7587	0.7563
权重体积	136.10 MB	136.10 MB	80.07 MB	80.06 MB

INT8 两边数值逐位一致（分组方式相同）；INT4 略有差异是预期的——ONNX 版为了避免动态 reshape 使用了「偶数维/奇数维分半」的列重排，MLX 版用标准的连续分组，因此每组的成员不同。

5. 复现

pip install -r requirements.txt

python convert_mlx.py                                   # 导出 int8 + int4
python convert_mlx.py --variants int4 --int4-group 64   # 换 group size
python verify_mlx.py                                     # 完整性 + 精度 + 检索评测

6. 依赖

numpy>=1.24
tokenizers>=0.15
safetensors>=0.4          # convert_mlx.py 读取源权重
model2vec>=0.4            # verify_mlx.py 的对拍基线（可选，缺失自动跳过）
mlx>=0.16                 # 仅推理需要，且仅 Apple Silicon 可用

7. 已知限制 / 待办

• 未在真机跑过 MLX：构建环境为 Intel macOS，以上 forward 结论来自 numpy 逐算子替身。 首次在 Apple Silicon 上使用时建议先跑一次 python verify_mlx.py（会走同一份 example.py）。
• MLX 生态里没有针对这类「静态查表 + 均值池化」模型的官方转换器，example.py 里的 forward 是 手写的等量实现，改动 example.py 请注意保持 output 语义。
• 统一内存优势带来的实际加速/能耗数据需要在 M 系列机器上另测（本环境的性能数字不可比，故未列出）。