Skip to content

Quantization-Aware Training — STE for Hardware Deployment

Train SNNs through quantization using straight-through estimators (STE). The missing link between training and FPGA deployment: weights are quantized in the forward pass but maintain full precision in the backward pass.

How STE Works

Standard quantization is non-differentiable (rounding has zero gradient almost everywhere). The straight-through estimator passes the gradient through quantization as if it weren't there:

  • Forward: W_q = round(W / scale) * scale (quantized)
  • Backward: ∂L/∂W = ∂L/∂W_q (identity, as if no quantization)

This trains weights to be robust to their own quantization noise. At export time, weights are already at target precision.

Components

  • QuantizedSNNLayer — SNN layer with quantization-aware forward pass.
Parameter Default Meaning
n_inputs (required) Input dimension
n_neurons (required) Output dimension
weight_bits 8 Target weight precision (2, 4, 8, 16)
threshold 1.0 LIF spike threshold
tau_mem 20.0 Membrane time constant
  • TernaryWeights — Ternary quantization: {-1, 0, +1}. 94% memory reduction. Weights with |w| < threshold_ratio * mean(|w|) become 0.
  • quantize_aware_train_step — One QAT training step with STE gradient flow. Returns {'output', 'loss'}.
  • _ste_quantize — Core quantization function. Supports symmetric and asymmetric modes.

Learned quantisers (PyTorch)

Higher-accuracy quantisers that learn their parameters during training instead of fixing them from the running range:

  • LSQLinear / LSQQuantizer — Learned Step Size Quantization (Esser et al. 2020). The quantiser step size is a trainable parameter, per-tensor or per-output-channel, learned jointly with the weights. export_quantized() returns integer codes plus the learned step(s).
  • PACTActivation — PArameterized Clipping acTivation (Choi et al. 2018). A learnable clipping bound alpha bounds the activation range before uniform quantisation, so low-bit activations no longer need a hand-tuned clip.
  • MinMaxObserver / PerChannelMinMaxObserver — Range observers that turn calibration statistics into (scale, zero_point), per-tensor or per-channel. Per-channel weight scales recover the accuracy a single per-tensor scale loses across channels of differing magnitude.
  • fake_quantize — Quantise/de-quantise helper (no STE) used to evaluate observer scales.
  • LSQPACTLIFNet — Feedforward LIF SNN wiring LSQ per-channel weight quantisation and a PACT-quantised analogue input end to end.

Usage

Python
from sc_neurocore.qat import QuantizedSNNLayer, quantize_aware_train_step, TernaryWeights
import numpy as np

# Create QAT layer
layer = QuantizedSNNLayer(n_inputs=784, n_neurons=128, weight_bits=8)

# Training loop with STE
for epoch in range(100):
    result = quantize_aware_train_step(layer, x_train, y_target, lr=0.01)
    print(f"Loss: {result['loss']:.4f}")

# Export hardware-ready weights (already quantized to 8-bit)
hw_weights = layer.export_weights()

# Ternary quantization for extreme compression
tw = TernaryWeights(threshold_ratio=0.7)
ternary = tw.quantize(layer.W)
print(f"Sparsity: {tw.sparsity(layer.W):.1%}")  # ~50-70% zeros

References: QP-SNN (ICLR 2025), SpikeFit (EurIPS 2025).

See Tutorial 77: QAT.

sc_neurocore.qat.quantize

Train SNNs through quantization using straight-through estimators.

Missing link between training and hardware deployment. No SNN library ships QAT as a reusable module.

Reference: QP-SNN (ICLR 2025), SpikeFit (EurIPS 2025)

QuantizedSNNLayer dataclass

SNN layer with quantization-aware forward pass.

During training: weights quantized in forward, full-precision in backward (STE). At export: weights are already at target precision.

Parameters

n_inputs : int n_neurons : int weight_bits : int Target weight precision (2, 4, 8, 16). threshold : float tau_mem : float

Source code in src/sc_neurocore/qat/quantize.py
Python
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
@dataclass
class QuantizedSNNLayer:
    """SNN layer with quantization-aware forward pass.

    During training: weights quantized in forward, full-precision in backward (STE).
    At export: weights are already at target precision.

    Parameters
    ----------
    n_inputs : int
    n_neurons : int
    weight_bits : int
        Target weight precision (2, 4, 8, 16).
    threshold : float
    tau_mem : float
    """

    n_inputs: int
    n_neurons: int
    weight_bits: int = 8
    threshold: float = 1.0
    tau_mem: float = 20.0

    def __post_init__(self) -> None:
        rng = np.random.RandomState(42)
        self.W = rng.randn(self.n_neurons, self.n_inputs) * np.sqrt(2.0 / self.n_inputs)
        self._v = np.zeros(self.n_neurons)

    def forward(self, x: np.ndarray[Any, Any], dt: float = 1.0) -> np.ndarray[Any, Any]:
        """Quantization-aware forward pass."""
        W_q = _ste_quantize(self.W, self.weight_bits)
        alpha = np.exp(-dt / self.tau_mem)
        current = W_q @ x
        self._v = alpha * self._v + (1 - alpha) * current
        spikes = (self._v >= self.threshold).astype(np.float64)
        self._v -= spikes * self.threshold
        return spikes

    def export_weights(self) -> np.ndarray[Any, Any]:
        """Export quantized weights for hardware deployment."""
        return _ste_quantize(self.W, self.weight_bits)

    def reset(self) -> None:  # pragma: no cover
        self._v = np.zeros(self.n_neurons)

forward(x, dt=1.0)

Quantization-aware forward pass.

Source code in src/sc_neurocore/qat/quantize.py
Python
 97
 98
 99
100
101
102
103
104
105
def forward(self, x: np.ndarray[Any, Any], dt: float = 1.0) -> np.ndarray[Any, Any]:
    """Quantization-aware forward pass."""
    W_q = _ste_quantize(self.W, self.weight_bits)
    alpha = np.exp(-dt / self.tau_mem)
    current = W_q @ x
    self._v = alpha * self._v + (1 - alpha) * current
    spikes = (self._v >= self.threshold).astype(np.float64)
    self._v -= spikes * self.threshold
    return spikes

export_weights()

Export quantized weights for hardware deployment.

Source code in src/sc_neurocore/qat/quantize.py
Python
107
108
109
def export_weights(self) -> np.ndarray[Any, Any]:
    """Export quantized weights for hardware deployment."""
    return _ste_quantize(self.W, self.weight_bits)

TernaryWeights

Ternary weight quantization: {-1, 0, +1}.

94% memory reduction. Each weight is one of three values. Threshold-based: weights with |w| < threshold become 0.

Parameters

threshold_ratio : float Fraction of max(|w|) below which weights are zeroed.

Source code in src/sc_neurocore/qat/quantize.py
Python
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
class TernaryWeights:
    """Ternary weight quantization: {-1, 0, +1}.

    94% memory reduction. Each weight is one of three values.
    Threshold-based: weights with |w| < threshold become 0.

    Parameters
    ----------
    threshold_ratio : float
        Fraction of max(|w|) below which weights are zeroed.
    """

    def __init__(self, threshold_ratio: float = 0.7):
        self.threshold_ratio = threshold_ratio

    def quantize(self, weights: np.ndarray[Any, Any]) -> np.ndarray[Any, Any]:
        threshold = self.threshold_ratio * np.mean(np.abs(weights))
        ternary = np.zeros_like(weights)
        ternary[weights > threshold] = 1.0
        ternary[weights < -threshold] = -1.0
        return ternary

    def sparsity(self, weights: np.ndarray[Any, Any]) -> float:
        t = self.quantize(weights)
        return float(np.mean(t == 0))

quantize_aware_train_step(layer, x, target, lr=0.01)

One QAT training step with STE.

Parameters

layer : QuantizedSNNLayer x : ndarray of shape (n_inputs,) target : ndarray of shape (n_neurons,) lr : float

Returns

dict with 'output', 'loss'

Source code in src/sc_neurocore/qat/quantize.py
Python
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
def quantize_aware_train_step(
    layer: QuantizedSNNLayer,
    x: np.ndarray[Any, Any],
    target: np.ndarray[Any, Any],
    lr: float = 0.01,
) -> dict[str, object]:
    """One QAT training step with STE.

    Parameters
    ----------
    layer : QuantizedSNNLayer
    x : ndarray of shape (n_inputs,)
    target : ndarray of shape (n_neurons,)
    lr : float

    Returns
    -------
    dict with 'output', 'loss'
    """
    output = layer.forward(x)
    error = output - target
    loss = 0.5 * float(np.sum(error**2))

    # STE: gradient flows through quantization as if it weren't there
    grad_W = np.outer(error, x)
    layer.W -= lr * grad_W

    return {"output": output, "loss": loss}

Learned Step Size Quantization

sc_neurocore.qat.lsq

Learned Step Size Quantization for quantisation-aware training.

LSQ makes the quantiser's step size s a trainable parameter learned jointly with the weights, rather than a fixed function of the running weight range. The forward pass quantises v onto the integer grid [qmin, qmax] at step s; the backward pass propagates a gradient to s itself:

Text Only
∂v̂/∂s = round(v/s) - v/s        for qmin < v/s < qmax
∂v̂/∂s = qmin                     for v/s <= qmin
∂v̂/∂s = qmax                     for v/s >= qmax

and a straight-through gradient of 1 to v inside the clip range (0 outside). The step-size gradient is scaled by 1 / sqrt(qmax * n) (n = elements per step) so its magnitude matches the weight gradients during joint optimisation. The step is initialised to 2 * mean(|v|) / sqrt(qmax).

A single scalar step gives per-tensor quantisation; one step per output channel gives per-channel quantisation, which pairs naturally with the per-channel observers in :mod:sc_neurocore.qat.observers.

Reference: Esser et al. 2020 — "Learned Step Size Quantization" (ICLR).

LSQLinear

Bases: Module

Linear layer whose weights are quantised by a learned step size.

Per-channel (per output neuron) quantisation is the default, matching the granularity that recovers most of the accuracy lost to low-bit weights.

Parameters

in_features, out_features : int Layer dimensions. n_bits : int Weight quantiser bit width. per_channel : bool Learn one step per output neuron (default) or a single scalar step. bias : bool Whether to include a full-precision bias.

Source code in src/sc_neurocore/qat/lsq.py
Python
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
class LSQLinear(nn.Module):
    """Linear layer whose weights are quantised by a learned step size.

    Per-channel (per output neuron) quantisation is the default, matching the
    granularity that recovers most of the accuracy lost to low-bit weights.

    Parameters
    ----------
    in_features, out_features : int
        Layer dimensions.
    n_bits : int
        Weight quantiser bit width.
    per_channel : bool
        Learn one step per output neuron (default) or a single scalar step.
    bias : bool
        Whether to include a full-precision bias.
    """

    def __init__(
        self,
        in_features: int,
        out_features: int,
        *,
        n_bits: int = 8,
        per_channel: bool = True,
        bias: bool = True,
    ) -> None:
        super().__init__()
        self.linear = nn.Linear(in_features, out_features, bias=bias)
        self.weight_quant = LSQQuantizer(
            n_bits,
            per_channel=per_channel,
            ch_axis=0,
            num_channels=out_features if per_channel else None,
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Apply the layer with LSQ-quantised weights."""
        w_q = self.weight_quant(self.linear.weight)
        return nn.functional.linear(x, w_q, self.linear.bias)

    def export_quantized(self) -> dict[str, Any]:
        """Export integer weights, the learned step(s), and the bias.

        Returns
        -------
        dict
            ``weight_int`` (int32 codes), ``step`` (per-tensor or per-channel),
            ``n_bits``, ``per_channel``, and optionally ``bias``.
        """
        codes, step = self.weight_quant.integer_weights(self.linear.weight.detach())
        result: dict[str, Any] = {
            "weight_int": codes,
            "step": step,
            "n_bits": self.weight_quant.n_bits,
            "per_channel": self.weight_quant.per_channel,
        }
        if self.linear.bias is not None:
            result["bias"] = self.linear.bias.detach()
        return result

forward(x)

Apply the layer with LSQ-quantised weights.

Source code in src/sc_neurocore/qat/lsq.py
Python
275
276
277
278
def forward(self, x: torch.Tensor) -> torch.Tensor:
    """Apply the layer with LSQ-quantised weights."""
    w_q = self.weight_quant(self.linear.weight)
    return nn.functional.linear(x, w_q, self.linear.bias)

export_quantized()

Export integer weights, the learned step(s), and the bias.

Returns

dict weight_int (int32 codes), step (per-tensor or per-channel), n_bits, per_channel, and optionally bias.

Source code in src/sc_neurocore/qat/lsq.py
Python
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
def export_quantized(self) -> dict[str, Any]:
    """Export integer weights, the learned step(s), and the bias.

    Returns
    -------
    dict
        ``weight_int`` (int32 codes), ``step`` (per-tensor or per-channel),
        ``n_bits``, ``per_channel``, and optionally ``bias``.
    """
    codes, step = self.weight_quant.integer_weights(self.linear.weight.detach())
    result: dict[str, Any] = {
        "weight_int": codes,
        "step": step,
        "n_bits": self.weight_quant.n_bits,
        "per_channel": self.weight_quant.per_channel,
    }
    if self.linear.bias is not None:
        result["bias"] = self.linear.bias.detach()
    return result

LSQQuantizer

Bases: Module

Learned-step-size fake quantiser for signed weights.

Parameters

n_bits : int Quantiser bit width (>= 2). The signed grid is [-2**(n_bits-1), 2**(n_bits-1) - 1]. per_channel : bool Learn one step per channel along ch_axis instead of a single scalar step. ch_axis : int Channel axis used when per_channel is set. num_channels : int, optional Channel count; required when per_channel is set.

Attributes

step : torch.nn.Parameter The learned step size(s). Lazily initialised from the first input.

Source code in src/sc_neurocore/qat/lsq.py
Python
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
class LSQQuantizer(nn.Module):
    """Learned-step-size fake quantiser for signed weights.

    Parameters
    ----------
    n_bits : int
        Quantiser bit width (``>= 2``). The signed grid is
        ``[-2**(n_bits-1), 2**(n_bits-1) - 1]``.
    per_channel : bool
        Learn one step per channel along ``ch_axis`` instead of a single
        scalar step.
    ch_axis : int
        Channel axis used when ``per_channel`` is set.
    num_channels : int, optional
        Channel count; required when ``per_channel`` is set.

    Attributes
    ----------
    step : torch.nn.Parameter
        The learned step size(s). Lazily initialised from the first input.
    """

    step: nn.Parameter
    _initialized: torch.Tensor

    def __init__(
        self,
        n_bits: int = 8,
        *,
        per_channel: bool = False,
        ch_axis: int = 0,
        num_channels: int | None = None,
    ) -> None:
        super().__init__()
        if n_bits < 2:
            raise ValueError(f"n_bits must be >= 2, got {n_bits}")
        if per_channel and num_channels is None:
            raise ValueError("per_channel quantisation requires num_channels")
        self.n_bits = n_bits
        self.per_channel = per_channel
        self.ch_axis = ch_axis
        self.qmin = -(1 << (n_bits - 1))
        self.qmax = (1 << (n_bits - 1)) - 1
        if per_channel:
            assert num_channels is not None  # validated above
            init = torch.ones(num_channels)
        else:
            init = torch.ones(())
        self.step = nn.Parameter(init)
        self.register_buffer("_initialized", torch.zeros((), dtype=torch.bool))

    def _init_step_from(self, x: torch.Tensor) -> None:
        """Initialise the step to ``2*mean(|x|)/sqrt(qmax)`` (Esser et al. 2020)."""
        with torch.no_grad():
            if self.per_channel:
                axis = self.ch_axis % x.dim()
                moved = x.detach().movedim(axis, 0).reshape(x.shape[axis], -1)
                mean_abs = moved.abs().mean(dim=1)
            else:
                mean_abs = x.detach().abs().mean()
            self.step.copy_(2.0 * mean_abs / math.sqrt(self.qmax))
            self._initialized.fill_(True)

    def _broadcast_step(self, ndim: int) -> torch.Tensor:
        """Reshape the per-channel step for broadcasting over ``ndim`` dims."""
        if not self.per_channel:
            return self.step
        axis = self.ch_axis % ndim
        shape = [1] * ndim
        shape[axis] = -1
        return self.step.reshape(shape)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Fake-quantise ``x`` at the learned step, with the LSQ gradient.

        Parameters
        ----------
        x : torch.Tensor
            Full-precision weights to quantise.

        Returns
        -------
        torch.Tensor
            The quantised-dequantised weights, differentiable in both ``x`` and
            the step size.
        """
        if not bool(self._initialized):
            self._init_step_from(x)
        n_elements = x.numel() // self.step.numel()
        grad_scale = 1.0 / math.sqrt(self.qmax * max(n_elements, 1))
        step = self._broadcast_step(x.dim()).clamp(min=1e-8)
        return cast(
            torch.Tensor,
            _LSQQuantize.apply(x, step, self.qmin, self.qmax, grad_scale),
        )

    def integer_weights(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        """Return integer codes and the step(s) for hardware export.

        Parameters
        ----------
        x : torch.Tensor
            Weights to encode at the learned step.

        Returns
        -------
        tuple of (torch.Tensor, torch.Tensor)
            Integer codes clamped to the grid, and the per-tensor or
            per-channel step size(s).
        """
        step = self._broadcast_step(x.dim()).clamp(min=1e-8)
        codes = torch.round(x / step).clamp(self.qmin, self.qmax)
        return codes.to(torch.int32), self.step.detach()

forward(x)

Fake-quantise x at the learned step, with the LSQ gradient.

Parameters

x : torch.Tensor Full-precision weights to quantise.

Returns

torch.Tensor The quantised-dequantised weights, differentiable in both x and the step size.

Source code in src/sc_neurocore/qat/lsq.py
Python
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
def forward(self, x: torch.Tensor) -> torch.Tensor:
    """Fake-quantise ``x`` at the learned step, with the LSQ gradient.

    Parameters
    ----------
    x : torch.Tensor
        Full-precision weights to quantise.

    Returns
    -------
    torch.Tensor
        The quantised-dequantised weights, differentiable in both ``x`` and
        the step size.
    """
    if not bool(self._initialized):
        self._init_step_from(x)
    n_elements = x.numel() // self.step.numel()
    grad_scale = 1.0 / math.sqrt(self.qmax * max(n_elements, 1))
    step = self._broadcast_step(x.dim()).clamp(min=1e-8)
    return cast(
        torch.Tensor,
        _LSQQuantize.apply(x, step, self.qmin, self.qmax, grad_scale),
    )

integer_weights(x)

Return integer codes and the step(s) for hardware export.

Parameters

x : torch.Tensor Weights to encode at the learned step.

Returns

tuple of (torch.Tensor, torch.Tensor) Integer codes clamped to the grid, and the per-tensor or per-channel step size(s).

Source code in src/sc_neurocore/qat/lsq.py
Python
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
def integer_weights(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
    """Return integer codes and the step(s) for hardware export.

    Parameters
    ----------
    x : torch.Tensor
        Weights to encode at the learned step.

    Returns
    -------
    tuple of (torch.Tensor, torch.Tensor)
        Integer codes clamped to the grid, and the per-tensor or
        per-channel step size(s).
    """
    step = self._broadcast_step(x.dim()).clamp(min=1e-8)
    codes = torch.round(x / step).clamp(self.qmin, self.qmax)
    return codes.to(torch.int32), self.step.detach()

PACT Activation

sc_neurocore.qat.pact

PACT: PArameterized Clipping acTivation for quantisation-aware training.

PACT replaces the unbounded ReLU with a clipping activation whose upper bound alpha is a trainable parameter, so the network learns the activation range that minimises quantisation error instead of relying on a fixed clip:

Text Only
y = clip(x, 0, alpha)          (parameterised clip)
y_q = round(y / s) * s,  s = alpha / (2**n_bits - 1)   (uniform quantise)

The clip gradient flows to alpha only where the input saturates the upper bound (x > alpha), which is the PACT contribution; rounding is handled with a straight-through estimator. Clipping the activation range is what makes low bit-width activation quantisation viable, complementing the learned-step weight quantiser in :mod:sc_neurocore.qat.lsq.

Reference: Choi et al. 2018 — "PACT: Parameterized Clipping Activation for Quantized Neural Networks".

PACTActivation

Bases: Module

Parameterised clipping activation with uniform quantisation.

Parameters

n_bits : int Activation quantiser bit width (>= 2). The activation grid has 2**n_bits - 1 positive levels over [0, alpha]. alpha_init : float Initial clipping bound.

Attributes

alpha : torch.nn.Parameter The learned clipping bound.

Source code in src/sc_neurocore/qat/pact.py
Python
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
class PACTActivation(nn.Module):
    """Parameterised clipping activation with uniform quantisation.

    Parameters
    ----------
    n_bits : int
        Activation quantiser bit width (``>= 2``). The activation grid has
        ``2**n_bits - 1`` positive levels over ``[0, alpha]``.
    alpha_init : float
        Initial clipping bound.

    Attributes
    ----------
    alpha : torch.nn.Parameter
        The learned clipping bound.
    """

    def __init__(self, n_bits: int = 8, alpha_init: float = 6.0) -> None:
        super().__init__()
        if n_bits < 2:
            raise ValueError(f"n_bits must be >= 2, got {n_bits}")
        self.n_bits = n_bits
        self.n_levels = (1 << n_bits) - 1
        self.alpha = nn.Parameter(torch.tensor(float(alpha_init)))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Clip ``x`` to ``[0, alpha]`` and quantise to ``n_bits`` levels.

        Parameters
        ----------
        x : torch.Tensor
            Pre-activation values.

        Returns
        -------
        torch.Tensor
            Clipped, quantised activations, differentiable in ``x`` and
            ``alpha``.
        """
        y = cast(torch.Tensor, _PACTClip.apply(x, self.alpha))
        scale = self.alpha.clamp(min=1e-8) / self.n_levels
        return _round_ste(y / scale).clamp(0, self.n_levels) * scale

    def quantize(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        """Return integer activation codes and the scale for export.

        Parameters
        ----------
        x : torch.Tensor
            Pre-activation values to encode with the learned clip.

        Returns
        -------
        tuple of (torch.Tensor, torch.Tensor)
            Integer codes in ``[0, n_levels]`` and the scalar activation scale.
        """
        with torch.no_grad():
            scale = self.alpha.clamp(min=1e-8) / self.n_levels
            y = torch.clamp(x, min=0.0).minimum(self.alpha)
            codes = torch.round(y / scale).clamp(0, self.n_levels)
            return codes.to(torch.int32), scale.detach()

    def extra_repr(self) -> str:
        """Return the compact module representation."""
        return f"n_bits={self.n_bits}, alpha={self.alpha.item():.3f}"

forward(x)

Clip x to [0, alpha] and quantise to n_bits levels.

Parameters

x : torch.Tensor Pre-activation values.

Returns

torch.Tensor Clipped, quantised activations, differentiable in x and alpha.

Source code in src/sc_neurocore/qat/pact.py
Python
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
def forward(self, x: torch.Tensor) -> torch.Tensor:
    """Clip ``x`` to ``[0, alpha]`` and quantise to ``n_bits`` levels.

    Parameters
    ----------
    x : torch.Tensor
        Pre-activation values.

    Returns
    -------
    torch.Tensor
        Clipped, quantised activations, differentiable in ``x`` and
        ``alpha``.
    """
    y = cast(torch.Tensor, _PACTClip.apply(x, self.alpha))
    scale = self.alpha.clamp(min=1e-8) / self.n_levels
    return _round_ste(y / scale).clamp(0, self.n_levels) * scale

quantize(x)

Return integer activation codes and the scale for export.

Parameters

x : torch.Tensor Pre-activation values to encode with the learned clip.

Returns

tuple of (torch.Tensor, torch.Tensor) Integer codes in [0, n_levels] and the scalar activation scale.

Source code in src/sc_neurocore/qat/pact.py
Python
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
def quantize(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
    """Return integer activation codes and the scale for export.

    Parameters
    ----------
    x : torch.Tensor
        Pre-activation values to encode with the learned clip.

    Returns
    -------
    tuple of (torch.Tensor, torch.Tensor)
        Integer codes in ``[0, n_levels]`` and the scalar activation scale.
    """
    with torch.no_grad():
        scale = self.alpha.clamp(min=1e-8) / self.n_levels
        y = torch.clamp(x, min=0.0).minimum(self.alpha)
        codes = torch.round(y / scale).clamp(0, self.n_levels)
        return codes.to(torch.int32), scale.detach()

extra_repr()

Return the compact module representation.

Source code in src/sc_neurocore/qat/pact.py
Python
121
122
123
def extra_repr(self) -> str:
    """Return the compact module representation."""
    return f"n_bits={self.n_bits}, alpha={self.alpha.item():.3f}"

Quantisation Observers

sc_neurocore.qat.observers

Range observers that turn weight/activation statistics into quantiser scales.

An observer watches tensors during calibration, tracks their value range, and converts that range into the (scale, zero_point) a uniform affine quantiser needs. Two granularities are provided:

Per-tensor One scale for the whole tensor (:class:MinMaxObserver).

Per-channel One scale per output channel along a chosen axis (:class:PerChannelMinMaxObserver). Per-channel weight quantisation absorbs the wide dynamic-range differences between filters that a single per-tensor scale would otherwise clip or under-resolve, which is the standard remedy for the accuracy loss of low-bit weight quantisation.

Both support a symmetric scheme (signed weights, zero_point == 0) and an affine scheme (arbitrary [min, max] mapped onto the integer grid). The observed range is a running min/max across every :meth:observe call, so a calibration loop can stream batches through the observer before the scales are read once via :meth:calculate_qparams.

MinMaxObserver

Bases: Module

Per-tensor running min/max range observer.

Parameters

n_bits : int Quantiser bit width the derived scale targets. symmetric : bool Use a symmetric (zero-centred) mapping — the default for weights. unsigned : bool Target an unsigned integer grid (e.g. non-negative activations). eps : float Scale floor guarding against a zero-width observed range.

Source code in src/sc_neurocore/qat/observers.py
Python
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
class MinMaxObserver(nn.Module):
    """Per-tensor running min/max range observer.

    Parameters
    ----------
    n_bits : int
        Quantiser bit width the derived scale targets.
    symmetric : bool
        Use a symmetric (zero-centred) mapping — the default for weights.
    unsigned : bool
        Target an unsigned integer grid (e.g. non-negative activations).
    eps : float
        Scale floor guarding against a zero-width observed range.
    """

    min_val: torch.Tensor
    max_val: torch.Tensor

    def __init__(
        self,
        n_bits: int = 8,
        *,
        symmetric: bool = True,
        unsigned: bool = False,
        eps: float = 1e-8,
    ) -> None:
        super().__init__()
        self.n_bits = n_bits
        self.symmetric = symmetric
        self.unsigned = unsigned
        self.eps = eps
        self.register_buffer("min_val", torch.tensor(float("inf")))
        self.register_buffer("max_val", torch.tensor(float("-inf")))

    def observe(self, x: torch.Tensor) -> torch.Tensor:
        """Fold ``x`` into the running range and return it unchanged.

        Parameters
        ----------
        x : torch.Tensor
            Calibration tensor.

        Returns
        -------
        torch.Tensor
            ``x`` unchanged, so the observer can be dropped into a forward pass.
        """
        x = x.detach()
        self.min_val = torch.minimum(self.min_val, x.min())
        self.max_val = torch.maximum(self.max_val, x.max())
        return x

    def calculate_qparams(self) -> tuple[torch.Tensor, torch.Tensor]:
        """Return the ``(scale, zero_point)`` for the observed range.

        Returns
        -------
        tuple of (torch.Tensor, torch.Tensor)
            Scalar scale and zero point.

        Raises
        ------
        RuntimeError
            If no tensor has been observed yet.
        """
        if not torch.isfinite(self.min_val) or not torch.isfinite(self.max_val):
            raise RuntimeError("MinMaxObserver.calculate_qparams called before any observation")
        return _qparams_from_range(
            self.min_val,
            self.max_val,
            n_bits=self.n_bits,
            symmetric=self.symmetric,
            unsigned=self.unsigned,
            eps=self.eps,
        )

    def quantize(self, x: torch.Tensor) -> torch.Tensor:
        """Fake-quantise ``x`` with the currently observed scale."""
        scale, zero_point = self.calculate_qparams()
        return fake_quantize(x, scale, zero_point, n_bits=self.n_bits, unsigned=self.unsigned)

observe(x)

Fold x into the running range and return it unchanged.

Parameters

x : torch.Tensor Calibration tensor.

Returns

torch.Tensor x unchanged, so the observer can be dropped into a forward pass.

Source code in src/sc_neurocore/qat/observers.py
Python
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
def observe(self, x: torch.Tensor) -> torch.Tensor:
    """Fold ``x`` into the running range and return it unchanged.

    Parameters
    ----------
    x : torch.Tensor
        Calibration tensor.

    Returns
    -------
    torch.Tensor
        ``x`` unchanged, so the observer can be dropped into a forward pass.
    """
    x = x.detach()
    self.min_val = torch.minimum(self.min_val, x.min())
    self.max_val = torch.maximum(self.max_val, x.max())
    return x

calculate_qparams()

Return the (scale, zero_point) for the observed range.

Returns

tuple of (torch.Tensor, torch.Tensor) Scalar scale and zero point.

Raises

RuntimeError If no tensor has been observed yet.

Source code in src/sc_neurocore/qat/observers.py
Python
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
def calculate_qparams(self) -> tuple[torch.Tensor, torch.Tensor]:
    """Return the ``(scale, zero_point)`` for the observed range.

    Returns
    -------
    tuple of (torch.Tensor, torch.Tensor)
        Scalar scale and zero point.

    Raises
    ------
    RuntimeError
        If no tensor has been observed yet.
    """
    if not torch.isfinite(self.min_val) or not torch.isfinite(self.max_val):
        raise RuntimeError("MinMaxObserver.calculate_qparams called before any observation")
    return _qparams_from_range(
        self.min_val,
        self.max_val,
        n_bits=self.n_bits,
        symmetric=self.symmetric,
        unsigned=self.unsigned,
        eps=self.eps,
    )

quantize(x)

Fake-quantise x with the currently observed scale.

Source code in src/sc_neurocore/qat/observers.py
Python
228
229
230
231
def quantize(self, x: torch.Tensor) -> torch.Tensor:
    """Fake-quantise ``x`` with the currently observed scale."""
    scale, zero_point = self.calculate_qparams()
    return fake_quantize(x, scale, zero_point, n_bits=self.n_bits, unsigned=self.unsigned)

PerChannelMinMaxObserver

Bases: Module

Per-channel running min/max range observer.

Tracks an independent min/max — and therefore an independent scale — for every slice along ch_axis. For a weight tensor shaped (out_features, in_features) the default ch_axis=0 yields one scale per output neuron.

Parameters

n_bits : int Quantiser bit width the derived scales target. ch_axis : int Axis whose length is the channel count. symmetric : bool Use a symmetric (zero-centred) mapping — the default for weights. unsigned : bool Target an unsigned integer grid. eps : float Scale floor guarding against a zero-width observed range.

Source code in src/sc_neurocore/qat/observers.py
Python
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
class PerChannelMinMaxObserver(nn.Module):
    """Per-channel running min/max range observer.

    Tracks an independent min/max — and therefore an independent scale — for
    every slice along ``ch_axis``. For a weight tensor shaped
    ``(out_features, in_features)`` the default ``ch_axis=0`` yields one scale
    per output neuron.

    Parameters
    ----------
    n_bits : int
        Quantiser bit width the derived scales target.
    ch_axis : int
        Axis whose length is the channel count.
    symmetric : bool
        Use a symmetric (zero-centred) mapping — the default for weights.
    unsigned : bool
        Target an unsigned integer grid.
    eps : float
        Scale floor guarding against a zero-width observed range.
    """

    min_vals: torch.Tensor
    max_vals: torch.Tensor

    def __init__(
        self,
        n_bits: int = 8,
        *,
        ch_axis: int = 0,
        symmetric: bool = True,
        unsigned: bool = False,
        eps: float = 1e-8,
    ) -> None:
        super().__init__()
        self.n_bits = n_bits
        self.ch_axis = ch_axis
        self.symmetric = symmetric
        self.unsigned = unsigned
        self.eps = eps
        self.register_buffer("min_vals", torch.empty(0))
        self.register_buffer("max_vals", torch.empty(0))

    def _per_channel_min_max(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        """Collapse every axis except ``ch_axis`` to per-channel min and max."""
        axis = self.ch_axis % x.dim()
        moved = x.movedim(axis, 0).reshape(x.shape[axis], -1)
        return moved.min(dim=1).values, moved.max(dim=1).values

    def observe(self, x: torch.Tensor) -> torch.Tensor:
        """Fold ``x`` into the running per-channel range and return it unchanged.

        Parameters
        ----------
        x : torch.Tensor
            Calibration tensor whose ``ch_axis`` length is the channel count.

        Returns
        -------
        torch.Tensor
            ``x`` unchanged.
        """
        x = x.detach()
        cur_min, cur_max = self._per_channel_min_max(x)
        if self.min_vals.numel() == 0:
            self.min_vals = cur_min
            self.max_vals = cur_max
        else:
            self.min_vals = torch.minimum(self.min_vals, cur_min)
            self.max_vals = torch.maximum(self.max_vals, cur_max)
        return x

    def calculate_qparams(self) -> tuple[torch.Tensor, torch.Tensor]:
        """Return the per-channel ``(scale, zero_point)`` vectors.

        Returns
        -------
        tuple of (torch.Tensor, torch.Tensor)
            1-D scale and zero-point tensors, one entry per channel.

        Raises
        ------
        RuntimeError
            If no tensor has been observed yet.
        """
        if self.min_vals.numel() == 0:
            raise RuntimeError(
                "PerChannelMinMaxObserver.calculate_qparams called before any observation"
            )
        return _qparams_from_range(
            self.min_vals,
            self.max_vals,
            n_bits=self.n_bits,
            symmetric=self.symmetric,
            unsigned=self.unsigned,
            eps=self.eps,
        )

    def _broadcast_shape(self, ndim: int) -> list[int]:
        """Shape that reshapes a per-channel vector for broadcasting over ``ndim`` dims."""
        axis = self.ch_axis % ndim
        shape = [1] * ndim
        shape[axis] = -1
        return shape

    def quantize(self, x: torch.Tensor) -> torch.Tensor:
        """Fake-quantise ``x`` with the observed per-channel scales."""
        scale, zero_point = self.calculate_qparams()
        shape = self._broadcast_shape(x.dim())
        return fake_quantize(
            x,
            scale.reshape(shape),
            zero_point.reshape(shape),
            n_bits=self.n_bits,
            unsigned=self.unsigned,
        )

observe(x)

Fold x into the running per-channel range and return it unchanged.

Parameters

x : torch.Tensor Calibration tensor whose ch_axis length is the channel count.

Returns

torch.Tensor x unchanged.

Source code in src/sc_neurocore/qat/observers.py
Python
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
def observe(self, x: torch.Tensor) -> torch.Tensor:
    """Fold ``x`` into the running per-channel range and return it unchanged.

    Parameters
    ----------
    x : torch.Tensor
        Calibration tensor whose ``ch_axis`` length is the channel count.

    Returns
    -------
    torch.Tensor
        ``x`` unchanged.
    """
    x = x.detach()
    cur_min, cur_max = self._per_channel_min_max(x)
    if self.min_vals.numel() == 0:
        self.min_vals = cur_min
        self.max_vals = cur_max
    else:
        self.min_vals = torch.minimum(self.min_vals, cur_min)
        self.max_vals = torch.maximum(self.max_vals, cur_max)
    return x

calculate_qparams()

Return the per-channel (scale, zero_point) vectors.

Returns

tuple of (torch.Tensor, torch.Tensor) 1-D scale and zero-point tensors, one entry per channel.

Raises

RuntimeError If no tensor has been observed yet.

Source code in src/sc_neurocore/qat/observers.py
Python
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
def calculate_qparams(self) -> tuple[torch.Tensor, torch.Tensor]:
    """Return the per-channel ``(scale, zero_point)`` vectors.

    Returns
    -------
    tuple of (torch.Tensor, torch.Tensor)
        1-D scale and zero-point tensors, one entry per channel.

    Raises
    ------
    RuntimeError
        If no tensor has been observed yet.
    """
    if self.min_vals.numel() == 0:
        raise RuntimeError(
            "PerChannelMinMaxObserver.calculate_qparams called before any observation"
        )
    return _qparams_from_range(
        self.min_vals,
        self.max_vals,
        n_bits=self.n_bits,
        symmetric=self.symmetric,
        unsigned=self.unsigned,
        eps=self.eps,
    )

quantize(x)

Fake-quantise x with the observed per-channel scales.

Source code in src/sc_neurocore/qat/observers.py
Python
339
340
341
342
343
344
345
346
347
348
349
def quantize(self, x: torch.Tensor) -> torch.Tensor:
    """Fake-quantise ``x`` with the observed per-channel scales."""
    scale, zero_point = self.calculate_qparams()
    shape = self._broadcast_shape(x.dim())
    return fake_quantize(
        x,
        scale.reshape(shape),
        zero_point.reshape(shape),
        n_bits=self.n_bits,
        unsigned=self.unsigned,
    )

fake_quantize(x, scale, zero_point, *, n_bits, unsigned)

Quantise then de-quantise x (simulated quantisation, no STE).

This is the inference-time / calibration-time fake-quant used to evaluate an observer's scales; for training use the learned-step quantisers in :mod:sc_neurocore.qat.lsq. scale and zero_point broadcast against x, so per-channel parameters must already be reshaped onto the channel axis by the caller.

Parameters

x : torch.Tensor Tensor to fake-quantise. scale, zero_point : torch.Tensor Quantiser parameters, broadcastable to x. n_bits : int Quantiser bit width. unsigned : bool Whether the integer grid is unsigned.

Returns

torch.Tensor The de-quantised approximation of x on the integer grid.

Source code in src/sc_neurocore/qat/observers.py
Python
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
def fake_quantize(
    x: torch.Tensor,
    scale: torch.Tensor,
    zero_point: torch.Tensor,
    *,
    n_bits: int,
    unsigned: bool,
) -> torch.Tensor:
    """Quantise then de-quantise ``x`` (simulated quantisation, no STE).

    This is the inference-time / calibration-time fake-quant used to evaluate
    an observer's scales; for training use the learned-step quantisers in
    :mod:`sc_neurocore.qat.lsq`. ``scale`` and ``zero_point`` broadcast against
    ``x``, so per-channel parameters must already be reshaped onto the channel
    axis by the caller.

    Parameters
    ----------
    x : torch.Tensor
        Tensor to fake-quantise.
    scale, zero_point : torch.Tensor
        Quantiser parameters, broadcastable to ``x``.
    n_bits : int
        Quantiser bit width.
    unsigned : bool
        Whether the integer grid is unsigned.

    Returns
    -------
    torch.Tensor
        The de-quantised approximation of ``x`` on the integer grid.
    """
    qmin, qmax = _quant_bounds(n_bits, unsigned=unsigned)
    codes = torch.round(x / scale + zero_point).clamp(qmin, qmax)
    return (codes - zero_point) * scale