RFC: Weight Clustering for torchao

[wwwind]

Description

RFC: Weight Clustering for torchao

Field	Value
Status	Proposed
Type	New feature
Author(s)	Elena Zhelezina
Initial location	torchao/prototype/clustering/

---

Summary

This RFC proposes adding weight clustering to torchao as a first-class model optimization technique.

Weight clustering, also known as weight sharing, constrains each supported weight tensor to use at most N distinct values per clustering group. During fine-tuning, each group is represented by:

• a small trainable centroid table, and
• a non-trainable assignment tensor mapping every weight position to one centroid.

The forward pass reconstructs the effective weight tensor from the centroids and assignments. After fine-tuning, strip_clustering() replaces the wrapper with an ordinary PyTorch module whose weight tensor contains the clustered values directly.

The initial implementation is intentionally narrow:

• mirrors the TensorFlow Model Optimization Toolkit (TFMOT) clustering model and user-facing controls where they map cleanly to PyTorch;
• supports weight clustering for common dense and convolutional layers;
• does not change inference graphs after stripping;
• does not introduce backend-specific lowering, runtime codebook lookup, new operators, or CQAT integration in the initial PR;
• lands under torchao/prototype/clustering/ until the API and behavior are validated.

The proposed API follows torchao conventions by using config objects derived from AOBaseConfig and filter_fn(module, fqn) predicates, while preserving the TFMOT clustering vocabulary: number_of_clusters, cluster_centroids_init, cluster_per_channel, and preserve_sparsity.

---

Motivation

Problem

Quantization reduces the numeric precision of weights, but it does not explicitly control how many different weight values appear in a layer. Weight clustering addresses a different compression axis: it constrains weight tensors to use a small number of shared values.

For example, a layer with number_of_clusters=16 uses no more than 16 distinct values per configured clustering group after stripping. This can be useful for:

1. Model compression. Low-unique-value tensors are more compressible by general-purpose compression, model packaging formats, or backend-specific weight encoders when those systems exploit repeated values.
2. Regularization and model simplification. Weight sharing can act as a structural constraint during fine-tuning.
3. Optimization pipeline parity for users migrating from TensorFlow. TFMOT already provides weight clustering for Keras models. torchao does not currently provide an equivalent native PyTorch workflow.

This RFC is backend-agnostic. It does not claim a guaranteed size or latency improvement for any particular deployment target. The output of the initial workflow is an ordinary PyTorch model with clustered dense weights; any deployment benefit depends on the downstream serialization, quantization, and runtime stack.

Why this belongs in torchao

torchao is the PyTorch ecosystem library for model optimization workflows including quantization, sparsity, prototype optimizations, and training-time optimization techniques. Weight clustering is a natural addition to this family because it is a weight-domain optimization that composes with pruning and quantization but is distinct from both.

The RFC proposes starting in torchao/prototype/clustering/, consistent with introducing a new optimization workflow before stabilizing its public API.

---

Goals

1. Provide a PyTorch-native implementation of weight clustering in torchao.
2. Mirror the main TFMOT clustering concepts and controls where they map cleanly to PyTorch.
3. Preserve torchao API conventions:
   • config object based on AOBaseConfig;
   • in-place transform named cluster_weights_();
   • optional non-in-place transform named cluster_weights();
   • filter_fn(module, fqn) selection predicate.
4. Produce ordinary PyTorch modules after strip_clustering().
5. Support common layer types in the initial implementation:
   • nn.Linear.weight;
   • nn.Conv1d.weight;
   • nn.Conv2d.weight;
   • nn.Conv3d.weight.
6. Preserve exact zeros when requested through preserve_sparsity=True, enabling prune → cluster workflows.
7. Provide strict validation utilities to verify that stripped weights satisfy the requested cluster count.
8. Leave CQAT as future work, with enough design context in this RFC to explain why it is useful.

---

Non-goals

The initial RFC does not include:

• clustering activations;
• runtime codebook lookup or on-device decompression;
• new PyTorch, ATen, TOSA, or backend operators;
• backend-specific compiler changes;
• online reassignment or periodic k-means during training;
• latent full-precision shadow weights;
• group-wise or sub-channel clustering beyond per-tensor and per-output-channel modes;
• CQAT implementation;
• ConvTranspose* support;
• bias clustering by default;
• an API guarantee that clustered models compress by a fixed ratio in any particular deployment format.

---

Prior art: TFMOT weight clustering

TFMOT provides a Keras cluster_weights() API that wraps a model or layer with clustering functionality. The TFMOT API constrains each clustered weight tensor to no more than number_of_clusters unique values and expects the model to be pretrained before clustering is applied.

The core TFMOT public API is:

cluster_weights(
    to_cluster,
    number_of_clusters,
    cluster_centroids_init=CentroidInitialization.KMEANS_PLUS_PLUS,
    **kwargs,
)

TFMOT exposes centroid initialization strategies including:

• LINEAR;
• 'RANDOM';
• DENSITY_BASED;
• KMEANS_PLUS_PLUS.

The centroid 'RANDOM' is not very useful, so we are going to drop it in PyTorch implementation.

TFMOT also documents more advanced clustering workflows, including per-channel clustering, sparsity-preserving clustering, and cluster-preserving quantization-aware training.

This torchao RFC intentionally mirrors the TFMOT mental model:

• wrap a pretrained model;
• fine-tune with clustering constraints active;
• strip clustering wrappers before deployment or further optimization.

The PyTorch API differs where Keras-specific mechanics do not apply. In particular, PyTorch does not need TFMOT's Keras layer reconstruction **kwargs passthrough.

---

Proposed API

Module placement

Phase 1 prototype:

torchao/prototype/clustering/
    __init__.py
    config.py
    centroids.py
    core.py
    api.py
    validation.py
    debug.py

Prototype exports:

from torchao.prototype.clustering import (
    ClusteringConfig,
    CentroidInitialization,
    cluster_weights_,
    cluster_weights,
    strip_clustering,
    validate_clustering,
)

Stable graduation location remains an open question. See "Open questions".

---

CentroidInitialization

from enum import Enum


class CentroidInitialization(str, Enum):
    """Strategy for initializing cluster centroids.

    Mirrors TFMOT names and values where possible.
    """

    LINEAR = "linear"
    DENSITY_BASED = "density_based"
    KMEANS_PLUS_PLUS = "kmeans_plus_plus"

Semantics:

• LINEAR: centroids are evenly spaced between the minimum and maximum weight value in each group.
• DENSITY_BASED: centroids are initialized from empirical weight quantiles so that denser regions receive more centroids.
• KMEANS_PLUS_PLUS: centroids are initialized with k-means++ seeding. This is the default.

---

ClusteringConfig

from dataclasses import dataclass
from torchao.core.config import AOBaseConfig


@dataclass
class ClusteringConfig(AOBaseConfig):
    """Configuration for torchao weight clustering.

    Args:
        number_of_clusters: Maximum number of unique values per clustering
            group. Required. Mirrors TFMOT's `number_of_clusters`.
        cluster_centroids_init: Centroid initialization strategy. Mirrors
            TFMOT's `cluster_centroids_init`.
        cluster_per_channel: If False, cluster each supported weight tensor
            as a single group. If True, cluster independently per output
            channel. Mirrors TFMOT's `cluster_per_channel` control.
        preserve_sparsity: If True, preserve exactly-zero weights by assigning
            them to a zero centroid and keeping that centroid fixed at zero.
            Mirrors TFMOT's `preserve_sparsity` control.
    """

    number_of_clusters: int
    cluster_centroids_init: CentroidInitialization = (
        CentroidInitialization.KMEANS_PLUS_PLUS
    )
    cluster_per_channel: bool = False
    preserve_sparsity: bool = False

    def __post_init__(self) -> None:
        if self.number_of_clusters < 1:
            raise ValueError(
                f"number_of_clusters must be >= 1, got {self.number_of_clusters}."
            )

Deliberately omitted from the initial config:

• zero_threshold: useful, but not part of TFMOT parity and can be added later if pruning workflows need it;
• reassign_every_n_steps: excluded from the MVP because assignment refresh semantics are subtle and should be benchmarked separately;
• kmeans_init_steps: excluded from the public API for the MVP to keep the initial surface close to TFMOT;
• bits: excluded because number_of_clusters is the TFMOT-compatible spelling;
• granularity: excluded because the MVP supports only per-tensor and per-output-channel clustering through the TFMOT-compatible cluster_per_channel bool.

---

cluster_weights_

from collections.abc import Mapping
from typing import Callable
import torch.nn as nn


def cluster_weights_(
    model: nn.Module,
    config: ClusteringConfig | Mapping[str, ClusteringConfig],
    filter_fn: Callable[[nn.Module, str], bool] | None = None,
) -> nn.Module:
    """Apply weight clustering wrappers to a model in place.

    Args:
        model: Pretrained model to modify in place.
        config: Either a single ClusteringConfig or a mapping from regex
            patterns to ClusteringConfig objects. Regex patterns match fully
            qualified module names.
        filter_fn: Optional predicate with signature `(module, fqn) -> bool`.
            If provided, clustering is applied only when the predicate returns
            True. The argument order follows torchao `quantize_` and
            `sparsify_` conventions.

    Returns:
        The same model object, modified in place.
    """

Selection rules:

1. If config is a single ClusteringConfig, it applies to all supported modules selected by filter_fn.
2. If config is a mapping, keys are regex patterns over fully qualified module names and values are per-pattern configs.
3. If both a mapping and filter_fn are provided, a module must match both the regex selection and the predicate.
4. Unsupported modules are skipped by default unless a user explicitly requests strict mode in a future API.

Example:

from torchao.prototype.clustering import (
    ClusteringConfig,
    CentroidInitialization,
    cluster_weights_,
)

cfg = ClusteringConfig(
    number_of_clusters=16,
    cluster_centroids_init=CentroidInitialization.KMEANS_PLUS_PLUS,
    cluster_per_channel=True,
)

cluster_weights_(
    model,
    cfg,
    filter_fn=lambda module, fqn: isinstance(module, (nn.Conv2d, nn.Linear)),
)

Per-module config example:

cluster_weights_(
    model,
    {
        r"features\..*": ClusteringConfig(
            number_of_clusters=16,
            cluster_per_channel=True,
        ),
        r"classifier\..*": ClusteringConfig(
            number_of_clusters=8,
            cluster_per_channel=False,
        ),
    },
)

---

cluster_weights

def cluster_weights(
    model: nn.Module,
    config: ClusteringConfig | Mapping[str, ClusteringConfig],
    filter_fn: Callable[[nn.Module, str], bool] | None = None,
) -> nn.Module:
    """Return a clustered deep copy of `model`.

    This is the non-in-place counterpart to `cluster_weights_`.
    """

---

strip_clustering

def strip_clustering(
    model: nn.Module,
    *,
    inplace: bool = False,
    validate: bool = True,
) -> nn.Module:
    """Remove clustering wrappers and bake clustered weights into plain modules.

    Args:
        model: Model containing clustering wrappers.
        inplace: If True, mutate `model`. If False, operate on a deep copy.
        validate: If True, assert that every stripped clustered weight group
            has at most `number_of_clusters` unique values.

    Returns:
        A model with ordinary PyTorch modules and dense weight tensors.
    """

strip_clustering() is the required boundary before export, serialization for deployment, or standard post-training quantization flows. The stripped model should contain no clustering wrappers, no centroid parameters, and no assignment buffers.

Validation is strict. A gather from N centroids should not produce more than N unique values per group. Validation failures should therefore be treated as implementation bugs, wrong grouping, missed wrappers, or post-strip mutation.

---

validate_clustering

def validate_clustering(model: nn.Module) -> dict[str, object]:
    """Return a per-layer report of clustered-weight uniqueness.

    The report includes, for each clustered or stripped layer:
        - fully qualified module name;
        - configured number of clusters;
        - pass/fail status.
    """

This utility is intended for tests, debugging, and tutorials. It should not be required in the common path because strip_clustering(validate=True) performs the main check.

---

Supported layer scope

The MVP supports weight clustering for:

Module	Parameter	Per-channel grouping
nn.Linear	weight	output features, dim 0
nn.Conv1d	weight	output channels, dim 0
nn.Conv2d	weight	output channels, dim 0
nn.Conv3d	weight	output channels, dim 0

The MVP does not cluster bias by default. Bias tensors are small, and clustering bias may harm accuracy. A future extension can add explicit parameter-level targeting for advanced users.

ConvTranspose* is excluded from the MVP because output-channel semantics are more complex, especially with grouped transposed convolutions. It can be added later with an explicit registry and tests.

---

User workflow

Basic workflow

import torch
import torch.nn as nn

from torchao.prototype.clustering import (
    ClusteringConfig,
    CentroidInitialization,
    cluster_weights_,
    strip_clustering,
)

# 1. Start from a pretrained model.
model = MyModel()
model.load_state_dict(torch.load("pretrained.pth"))

# 2. Apply clustering wrappers.
cluster_weights_(
    model,
    ClusteringConfig(
        number_of_clusters=16,
        cluster_centroids_init=CentroidInitialization.KMEANS_PLUS_PLUS,
        cluster_per_channel=True,
    ),
    filter_fn=lambda module, fqn: isinstance(module, (nn.Conv2d, nn.Linear)),
)

# 3. Fine-tune normally.
model.train()
opt = torch.optim.Adam(model.parameters(), lr=1e-5)

for x, y in train_loader:
    opt.zero_grad(set_to_none=True)
    loss = criterion(model(x), y)
    loss.backward()
    opt.step()

# 4. Strip wrappers back to ordinary PyTorch modules.
model = strip_clustering(model, validate=True)

# 5. Continue with normal PyTorch export, serialization, or quantization.
model.eval()

Sparsity-preserving workflow

cluster_weights_(
    model,
    ClusteringConfig(
        number_of_clusters=16,
        cluster_per_channel=True,
        preserve_sparsity=True,
    ),
)

When preserve_sparsity=True, exactly-zero weights at clustering initialization are assigned to a dedicated zero centroid. That centroid is kept at exactly zero during training and after stripping.

This preserves exact zeros only. The MVP does not include a near-zero threshold.

---

Technical design

Core representation

For each supported weight tensor, clustering reshapes the weight into grouped layout:

W_grouped: [G, K]
centroids: [G, N]
assignments: [G, K]

where:

• G = 1 for per-tensor clustering;
• G = out_channels for per-channel clustering;
• K is the number of weights per group;
• N = number_of_clusters.

The clustered weight is reconstructed as:

W_clustered[g, k] = centroids[g, assignments[g, k]]

The reconstructed grouped tensor is reshaped back to the original weight shape and used for the wrapped module's forward pass.

Handling small groups

Some groups may contain fewer unique values than number_of_clusters; for example, a depthwise 3x3 convolution channel has only 9 weights. The implementation should allow unused or duplicate centroids rather than failing.

The contract is:

unique_values_per_group <= number_of_clusters

not:

unique_values_per_group == number_of_clusters

Wrapper design

ClusterWeights is an nn.Module wrapper around a supported layer. It owns:

• trainable centroid parameters;
• non-trainable assignment buffers;
• the wrapped module and its non-clustered parameters, such as bias.

The original weight parameter is frozen while clustering is active. The effective weight used in forward is reconstructed from centroids and assignments.

The MVP should prioritize correctness over memory minimization. Follow-up work can reduce persistent memory if needed.

Forward pass

The forward pass:

1. reconstructs the clustered weight with a gather from centroids;
2. reshapes it to the original weight shape;
3. calls the wrapped layer with the reconstructed weight;
4. does not update assignments;
5. does not mutate model state except for ordinary autograd bookkeeping.

Gradient behavior

Gradients flow from the reconstructed weight back to centroids. For aggregation, gradients assigned to the same centroid are accumulated.

class _ClusterGather(torch.autograd.Function):
    @staticmethod
    def forward(ctx, centroids, assignments, counts, use_avg: bool):
        ctx.save_for_backward(assignments, counts)
        ctx.centroid_shape = centroids.shape
        return torch.gather(centroids, 1, assignments)

    @staticmethod
    def backward(ctx, grad_out):
        assignments, counts = ctx.saved_tensors
        grad_centroids = grad_out.new_zeros(ctx.centroid_shape)
        grad_centroids.scatter_add_(1, assignments, grad_out)
        return grad_centroids, None, None, None

We could use native torch.gather backward.

Assignment dtype

torch.gather uses integer indices. The simplest implementation stores assignments as torch.long buffers. This is straightforward but memory-heavy.

A memory-optimized implementation may store assignments compactly, for example as uint8, int16, or int32 depending on number_of_clusters, and cast to torch.long only for the gather. This optimization should not affect the public API.

The RFC does not claim that clustering reduces training-time memory. Training-time clustering adds state. The benefit is the constrained value set after stripping.

Centroid initialization

All initialization runs independently per clustering group.

LINEAR

centroids = linspace(min(W_group), max(W_group), N)

DENSITY_BASED

Initialize centroids from empirical quantiles of the group. This places more centroids in denser regions of the weight distribution.

KMEANS_PLUS_PLUS

Use k-means++ seeding over scalar weight values in each group.

Preserving sparsity

When preserve_sparsity=True:

1. Each group reserves one centroid for zero.
2. Weights that are exactly zero at clustering initialization are assigned to that centroid.
3. The zero centroid is forced to exactly 0.0 in the effective forward weight.
4. Gradients to the zero centroid are masked.
5. strip_clustering() writes exact zeros for those positions.

The MVP preserves exact zeros only. It does not classify near-zero weights as zero.

No assignment refresh in MVP

Assignments are fixed after initialization. This mirrors the TFMOT mental model and keeps the initial implementation predictable.

Periodic reassignment is intentionally excluded because its correct semantics are not obvious. Possible future variants include reassignment against frozen original weights, reassignment against latent full-precision shadow weights, or full online k-means. Each has different optimization behavior and should be evaluated separately before becoming public API.

State dict behavior

A clustered model's state dict should include enough information to resume clustering fine-tuning:

• centroid parameters;
• assignment buffers;
• wrapped module parameters not controlled by clustering, such as bias;
• clustering metadata needed to strip or validate.

A stripped model's state dict should look like a normal PyTorch model state dict and should not contain clustering-specific centroids or assignments.

Export behavior

Users should call strip_clustering() before export. The exported model should not contain clustering wrappers, gathers, centroid parameters, or assignment buffers.

This RFC does not require torch.export to understand clustering wrappers.

---

Composability

With pruning and sparsity

Weight clustering can be applied after a sparsity or pruning pass. In that workflow, preserve_sparsity=True pins one centroid in each clustering group to exactly zero and assigns zero-valued weights to that centroid. This prevents clustering fine-tuning from moving previously pruned zeros to nonzero centroid values.

The intended ordering is:

1. train baseline model;
2. apply pruning/sparsity workflow;
3. fine-tune to recover accuracy;
4. make pruning permanent, if the pruning workflow requires a strip/squash step;
5. apply clustering with preserve_sparsity=True;
6. fine-tune with clustering active;
7. call strip_clustering(validate=True).

The key requirement is that the weights intended to remain sparse are exactly zero when clustering is initialized.

With post-training quantization

Recommended order:

1. train baseline model;
2. apply clustering;
3. fine-tune;
4. strip clustering;
5. apply PTQ or another standard quantization flow.

PTQ can change the final set of unique values, especially when quantization scales are shared across groups differently from clustering groups. Therefore, there are two distinct validation questions:

1. Does the stripped floating-point model satisfy the clustering contract?
2. Does the quantized model retain the intended number of unique integer values per group?

This RFC implements only the first validation as part of strip_clustering(validate=True). Quantized uniqueness reporting can be added as a debugging utility, but it is not part of the initial API contract.

Future work: CQAT

CQAT means cluster-preserving quantization-aware training.

Standard quantization-aware training simulates quantization effects during training, usually by inserting fake-quantization operations. This helps the model adapt to the quantization grid before conversion to an actual quantized representation.

However, if ordinary QAT is applied after clustering without preserving the clustering constraint, weights can drift independently and the cluster structure can be lost. In other words, the model may remain quantization-aware but no longer have the intended small number of shared weight values.

CQAT combines the two constraints:

centroids + assignments -> reconstructed clustered weight -> fake quantization -> layer forward

Only the centroid values are trained, while assignments preserve the weight-sharing structure. Fake quantization exposes those centroids to quantization effects during training. This is useful when:

• post-training quantization after clustering causes unacceptable accuracy loss;
• the target deployment format requires quantized weights;
• preserving the cluster structure after quantization is important for compression or analysis.

TFMOT includes cluster-preserving QAT as part of its collaborative optimization workflows. torchao should consider an equivalent in a separate RFC after the base clustering prototype is implemented and validated.

The initial clustering RFC does not add CQAT APIs, does not modify torchao QAT configs, and does not define a combined clustering-plus-QAT prepare path.

---

Validation and debugging

Strip-time validation

strip_clustering(validate=True) checks each clustered weight after stripping:

• per-tensor config: global unique count must be <= number_of_clusters;
• per-channel config: each output channel's unique count must be <= number_of_clusters.

Failure should raise ClusteringValidationError with a per-layer report.

Debug report

validate_clustering(model) returns a structured report that can be printed or consumed by tests.

Possible report schema:

{
    "features.0": {
        "module_type": "Conv2d",
        "number_of_clusters": 16,
        "cluster_per_channel": True,
        "max_unique_values_per_group": 16,
        "passed": True,
    },
}

---

Test plan

All provided functionality will be covered by unit tests, with integration and smoke tests added to validate the end-to-end clustering workflow.

---

Documentation plan

Add:

1. API reference for torchao.prototype.clustering.
2. A tutorial: "Weight clustering a pretrained PyTorch model".
3. A migration note for users familiar with TFMOT:

TFMOT concept	torchao proposal
cluster_weights()	cluster_weights() / cluster_weights_()
strip_clustering()	strip_clustering()
number_of_clusters	number_of_clusters
cluster_centroids_init	cluster_centroids_init
CentroidInitialization	CentroidInitialization
cluster_per_channel	cluster_per_channel
preserve_sparsity	preserve_sparsity
Keras **kwargs	no equivalent; not needed in PyTorch

The documentation should state clearly that compression benefits are deployment-format-dependent. The clustering contract is about the number of unique weight values, not a guaranteed file-size ratio.

Open questions

Stable package location

Where should the stable API live after graduation?

Options:

• torchao.quantization.clustering: emphasizes common composition with PTQ/QAT.
• torchao.sparsity.clustering: emphasizes weight-domain compression and relationship to pruning.
• torchao.clustering: avoids categorizing clustering as either quantization or sparsity.

---

Implementation plan

PR 1: prototype core

Repository: pytorch/ao

Contents:

• ClusteringConfig;
• CentroidInitialization;
• ClusterGradientAggregation;
• centroid initialization utilities;
• ClusterWeights wrapper;
• cluster_weights_();
• cluster_weights();
• strip_clustering();
• validate_clustering();
• unit tests for supported layers, validation, state dict, and sparsity preservation.

PR 2: documentation and tutorial

Repository: pytorch/ao

Contents:

• prototype API docs;
• TFMOT migration guide;
• tutorial notebook or script;
• debugging examples.

PR 3: benchmark and graduation criteria

Repository: pytorch/ao

Contents:

• slow benchmark scripts;
• model-family coverage;
• documented accuracy behavior;

Future separate RFC: CQAT

A separate RFC should define cluster-preserving quantization-aware training in torchao.

---

References

• TFMOT cluster_weights API: https://www.tensorflow.org/model_optimization/api_docs/python/tfmot/clustering/keras/cluster_weights
• TFMOT weight clustering guide: https://www.tensorflow.org/model_optimization/guide/clustering
• TFMOT comprehensive clustering guide: https://www.tensorflow.org/model_optimization/guide/clustering/clustering_comprehensive_guide
• TFMOT cluster-preserving QAT example: https://www.tensorflow.org/model_optimization/guide/combine/cqat_example
• torchao sparsify_ API: https://docs.pytorch.org/ao/stable/api_reference/generated/torchao.sparsity.sparsify_.html
• torchao quantize_ API: https://docs.pytorch.org/ao/stable/api_reference/generated/torchao.quantization.quantize_.html
• torchao AOBaseConfig: https://docs.pytorch.org/ao/stable/api_reference/generated/torchao.core.config.AOBaseConfig.html

Proposed API

Target hardware

None