• AI 应用
• 机器学习
• 模型优化

Palettization（也称为 weight clustering / 权重重聚类）是 Core ML 模型压缩的核心手段之一：它将模型的 float 权重聚类为少量簇心（centroid），构建查找表（Lookup Table，LUT），并用指向簇心的索引代替原始权重，从而显著减小模型体积。本文以 docs-guides/source/opt-palettization-overview.md 为骨架，结合coremltools.optimize的 Core ML 与 PyTorch 两侧源码实现，系统讲解 palettization 的数学原理、粒度配置、向量聚类、LUT 量化，以及 K-Means / Sensitive K-Means / Differentiable K-Means 三套算法的选型与完整实战代码。读完本文，你将掌握在 coremltools 中为 mlprogram 模型或 PyTorch 模型配置并执行 1/2/3/4/6/8 位权重聚类压缩的完整方案。

Palettization 是什么：用簇心索引替代 float 权重

Palettization 通过聚类把模型的 float 权重压缩成「查找表（LUT）+ 索引」两部分：数值相近的权重被归入同一个簇，统一用该簇的簇心（centroid）表示。原始权重矩阵因此被转换为一张索引表，表中每个元素指向其所属簇在 LUT 中的位置。

以一个简单例子说明（出自 palettize_weights 函数文档）：float 权重向量{0.3, 0.3, 0.5, 0.5}可以用 1-bit LUT{0.3, 0.5}压缩，权重向量被替换为 1-bit 索引{0, 0, 1, 1}。运行时，Core ML 依据 LUT 与索引现场重建 float 权重再参与计算，权重精度由 float 32/16 降低为非均匀的簇心值——这正是「非均匀精度降低（non-uniform lowering of precision）」的含义。

支持的位宽：N = {1, 2, 3, 4, 6, 8}

N为 palettization 使用的比特数，每个 LUT 拥有2^N个条目。coremltools 支持N={1,2,3,4,6,8}，这一集合在源码中作为硬约束存在：Core ML 侧 _quantization_passes.py 定义了_SUPPORTED_NBITS = (1, 2, 3, 4, 6, 8)，PyTorch 侧 K-Means 实现同样遵循该集合。当传入集合之外的nbits时会抛出Invalid nbits. Support (1, 2, 3, 4, 6, 8)类错误。

粒度（Granularity）：从 per_tensor 到 per_grouped_channel

上图展示的是per_tensor粒度：整张权重张量共享一个LUT。对于大矩阵，单一 LUT 会带来较高的近似误差。为此，从iOS18/macOS15起，coremltools 引入per_grouped_channel模式：由参数group_size指定的一组输出通道共享一个 LUT，整个权重矩阵拥有多个LUT。例如形状为(1024, 1024)的权重，若group_size=16，则共有1024 / 16 = 64个 LUT，簇心数量大幅增加，聚类近似精度显著提升。

源码中的粒度配置与校验

粒度在 Core ML 侧由 OpPalettizerConfig 控制，关键参数：

在 PyTorch 侧（如 PostTrainingPalettizerConfig 与 ModuleDKMPalettizerConfig），group_size默认是None，并且有严格的互斥校验：当granularity == "per_grouped_channel"时必须提供group_size（group_size has to be specified along with per_grouped_channel granularity），且group_size > 0；反之，per_tensor模式下指定group_size会直接报错。channel_axis默认指向输出通道轴（DKM 目前只支持输出通道轴，源码校验in_([0])，见 palettization_config.py）。

以(16, 25)形状的权重矩阵为例，若group_size=8，LUT 的形状即为(2, 2^n_bits)——每 8 个通道共享一个含2^n_bits个条目的 LUT。

逐通道缩放（Per-channel Scale）

启用enable_per_channel_scale=True后，权重会先沿输出通道维度按逐通道比例因子（per-channel scale）归一化，再执行聚类。归一化后的权重在相似数值范围内分布更集中，可进一步改善聚类质量；该模式可与per_grouped_channel组合使用。

源码层面的约束值得注意：enable_per_channel_scale=True与cluster_dim > 1（向量聚类）互斥。在 post_training_palettization.py 与 sensitive_k_means.py 中均有校验，违反时会抛出Enabling per_channel_scale is not supported for cluster_dim=... larger than 1的断言错误。此外，enable_fast_kmeans_mode开启时会将 fp32 权重先转 fp16 再取整以加速聚类，该优化同样不与cluster_dim > 1兼容。

向量聚类（Vector Palettization）：cluster_dim 详解

簇心既可以是标量，也可以是向量，由参数cluster_dim控制：

• cluster_dim = 1（默认）：标量聚类，LUT 中每个条目是一个标量；
• cluster_dim > 1：2-D 聚类，沿输出通道方向每cluster_dim长度的权重向量共享同一个 2-D 簇心，LUT 中每个条目的长度等于cluster_dim。

向量聚类能够在更低的位宽下保留相邻元素间的结构信息，适合对精度要求更高的场景。Core ML 侧 OpPalettizerConfig 对cluster_dim的描述与 PyTorch 侧ModulePostTrainingPalettizerConfig.cluster_dim的语义完全一致（见 post_training_palettization.py）。

量化 LUT：把查找表也压一压

LUT 中的值默认以 float 存储。从iOS18/macOS15起，LUT 本身也可以按 8-bit 精度存储，实现进一步的压缩。

• Core ML 侧：LUT 量化由palettize_weights图 pass 的joint_compression选项及后续量化 pass 支持，可将constexpr_lut_to_dense中的 LUT 元素进一步量化（见 _quantization_passes.py 中 "palettization + quantization" 的联合压缩说明）。
• PyTorch 侧：通过lut_dtype参数控制，支持的取值为torch.int8、torch.uint8；默认None表示不量化 LUT（见 post_training_palettization.py）。在 DKM 配置中对应lut_dtype字符串选项'i8' / 'u8' / 'f16' / 'bf16' / 'f32'，默认'f32'（见 palettization_config.py）。

可用性（Feature Availability）

Core MLmlprogram模型的 palettized 权重表示自iOS16 / macOS13 / watchOS9 / tvOS16及更新的部署目标（deployment target）可用；per_grouped_channel粒度与 8-bit LUT 则要求iOS18 / macOS15及以上，转换时需通过minimum_deployment_target=ct.target.iOS18显式声明。

三种 Palettization 算法选型

coremltools.optimize提供三套算法，它们为相同压缩率下的精度表现不同，选择时应结合数据、算力与精度诉求：

K-Means：无数据后训练聚类

最易上手的算法：直接用 k-means 聚类权重，簇心即 LUT。只需要模型权重，无需任何数据。官方精度基准（见 opt-palettization-algos.md）显示，MobileNetv2-1.0 经 6-bit K-Means 后 Top-1 精度 70.89（压缩率 2.54×），8-bit 为 71.80（1.97×），而 float16 基线为 71.86——可见较高位宽下压缩-精度权衡良好。对应 API：cto.coreml.palettize_weights与cto.torch.palettization.PostTrainingPalettizer。

Sensitive K-Means：基于 Fisher 信息的加权聚类

SKM 源自 SqueezeLLM: Dense-and-Sparse Quantization，通过_register_grad_square_hooks在 forward 过程中累计梯度平方）。

元素越敏感，扰动（此处即聚类）它对模型损失的影响越大；加权 k-means 会把簇心向敏感权重值靠近，使其表示得更精确，从而降低精度劣化。实践中通常128 个校准样本即可（SKMPalettizerConfig.calibration_nsamples默认值即为 128，见 sensitive_k_means.py），且对大型 transformer 架构效果显著优于无数据 K-Means。对应 API：cto.torch.palettization.SKMPalettizer。

Differentiable K-Means：训练期可微分聚类

DKM 是训练期算法：每个训练步对权重张量执行可微的软 k-means 聚类分配，使流程中每个操作都可微分，从而在训练过程中一边更新权重、一边学习 LUT 与索引（见 palettizer.py 的类注释）。实现方式是在模型内插入模拟聚类的子模块（FakePalettize层），并用注意力机制计算可微聚类；训练结束后通过finalize()将 LUT 与索引折叠回权重。

它在所有算法中压缩-精度权衡最佳，可支撑 2-bit 等极低位宽：官方基准中 ResNet50 经 2-bit DKM 压缩率达 7.63×、精度 75.47 ± 0.05（float16 基线 76.14），MobileNetv3-small 2-bit 压缩率 5.82×、精度 59.82 ± 0.98。代价是时间与数据开销最大——涉及距离矩阵与注意力矩阵计算，内存占用可观。对应 API：cto.torch.palettization.DKMPalettizer。milestone参数可控制从第几步开始启用 palettization（默认 0，即训练一开始就启用），源码中通过_milestones字典跟踪（见 palettizer.py）。

实战：压缩一个 Core ML mlprogram 模型

基础用法：全局 6-bit K-Means

对已转换为 Core ML 的模型，cto.coreml.palettize_weights会遍历 mlprogram 中所有权重，按mode指定的算法构建 LUT，将存储权重的constop 替换为constexpr_lut_to_denseop（见 _post_training_quantization.py 与 _quantization_passes.py 的图变换说明）。下面的例子对所有参数数量超过 512 的 op 应用 6-bit 聚类：

import coremltools as ct import coremltools.optimize as cto # 加载模型 mlmodel = ct.models.MLModel(uncompressed_model_path) # 定义 op 配置：6-bit，仅压缩元素数超过 512 的权重 op_config = cto.coreml.OpPalettizerConfig(nbits=6, weight_threshold=512) # 将 op 配置全局应用于所有 op config = cto.coreml.OptimizationConfig(global_config=op_config) # 执行权重聚类压缩 compressed_mlmodel = cto.coreml.palettize_weights(mlmodel, config)

按 op 类型 / 名称定制

OptimizationConfig支持op_type_configs与op_name_configs实现细粒度控制；置为None表示跳过该 op。下面的例子对除linear外的全部 op 应用 6-bit，linearop 用 8-bit，并跳过名为conv1、conv3的两个卷积：

import coremltools as ct import coremltools.optimize as cto mlmodel = ct.models.MLModel(uncompressed_model_path) global_config = cto.coreml.OpPalettizerConfig(nbits=6) linear_config = cto.coreml.OpPalettizerConfig(nbits=8) config = cto.coreml.OptimizationConfig( global_config=global_config, op_type_configs={"linear": linear_config}, op_name_configs={"conv1": None, "conv3": None}, ) compressed_mlmodel = cto.coreml.palettize_weights(mlmodel, config)

若想按 op 名定制，可先用 get_weights_metadata() 工具获取网络中所有权重及其所喂入 op 的详细信息，再挑选要定制的 op 名。

Core ML 侧 LUT 构建的四种 mode

OpPalettizerConfig 通过mode决定 LUT 的构建方式：

• "kmeans"（默认）：k-means 聚类生成 LUT，需提供nbits；
• "uniform"：线性直方图，LUT 为[v_min, v_min+scale, ..., v_max]，其中scale = (v_max - v_min) / (1 << nbits - 1)，需提供nbits；
• "unique"：直接用权重中的唯一值构成 LUT（权重假设位于离散格点但以 float 存储），不能提供nbits，系统自动从{1, 2, 4, 6, 8}中取能容纳唯一值数的最小值；若唯一值超过 256 个则跳过压缩；
• "custom"：自定义lut_function(weight) -> (lut, indices)计算 LUT 与索引，不能提供nbits。

注意nbits的互斥校验（_config.py）：kmeans/uniform必须提供nbits，unique/custom必须不提供，否则抛 ValueError。

实战：压缩一个 PyTorch 模型

Post-Training Palettization（数据无关）

PyTorch 侧流程与 Core ML 侧一致，但作用于 Torch 模型：先构造配置对象（支持代码内字典或从磁盘加载 YAML），再调用compress()。下面的例子对全部 op 应用 4-bit，并启用per_grouped_channel（group_size=4，即每 4 个通道共享一个 LUT）：

from model_utilities import get_torch_model from coremltools.optimize.torch.palettization import PostTrainingPalettizer, \ PostTrainingPalettizerConfig # 加载模型 torch_model = get_torch_model() palettization_config_dict = { "global_config": {"n_bits": 4, "granularity": "per_grouped_channel", "group_size": 4}, } palettization_config = PostTrainingPalettizerConfig.from_dict(palettization_config_dict) palettizer = PostTrainingPalettizer(torch_model, palettization_config) palettized_torch_model = palettizer.compress()

compress()内部把每个受支持模块的参数转成KMeansConfig（携带n_bits、axis、lut_dtype、block_size即group_size、cluster_dim、enable_per_channel_scale等），随后调用SequentialKMeans或ParallelKMeans完成聚类（见 post_training_palettization.py）。compress(num_kmeans_workers=N)支持多进程并行聚类以加速，推荐多 CPU 环境下使用。

Sensitive K-Means：YAML 配置 + 校准数据

所有算法的配置都可以写成 YAML 从磁盘加载。以下palettization_config.yaml定义了全局 4-bit、Linear模块 6-bit、跳过conv1/conv3、使用 64 个校准样本：

global_config: n_bits: 4 module_type_configs: Linear: n_bits: 6 module_name_configs: conv1: null conv3: null calibration_nsamples: 64

对应的 Python 脚本加载 YAML 并执行 SKM：

from model_utilities import get_torch_model from data_utils import get_calibration_data from coremltools.optimize.torch.palettization import (SKMPalettizer, SKMPalettizerConfig) import torch.nn.functional as F torch_model = get_torch_model() palettization_config = SKMPalettizerConfig.from_yaml('palettization_config.yaml') # 定义损失函数：前向传播后计算训练损失 loss_fn = lambda mod, dat: F.nll_loss(mod(dat[0]), dat[1]) calibration_data = get_calibration_data() palettizer = SKMPalettizer(torch_model, palettization_config) palettized_torch_model = palettizer.compress(dataloader=calibration_data, loss_fn=loss_fn)

SKM 的compress()比 post-training 版本多两个必填参数：

• dataloader：可迭代对象，每个元素是模型的一个输入，用于计算模型权重的梯度；
• loss_fn：可调用对象，接收模型与数据、执行前向并计算训练损失。

module_name_configs的键支持正则或完全限定名（见 sensitive_k_means.py），比 Core ML 侧的精确 op 名匹配更灵活。

Differentiable K-Means：插入训练循环

DKM 直接嵌入训练流程，一般需要更多数据。标准流程是五个步骤：

1. 用DKMPalettizerConfig定义聚类参数；
2. 用DKMPalettizer初始化 palettizer；
3. 调用prepare(inplace=True)，将原模型更新为含FakePalettize层的「palettization 友好」模型；
4. 运行常规训练循环，并在每个 step 额外调用palettizer.step()；
5. 模型收敛后调用finalize()，得到可转换到 Core ML 的模型。

下面的例子对全部 op 应用 2-bit DKM，并指定milestone=10——从第 10 个训练步开始启用聚类：

from model_utilities import get_torch_model from training_utilities import train_step from data_utils import get_dataloader from coremltools.optimize.torch.palettization import (DKMPalettizer, DKMPalettizerConfig) import torch.nn as nn torch_model = get_torch_model() dataloader = get_dataloader() num_palettization_epochs = 2 palettization_config_dict = { "global_config": {"n_bits": 4, "milestone": 10}, } palettization_config = DKMPalettizerConfig.from_dict(palettization_config_dict) palettizer = DKMPalettizer(torch_model, palettization_config) # 向模型中插入 palettization 友好模块 palettizer.prepare(inplace=True) torch_model.train() for epoch in range(num_palettization_epochs): for batch_idx, (data, target) in enumerate(dataloader): train_step(data, target, torch_model) palettizer.step() palettized_torch_model = palettizer.finalize()

ModuleDKMPalettizerConfig还提供大量进阶旋钮，绝大多数情况下保持默认即可：kmeans_max_iter（默认 3）、palett_tau（softmax 温度，默认 0.0001）、enforce_zero（是否强制把最接近原点的簇心固定为 0）、palett_lambda（抑制离群点对簇心计算的影响）、palett_unique/palett_shard/palett_batch_mode（面向大语言模型的 eDKM 高效选项，见 palettization_config.py 的默认值表）。若module_type_configs的值是配置列表，还会按weight_threshold排序后匹配不同尺寸的权重（见 palettization_config.py）。

将 Palettized PyTorch 模型转换为 Core ML

使用ct.optimize.torch.*API 压缩过的 PyTorch 模型，coremltools 8 可直接转换，无需额外参数；但若用到了per_grouped_channel等仅 iOS18/macOS15 可用的新特性，必须显式指定部署目标：

import torch import coremltools as ct palettized_torch_model.eval() traced_model = torch.jit.trace(palettized_torch_model, example_input) palettized_coreml_model = ct.convert(traced_model, inputs=...) # 若聚类时使用了 iOS18 特性 palettized_coreml_model = ct.convert(traced_model, inputs=.., minimum_deployment_target=ct.target.iOS18)

如果使用其他非 coremltools 的聚类工具或 coremltools 8 以下版本，需参考 opt-conversion.md 中转换带正确 palettized op 的 Core ML 模型的流程。

性能与精度取舍

Palettization 减小了每个权重值的体积，推理时数据搬运量下降，对内存瓶颈型（memory-bottlenecked）模型收益明显；该延迟优势仅在权重被「按需即时解压（just-in-time decompression）」时成立。从 iOS17/macOS14 起，这一行为更可能发生在主要运行于 Neural Engine 后端的模型上（见 opt-palettization-perf.md）。

per_grouped_channel模式（iOS18/macOS15）有一个明确的权衡：LUT 数量越多，运行时性能越可能下降。经验表明，group_size取 8 或 16 通常能在保持精度的同时获得相对未压缩模型的加速。官方基准中 MobileNetv2-1.0 的 2-bit DKM 模型（压缩率 5.92×）在 iPhone 14 Pro（iOS17）上延迟约 0.47 ms/批，与 float16 基线（0.48 ms/批）基本持平（详见 opt-palettization-perf.md 的完整延迟表）。

测试与验证参考

仓库的测试覆盖了本文介绍的三种 API：test_post_training_palettization.py 验证了PostTrainingPalettizer（如test_no_config断言无配置时默认 4-bit 聚类后权重唯一值数为 16，即2^4个簇心）与SKMPalettizer；test_palettizer.py 覆盖DKMPalettizer；test_palettization_api.py 覆盖配置 API 行为。这些测试是验证配置参数与算法行为的最佳可运行样例，也可作为自定义压缩流程的起点。

小结

Palettization 是 Core ML 生态中与线性量化、剪枝并列的三大模型压缩手段。其核心由「簇心 LUT + 索引」组成，N={1,2,3,4,6,8}位宽、per_tensor/per_grouped_channel粒度、cluster_dim向量聚类、enable_per_channel_scale逐通道缩放与 8-bit LUT 量化共同构成了完整的配置空间。按数据与算力条件，可在无数据 K-Means、基于校准数据的 Sensitive K-Means 与训练期 Differentiable K-Means 之间选择，它们分别对应 Core ML 与 PyTorch 两套 API 族。压缩完成后，只需在转换时正确声明minimum_deployment_target，即可得到可直接部署到 Apple 平台的 mlprogram 模型。

• AI 应用
• 机器学习
• 模型优化