BalancedTernaryEngine.
Copy Markdown
Quantization-Aware Training layer with ternary weights.
Maintains FP32 shadow weights that are ternarized on each forward pass. Uses a straight-through estimator (STE) so gradients flow to the shadow weights during training: the forward pass uses ternary weights, while the gradient is applied to the shadow weights directly.
This module provides the primitives (ternarize, forward, update, freeze). The training loop composes these with a loss and an optimizer step.
Summary
Functions
Forward pass: output = (ternary_weight @ input) + bias.
Freeze a QAT layer into a standard BalancedTernaryEngine.Layer for inference.
Create a QAT layer from a pretrained weight tensor.
Ternarize shadow weights. Returns {effective, shadow} where effective
is the ternary (scaled) weight used in the forward pass and shadow is
the FP32 shadow weight retained for gradient updates (STE).
Apply a straight-through training step using a finite-difference gradient of the supplied loss function.
SGD update of shadow weights: w_shadow = w_shadow - lr * grad.
Freeze and validate a QAT layer.
Types
@type t() :: %BalancedTernaryEngine.QATLayer{ bias: Nx.Tensor.t(), delta: float(), name: String.t(), scale: Nx.Tensor.t(), weight_shadow: Nx.Tensor.t() }
Functions
@spec forward(t(), Nx.Tensor.t()) :: Nx.Tensor.t()
Forward pass: output = (ternary_weight @ input) + bias.
@spec freeze(t()) :: BalancedTernaryEngine.Layer.t()
Freeze a QAT layer into a standard BalancedTernaryEngine.Layer for inference.
@spec from_pretrained(Nx.Tensor.t(), String.t(), keyword()) :: t()
Create a QAT layer from a pretrained weight tensor.
@spec ternarize(t()) :: {Nx.Tensor.t(), Nx.Tensor.t()}
Ternarize shadow weights. Returns {effective, shadow} where effective
is the ternary (scaled) weight used in the forward pass and shadow is
the FP32 shadow weight retained for gradient updates (STE).
@spec train_step(t(), Nx.Tensor.t(), Nx.Tensor.t(), keyword()) :: {t(), float()}
Apply a straight-through training step using a finite-difference gradient of the supplied loss function.
loss_fn receives the layer and returns a scalar loss. The gradient
w.r.t. the shadow weights is approximated and used to update them.
Returns the updated layer and the loss value.
@spec update(t(), Nx.Tensor.t(), float()) :: t()
SGD update of shadow weights: w_shadow = w_shadow - lr * grad.
@spec validate_frozen(t()) :: {:ok, BalancedTernaryEngine.Layer.t()} | {:error, term()}
Freeze and validate a QAT layer.