Compressing large neural networks with minimal performance loss is crucial to
enabling their deployment on edge devices. (Cho et al., 2022) proposed a weight
quantization method that uses an attention-based clustering algorithm called
differentiable $k$-means (DKM). Despite achieving state-of-the-art results,
DKM's performance is constrained by its heavy memory dependency. We propose an
implicit, differentiable $k$-means algorithm (IDKM), which eliminates the major
memory restriction of DKM. Let $t$ be the number of $k$-means iterations, $m$
be the number of weight-vectors, and $b$ be the number of bits per cluster
address. IDKM reduces the overall memory complexity of a single $k$-means layer
from $\mathcal{O}(t \cdot m \cdot 2^b)$ to $\mathcal{O}( m \cdot 2^b)$. We also
introduce a variant, IDKM with Jacobian-Free-Backpropagation (IDKM-JFB), for
which the time complexity of the gradient calculation is independent of $t$ as
well. We provide a proof of concept of our methods by showing that, under the
same settings, IDKM achieves comparable performance to DKM with less compute
time and less memory. We also use IDKM and IDKM-JFB to quantize a large neural
network, Resnet18, on hardware where DKM cannot train at all.

使用隐式可微分 k-means 算法（IDKM）和无雅可比传播（IDKM-JFB）的变种，在与同样设置下，IDKM 在更短的计算时间和较少的内存使用下实现与 DKM 相当的性能，并将其应用于无法在 DKM 上训练的 Resnet18 神经网络。