Keras quantization-aware training support for custom layers with learnable activations

Short answer

A custom Keras layer with a learnable activation parameter can participate in quantization-aware training, but tfmot.quantization.keras.quantize_model requires the layer to expose a quantization configuration—not merely a serializable get_config(). The get_config() method governs Keras model serialization and checkpoint round-tripping; tfmot needs separate information about how to insert fake-quantization nodes around weights and outputs.

Confirmed behavior vs. open questions

Well-established: tfmot clones the model and wraps layers that provide a QuantizeConfig, typically through get_quantize_config() or the quantize_annotate_layer / quantize_apply APIs. quantize_scope is required when saving or loading a quantized model so that custom configuration objects can be resolved. Post-training dynamic range quantization is a separate path that quantizes weights but leaves activations in floating point.

Verify for your version: The exact acceptance criteria and error messages for custom layers in TensorFlow 2.13+ have shifted across releases. Do not assume that a layer containing only standard ops is automatically instrumented; confirm against the installed tensorflow-model-optimization package.

Minimum configuration for QAT

1. Provide a quantization configuration. Implement get_quantize_config() on the custom layer, or annotate the layer with quantize_annotate_layer before calling quantize_apply. The config describes which weights and outputs receive fake-quant nodes.
2. Keep get_config() returning the activation callable and trainable parameters (for example, beta) as JSON-serializable entries so that quantize_scope can deserialize the model correctly.
3. Ensure call() uses standard TensorFlow ops. If the layer depends on custom C++ ops or Python control flow, tfmot cannot insert fake-quant nodes and that subgraph will remain floating point.
4. Convert with the TFLite converter and inspect the result for integer-only ops. Any unquantized subgraph limits achievable CPU speed-up.

When dynamic range is the fallback

If the layer cannot supply a quantization configuration, post-training dynamic range quantization remains viable. It requires no tfmot wrapper and preserves the learnable activation, but because activations stay in floating point, latency improvements on CPU are usually modest.

Verification commands

import tensorflow_model_optimization as tfmot

quantized = tfmot.quantization.keras.quantize_model(base_model)
print(quantized.get_config())  # confirm custom layer survives cloning

converter = tf.lite.TFLiteConverter.from_keras_model(quantized)
tflite_model = converter.convert()

If quantize_model raises an exception or the converter silently drops the custom layer, the exact error text determines whether the fix is adding a quantize config, switching to annotation, or falling back to dynamic range.