Learned Round

Learned Round is a post-training quantization (PTQ) technique that improves quantization quality by learning per-weight rounding decisions, instead of relying on fixed round-to-nearest (RTN). It unifies methods such as AdaRound [1] and SignRound [2] under a single, configurable framework integrated into Brevitas’ PTQ pipelines.

💻 Examples

About the Algorithm

Motivation

Quantization mappings generally require a rounding operator, for which round‑to‑nearest (RTN) is the standard choice.

For example, in symmetric integer quantization the mapping is typically written as:

\[\mathcal{Q}(W) := s \cdot \left( \text{clip}\left( \left\lceil \frac{W}{s} \right\rfloor + z, \min \mathcal{A}, \max \mathcal{A} \right) - z \right).\]

RTN is optimal when minimizing the weight reconstruction error

\[\lVert W - \mathcal{Q}(W) \rVert_2,\]

but this optimality does not generally hold when considering the layer (or block‑wise) output reconstruction loss

\[\lVert XW - X\mathcal{Q}(W) \rVert_2,\]

which is commonly used as a proxy for downstream accuracy degradation in PTQ.

This observation motivates learned rounding, where each weight is allowed to round up or down in a data‑driven way.

Rounding Optimization

Methods such as AdaRound [1] and SignRound [2] formulate rounding as a binary optimization problem, selecting either the floor or the ceiling of the quantization grid for each weight. Although the resulting discrete problem is NP‑hard, it can be relaxed into a continuous optimization by introducing learnable parameters inside the rounding operator and optimizing them using calibration data.

In contrast to greedy solvers such as GPTQ [3] and Qronos [4], which typically solve closed‑form layer‑wise objectives sequentially, learned rounding methods:

• jointly optimize rounding decisions (per layer or per block),

• rely on gradient‑based optimization over calibration data,

• restrict the search space to a limited subset of quantization grid points.

By jointly correcting quantization error across all weights within a block in a constrained manner, this approach more effectively reduces block output error while mitigating overfitting to calibration data. However, compared to GPTQ and Qronos, learned rounding typically requires greater compute and hyperparameter tuning.

Learned Round in Brevitas

In Brevitas, these approaches are unified under the name Learned Round, providing:

• a common abstraction for learned rounding,

• flexible choices of rounding parameterization and optimization strategy,

• seamless integration with existing PTQ pipelines (LLM and ImageNet entrypoints).

Learned Round is compatible with all quantized data types currently supported by Brevitas, including:

• integer quantization (e.g. INT2 / INT4 / INT8),

• weight‑only, weight‑and‑activation, and KV‑cache quantization,

• advanced formats such as MXFP4.

It is also composable with other PTQ techniques, including QuaRot [5], SpinQuant [6], and MagR [7].

Implementation Overview

At a high level, Learned Round performs block‑wise post‑training optimization of rounding decisions, following these steps:

1. Prepare the model (optional preprocessing, e.g. disabling internal caches).

2. Insert learnable rounding parameters into the quantization operators.

3. Decompose the model into blocks.

4. For each block: a. Cache block inputs (and reference outputs) using calibration data. b. Optimize rounding parameters (and optionally scales) via a local reconstruction loss. c. Freeze the optimized rounding decisions.

5. Optionally reuse cached activations to accelerate block‑to‑block optimization.

6. Restore the original model configuration for inference.

LearnedRoundTrainer orchestrates this block‑wise optimization by wiring together:

• a learned rounding parameterization (e.g. LearnedRoundIdentity),

• block‑level reconstruction losses (e.g. MSELoss, RoundRegularisationLoss),

• optimizers and learning‑rate schedulers,

• training configuration (batch size, iterations, AMP settings, etc.).

Following, an example configuration matching the SignRound [2] setup (without scale optimization) is provided:

brevitas_examples/common/learned_round/learned_round_trainer.py

 learned_round_trainer = LearnedRoundTrainer(
     config=Config(
         trainer=TrainerConfig(
             training_args=TrainingArgs(
                 optimizers_args=[
                     OptimizerArgs(
                         target_params="learned_round",
                         optimizer_cls="SignSGD",
                         lr=5e-3,
                         lr_scheduler_args=LRSchedulerArgs(
                             lr_scheduler_cls="LinearLR",
                             lr_scheduler_kwargs={
                                 "start_factor": 1.0,
                                 "end_factor": 0.0,
                                 "total_iters": 200}))],
                 batch_size=8,
                 iters=200,
                 losses_args=[LossArgs(cls="mse")],
                 loss_scaling_factor=1000.0,
                 use_best_model=True,
                 use_amp=True,
                 amp_dtype="float16",
                 fast_update=False),
             training_handlers=[
                 HandlerSpec(
                     name="learned_round",
                     config=LearnedRoundArgs(
                         learned_round_param=LearnedRoundImplType.IDENTITY))])))

Entrypoint Integration

Learned Round is built into Brevitas’ LLM and ImageNet PTQ entrypoints. When using these entrypoints, caches, block forward functions, and block extraction logic are handled internally so you only need to pass the appropriate CLI flags.

LLM entrypoint. The brevitas_ptq_llm command enables learned round through the --learned-round flag (currently accepts identity). The --gpxq-block-name flag must be set to the transformer block attribute path (e.g., model.model.layers for Qwen and LLaMA‑family models). The following example applies SignRound‑style learned round to a Qwen model:

brevitas_ptq_llm \
    --model Qwen/Qwen3-1.7B \
    --learned-round identity \
    --gpxq-block-name model.model.layers \
    --learned-round-iters 200 \
    --learned-round-lr 5e-3 \
    --weight-bit-width 4 \
    --weight-quant-granularity per_group \
    --weight-group-size 128

ImageNet entrypoint. The brevitas_ptq_imagenet_val command supports learned round through the --learned-round flag, which accepts identity, sigmoid, or hard_sigmoid. The --target-backend layerwise flag is required. The loss function can be set with --learned-round-loss (choices: regularised_mse, mse; default: regularised_mse). The following example uses an AdaRound‑style configuration with sigmoid rounding and regularized MSE:

brevitas_ptq_imagenet_val \
    --calibration-dir /path/to/imagenet/train \
    --validation-dir /path/to/imagenet/val \
    --learned-round sigmoid \
    --target-backend layerwise \
    --learned-round-mode layerwise \
    --learned-round-loss regularised_mse \
    --learned-round-iters 1000 \
    --learned-round-lr 1e-3

When using regularised_mse, the loss combines MSE with AdaRound’s [1] round regularization term (defaults: weight 0.01, temperature annealing from 20 to 2, 20% warmup).

More examples on how to use learned round through the LLM entrypoint are provided in LLM Learned Round Examples.

Extending Learned Round

Learned Round is designed to be extensible, supporting:

• custom learned‑round parameterizations,

• optimization of additional parameters (e.g. scales),

• integration with custom models and datasets.

This section targets advanced users.

Rounding Parameterizations

Learned Round expresses rounding as:

\[\text{round}(w) = \mathcal{R}(f(w; p)) + g(p), \quad \mathcal{R} \in \{\lfloor \cdot \rceil,\ \lfloor \cdot \rfloor,\ \lceil \cdot \rceil\},\]

where \(p\) denotes learnable parameters controlling the rounding behavior (typically only one of \(f\) or \(g\) is used).

Brevitas provides several implementations in brevitas/core/function_wrapper/learned_round.py, including:

• Sigmoid (AdaRound‑style):

  \[\text{round}(p; w, T) = \lfloor w \rfloor + \sigma(p / T)\]

• Identity (SignRound‑style):

  \[\text{round}(p; w) = \left\lfloor w + \text{clip}(p, -0.5, 0.5) \right\rceil\]

To add a custom rounding parameterization:

1. Define a class implementing forward and round_forward similarly to existing implementations in brevitas/core/function_wrapper/learned_round.py.

2. Register the implementation in: - LearnedRoundImplType (brevitas/inject/enum.py) - learned_round_impl (brevitas/quant/solver/common.py)

Extending to Custom Models or Datasets

To use Learned Round with custom models or datasets outside of the supported entrypoints, four components need to be defined. The following walks through each one and illustrates them with a self-contained MLP example on synthetic data.

1. Model and blocks

Learned Round optimizes the model one block at a time. A block is the unit of optimization: a repeated structural pattern in the network architecture. Typical examples include a ResNet block (conv‑bn‑relu‑conv‑bn) in vision models, an Attention+MLP layer in transformer‑based LLMs, or even an individual layer when fine‑grained per‑layer optimization is preferred. Blocks must be accessible as named submodules so they can be extracted programmatically.

The model below is a 3‑block quantized MLP for regression. Each block contains two QuantLinear layers with a QuantReLU activation, named block_0 through block_2.

import torch
from torch import nn
import brevitas.nn as qnn

class QuantBlock(nn.Module):
    def __init__(self, in_features, hidden_dim, out_features):
        super().__init__()
        self.linear1 = qnn.QuantLinear(in_features, hidden_dim, weight_bit_width=3)
        self.relu = qnn.QuantReLU(return_quant_tensor=True)
        self.linear2 = qnn.QuantLinear(hidden_dim, out_features, weight_bit_width=3)

    def forward(self, x):
        return self.linear2(self.relu(self.linear1(x)))

class QuantMLP(nn.Module):
    def __init__(self, in_features, hidden_dim, out_features):
        super().__init__()
        self.block_0 = QuantBlock(in_features, hidden_dim, hidden_dim)
        self.block_1 = QuantBlock(hidden_dim, hidden_dim, hidden_dim)
        self.block_2 = QuantBlock(hidden_dim, hidden_dim, out_features)

    def forward(self, x):
        return self.block_2(self.block_1(self.block_0(x)))

2. Cache

A cache captures block inputs and reference outputs during a calibration forward pass so they can be replayed during optimization. It must inherit from Cache (a Dataset subclass defined in learned_round_utils.py) and implement store_inputs, store_output, reset_cache, __getitem__, __len__, and collate_fn. Inputs are typically split along the batch dimension so each sample is stored individually.

from typing import Any, Dict, Iterable, List, Tuple
from brevitas_examples.common.learned_round.learned_round_utils import Cache

class CacheMLP(Cache[torch.Tensor, torch.Tensor]):

    def __init__(self):
        self.inputs: List[torch.Tensor] = []
        self.outputs: List[torch.Tensor] = []

    def store_inputs(self, args: Tuple[torch.Tensor, ...], kwargs: Dict[str, Any]) -> None:
        self.inputs.extend(torch.split(args[0], 1, dim=0))

    def store_output(self, output: Any) -> None:
        self.outputs.extend(torch.split(output, 1, dim=0))

    def reset_cache(self) -> None:
        self.inputs, self.outputs = [], []

    def __len__(self) -> int:
        return len(self.inputs)

    def __getitem__(self, index: int) -> Tuple[torch.Tensor, torch.Tensor]:
        return self.inputs[index], self.outputs[index]

    def collate_fn(
            self,
            batch: Iterable[Tuple[torch.Tensor, torch.Tensor]],
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        inputs, outputs = zip(*batch)
        return torch.cat(inputs, dim=0), torch.cat(outputs, dim=0)

3. Forward functions

Two forward functions are needed:

• A model forward function (ModelForwardFn protocol) that runs the full model on a calibration batch. This is used to populate the cache. Note that it receives raw batches from the DataLoader — for example, TensorDataset yields a list of tensors, so the input must be unpacked.

• A block forward function (BlockForwardFn protocol) that runs a single block on cached inputs and returns its output, used during per‑block optimization.

from accelerate.utils.operations import send_to_device

def mlp_forward(model: nn.Module, inputs: List[torch.Tensor]) -> None:
    device = next(model.parameters()).device
    # TensorDataset yields [tensor], so unpack the first element
    model(send_to_device(inputs[0], device))

def mlp_block_forward(block: nn.Module, inputs: torch.Tensor) -> torch.Tensor:
    device = next(block.parameters()).device
    return block(send_to_device(inputs, device))

4. Block extraction function

A block extraction function returns the ordered list of blocks (units of optimization) from the model. It typically delegates to get_blocks with a check function that identifies blocks by name or type. The blocks are optimized sequentially in the order returned.

from brevitas_examples.common.learned_round.learned_round_trainer import get_blocks

def get_mlp_blocks(model: nn.Module) -> List[nn.Module]:
    return get_blocks(model, lambda module, name: name.startswith("block_"))

Putting it all together

With the four components defined, create a calibration DataLoader, configure the trainer with TrainerConfig, and call trainer.train() to run block‑wise optimization. The configuration below uses SignSGD with a linear LR decay, MSE loss, and the identity (SignRound‑style) rounding parameterization.

from torch.utils.data import DataLoader, TensorDataset
from brevitas_examples.common.learned_round.learned_round_trainer import LearnedRoundTrainer
from brevitas_examples.common.learned_round.learned_round_args import (
    HandlerSpec, LossArgs, LRSchedulerArgs, OptimizerArgs,
    TrainerConfig, TrainingArgs)
from brevitas_examples.common.learned_round.learned_round_method import LearnedRoundArgs

# Synthetic calibration data
calib_loader = DataLoader(TensorDataset(torch.randn(64, 8)), batch_size=8)

# Model
model = QuantMLP(in_features=8, hidden_dim=16, out_features=1)

# Trainer configuration
config = TrainerConfig(
    training_args=TrainingArgs(
        optimizers_args=[
            OptimizerArgs(
                target_params="learned_round",
                optimizer_cls="SignSGD",
                lr=5e-3,
                lr_scheduler_args=LRSchedulerArgs(
                    lr_scheduler_cls="LinearLR",
                    lr_scheduler_kwargs={
                        "start_factor": 1.0,
                        "end_factor": 0.0,
                        "total_iters": 100}))],
        batch_size=8,
        iters=100,
        losses_args=[LossArgs(cls="mse")],
        loss_scaling_factor=1000.0),
    training_handlers=[
        HandlerSpec(
            name="learned_round",
            config=LearnedRoundArgs(learned_round_param="identity"))])

# Run learned round optimization
trainer = LearnedRoundTrainer(config=config)
trainer.train(
    model=model,
    model_forward=mlp_forward,
    block_forward=mlp_block_forward,
    data_loader=calib_loader,
    cache=CacheMLP(),
    get_blocks_fn=get_mlp_blocks,
    keep_gpu=True)

Next Steps

Learned Round has been evaluated in the LLM entrypoint across multiple quantization scenarios, including weight-only and weight-and-activation PTQ, and in combination with outlier suppression techniques. For detailed results, as well as instructions on how to reproduce them, see brevitas_examples/papers/learned_round/README.md.

References

[1] (1,2,3)

Nagel, M., Amjad, R. A., Van Baalen, M., Louizos, C., & Blankevoort, T. (2020, November). Up or down? adaptive rounding for post-training quantization. In International conference on machine learning (pp. 7197-7206). PMLR.

[2] (1,2,3)

Cheng, W., Zhang, W., Shen, H., Cai, Y., He, X., Kaokao, L., & Liu, Y. (2024, November). Optimize weight rounding via signed gradient descent for the quantization of llms. In Findings of the Association for Computational Linguistics: EMNLP 2024 (pp. 11332-11350).

[3]

Frantar, E., Ashkboos, S., Hoefler, T., & Alistarh, D. (2022). Gptq: Accurate post-training quantization for generative pre-trained transformers. arXiv preprint arXiv:2210.17323.

[4]

Zhang, S., Zhang, H., Colbert, I., & Saab, R. (2025). Qronos: Correcting the Past by Shaping the Future… in Post-Training Quantization. arXiv preprint arXiv:2505.11695.

[5]

Ashkboos, S., Mohtashami, A., Croci, M. L., Li, B., Cameron, P., Jaggi, M., … & Hensman, J. (2024). Quarot: Outlier-free 4-bit inference in rotated llms. Advances in Neural Information Processing Systems, 37, 100213-100240.

[6]

Liu, Z., Zhao, C., Fedorov, I., Soran, B., Choudhary, D., Krishnamoorthi, R., … & Blankevoort, T. (2024). Spinquant: Llm quantization with learned rotations. arXiv preprint arXiv:2405.16406.

[7]

Zhang, A., Wang, N., Deng, Y., Li, X., Yang, Z., & Yin, P. (2024). Magr: Weight magnitude reduction for enhancing post-training quantization. Advances in neural information processing systems, 37, 85109-85130.