Add Differential K-Means implementation

changelog.d/+dkm.added

@@ -0,0 +1 @@
+Added Differentiable K-Means (DKM) training-time palettization (`coreai_opt.palettization.kmeans.dkm.DKMTrainingSpec`).

docs/src/examples/toy_models.md

@@ -7,6 +7,7 @@ These tutorials demonstrate model compression techniques using simple toy models
 
 /tutorials/mnist_quantization
 /tutorials/mnist_palettization
+/tutorials/mnist_dkm_palettization
 /tutorials/mnist_palettization_and_activation_quantization
 /tutorials/mnist_pruning
 ```

docs/src/palettization/config.md

@@ -346,3 +346,114 @@ kmeans_palettization_config:
       n_init: 1
       max_iter: 1
 ```
+
+## Training strategy: DKM (`training_strategy_spec`)
+
+Each {class}`~coreai_opt.palettization.spec.PalettizationSpec` carries a `training_strategy_spec` field that selects the training-time behavior for that weight. It defaults to vanilla (frozen) k-means, where the codebook and index assignments computed at `prepare()` stay fixed during training. Setting it to a `DKMTrainingSpec` instead enables [Differentiable K-Means](overview.md#differentiable-k-means-dkm-api), so the codebook and assignments are learned through gradient-based fine-tuning.
+
+`DKMTrainingSpec` exposes the following settings:
+
+- `tau`: Softmax temperature for the soft centroid assignment (default `1e-4`). Smaller values sharpen toward hard k-means; larger values soften it. The loss-minimizing `tau` is model-dependent — **smaller models and lower bit-widths tend to prefer a larger `tau`**. A good starting point is `1e-5` for LLMs and the default `1e-4` for smaller models.
+- `kmeans_max_iter`: Number of soft-assignment EM iterations per DKM training forward pass (default `3`).
+- `enforce_zero`: If `True`, pins the centroid nearest the origin to exactly zero (default `False`). Useful for joint pruning and palettization.
+- `zero_threshold`: Numerical-stability floor applied to the softmax assignment weights and per-centroid mass at small `tau` (default `1e-7`).
+
+### Apply DKM to all modules
+
+Set `training_strategy_spec=DKMTrainingSpec()` on the `global_config` spec to enable DKM for every palettized weight in the model:
+
+```python
+from coreai_opt.palettization import (
+    KMeansPalettizerConfig,
+    ModuleKMeansPalettizerConfig,
+    PalettizationSpec,
+)
+from coreai_opt.palettization.kmeans.dkm import DKMTrainingSpec
+
+config = KMeansPalettizerConfig(
+    global_config=ModuleKMeansPalettizerConfig(
+        op_state_spec={
+            "weight": PalettizationSpec(
+                n_bits=4,
+                training_strategy_spec=DKMTrainingSpec(),
+            )
+        },
+    )
+)
+```
+
+```yaml
+# yaml
+kmeans_palettization_config:
+  global_config:
+    op_state_spec:
+      weight:
+        n_bits: 4
+        training_strategy_spec:
+          type: dkm
+```
+
+### Apply DKM to a subset of modules, vanilla k-means to the rest
+
+Because `training_strategy_spec` lives on the per-weight `PalettizationSpec`, DKM can be enabled for selected modules while the rest keep the default vanilla k-means. Below, all layers are palettized 4-bit; two named modules train their palettization with DKM, and every other module (covered by `global_config`) uses vanilla k-means simply by not setting a `training_strategy_spec`.
+
+```python
+from coreai_opt.palettization import (
+    KMeansPalettizerConfig,
+    ModuleKMeansPalettizerConfig,
+    PalettizationSpec,
+)
+from coreai_opt.palettization.kmeans.dkm import DKMTrainingSpec
+
+config = KMeansPalettizerConfig(
+    # Default for all layers: vanilla 4-bit k-means (no training_strategy_spec).
+    global_config=ModuleKMeansPalettizerConfig(
+        op_state_spec={"weight": PalettizationSpec(n_bits=4)},
+    ),
+    module_name_configs={
+        "encoder.layer.0.attention": ModuleKMeansPalettizerConfig(
+            op_state_spec={
+                "weight": PalettizationSpec(
+                    n_bits=4,
+                    training_strategy_spec=DKMTrainingSpec(tau=1e-4, kmeans_max_iter=3),
+                )
+            },
+        ),
+        "encoder.layer.1.attention": ModuleKMeansPalettizerConfig(
+            op_state_spec={
+                "weight": PalettizationSpec(
+                    n_bits=4,
+                    training_strategy_spec=DKMTrainingSpec(tau=5e-5),
+                )
+            },
+        ),
+    },
+)
+```
+
+```yaml
+# yaml
+kmeans_palettization_config:
+  global_config:
+    op_state_spec:
+      weight:
+        n_bits: 4
+  module_name_configs:
+    encoder.layer.0.attention:
+      op_state_spec:
+        weight:
+          n_bits: 4
+          training_strategy_spec:
+            type: dkm
+            tau: 1.0e-4
+            kmeans_max_iter: 3
+    encoder.layer.1.attention:
+      op_state_spec:
+        weight:
+          n_bits: 4
+          training_strategy_spec:
+            type: dkm
+            tau: 5.0e-5
+```
+
+For a DKM spec with the defaults, the `default_dkm_weight_palettization_spec()` helper (importable from `coreai_opt.palettization.kmeans.dkm`) returns a ready-made 4-bit, per-tensor `PalettizationSpec` with `training_strategy_spec=DKMTrainingSpec()`.

docs/src/palettization/overview.md

@@ -140,7 +140,7 @@ with palettizer.training_mode():
 
 ### Using PATSchedule
 
-Users may also choose to enable palettization of parameters at a future point during training instead of right from the start. Weights which are marked for palettization, but have not yet had their palettizers enabled, will still be able to update during training. When a weight's palettizer is switched from disabled to enabled, its codebook and index assignments are recomputed using the most current weight values.
+Users may also choose to enable palettization of parameters at a future point during training instead of right from the start. Weights which are marked for palettization, but have not yet had their palettizers enabled, will still be able to update during training. When a weight's palettizer is switched from disabled to enabled, its codebook and index assignments are recomputed using the most current weight values. `PATSchedule` applies to any training strategy, including [DKM](#differentiable-k-means-dkm-api).
 
 Since palettizer enabling can be configured on a module by module basis, this gives users more flexibility on training to be palettization-aware by allowing weights to update in response to palettization noise from other palettized weights before themselves being palettized.
 
@@ -149,7 +149,7 @@ For example, users could choose to apply the following workflow:
 1. Begin training with all parameters marked to be palettized but with all palettizers disabled
 2. Fine-tune the model with no palettization noise for a few epochs
 3. Enable a subset of palettizers while keeping the rest of the palettizers disabled
-4. Fine-tune the model with palettization noise from the subset of enabled palettizers (any weights with enabled palettizers will not update further, but any weights with palettizers not yet enabled will continue training)
+4. Fine-tune the model with palettization noise from the subset of enabled palettizers (with the default strategy, any weights with enabled palettizers will not update further, but any weights with palettizers not yet enabled will continue training)
 5. Enable the rest of palettizers in the model and evaluate the final accuracy
 
 Steps 3 and 4 can be repeated as many times as desired with more and more palettizers being enabled.
@@ -196,3 +196,45 @@ with palettizer.training_mode():
 Each time `palettizer.step()` is called, the palettizer checks each module's `PATSchedule` to see whether any palettizer should be enabled. With the above config settings of `enable_fake_palettize` = 2 and 4 for `module1` and `module2` respectively, this results in both palettizers being disabled to start, `module1`'s palettizer enabled after `palettizer.step()` has been called twice, and `module2`'s palettizer enabled after `palettizer.step()` has been called 4 times in total.
 
 Note that `palettizer.step()` can be called whenever the user desires, whether it be per-batch, per-epoch, or on some other cadence.
+
+## Differentiable K-Means (DKM) API
+
+DKM palettization allows the palettization codebook and cluster assignments to be learned through training, removing the inability for already palettized weights to update their codebook and index assignments. This is achieved through calculation of attention scores between all weights and all centroids, along with utilizing a soft-assignment instead of hard-assignment during training forward, allowing gradients to propagate through palettizers to the original weight parameters.
+
+For more details on DKM, refer to [DKM (Cho et al., 2022)](https://arxiv.org/abs/2108.12659).
+
+Enabling DKM simply requires setting `training_strategy_spec` of `PalettizationSpec` to the `DKMTrainingSpec` class. The same `KMeansPalettizer` class is used for palettization. For example, the below code shows enabling DKM for all weights in the model:
+
+```python
+from coreai_opt.palettization import (
+    KMeansPalettizer,
+    KMeansPalettizerConfig,
+    ModuleKMeansPalettizerConfig,
+    PalettizationSpec,
+)
+from coreai_opt.palettization.kmeans.dkm import DKMTrainingSpec
+
+kmeans_palettizer_config = KMeansPalettizerConfig(
+    global_config=ModuleKMeansPalettizerConfig(
+        op_state_spec={
+            "weight": PalettizationSpec(
+                ..., training_strategy_spec=DKMTrainingSpec(...)
+            )
+        }
+    )
+)
+
+palettizer = KMeansPalettizer(model, kmeans_palettizer_config)
+prepared_model = palettizer.prepare(example_inputs)
+
+with palettizer.training_mode():
+    for i in range(...):
+        # Perform whatever training forward pass / loss calculation / optimizer gradient application is desired
+        ...
+
+finalized_model = palettizer.finalize(...)
+```
+
+DKM works with `PATSchedule` as well. A schedule decides when a module's palettizer becomes active, while the training strategy decides what that palettizer's forward pass does once it is. For example, you can leave DKM modules unpalettized for the first few hundred steps and then turn on soft-assignment training. To do this, set `pat_schedule` on the module config and call `palettizer.step()` in the training loop as described in [Using PATSchedule](#using-patschedule). Unlike with the default strategy, a DKM weight keeps updating after its palettizer is enabled, because gradients flow through the soft assignment.
+
+As `training_strategy_spec` is part of `PalettizationSpec`, which in turn can be configured on a per-op or module basis in `KMeansPalettizerConfig`, users may choose to enable DKM for a subset of palettized weights in the model while leaving another portion of palettized weights using non-DKM `Vanilla K-means`.