# EMPnMOMA

Official code for the AAAI 2026 paper **"On the Impact of Weight Quantization on Deep Neural Network Uncertainty"**.

Github link: https://github.com/Mark-Liangs/EMPnMOMA

![Overview of EMP and MOMA](overview.png)

This repo contains two sub-projects:

- **EMP/** : Exact Moment Propagation (EMP) algorithm for fairly quantifying the uncertainty before and after weight quantization.
- **MOMA/** : MOMent Alignment (MOMA) algorithm for reducing the uncertainty induced by weight quantization.

---

## 1. Set Up the Environment

Requires Python 3.11.5 and CUDA 12.1. Conda is recommended:

```bash
conda create -n dl python=3.11.5 -y
conda activate dl
pip install -r requirements.txt --extra-index-url https://download.pytorch.org/whl/cu121
```

---

## 2. EMP

Main entry: [EMP/main.py](EMP/main.py). Just run it:

```bash
cd EMP
python main.py
```

> **Note**: `n_runs = 5` is set in [EMP/main.py:178](EMP/main.py#L178) for **quick testing only**. To reproduce the paper's final results, change it to `10000`.

Results are written to [EMP/output/](EMP/output/):

- `table_fp.csv` — variance ratio table for FP-Models like FP-MLP-2
- `table_tp.csv` — variance ratio table for T-Models like T-MLP-2
- `table_tp_moma.csv` — variance ratio table for T-Models-MOMA like T-MLP-2-MOMA
- `uncertainty.csv` — averaged uncertainty over dataset across FP-Models, T-Models, and T-Models-MOMA

---

## 3. MOMA

Main entry: [MOMA/model_train.py](MOMA/model_train.py). Two modes are supported via a CLI flag:

### Mode 1: MOMA enabled (default)

Ternary training with MOMA:

```bash
cd MOMA
python model_train.py
```
We recommand the configurations of hyperparameters as `c=0.9 and \lambda=0.1` for better uncertainty reduction performance.

Output dir: `MOMA/mlp_tn_moma(c=0.9)/lam0.1/`

### Mode 2: MOMA disabled

Ternary training without MOMA:

```bash
cd MOMA
python model_train.py --no-moma
```

Output dir: `MOMA/mlp_tn/run/`

### Inspecting per-layer ternary distributions

Each output dir contains a `ternary_dist.csv` recording the empirical ternary distribution of every `TnLinear` layer after each training round:

| Column | Meaning |
| --- | --- |
| `round` | Training round index |
| `accuracy` | Test accuracy at this round |
| `layer_name` | Layer name |
| `p(-1)` | Fraction of weights ternarized to −1 |
| `q(0)` | Fraction of weights ternarized to 0 |
| `numel` | Total number of weights in the layer |

Open `ternary_dist.csv` in the corresponding output directory to see how the (p, q, 1-p-q) distribution of each layer evolves across rounds.

---

## Citation

If you find this work useful, please cite:

```bibtex
@inproceedings{liang2026on,
  author       = {Shuang Liang and
                  Xun Lu and
                  Zi{-}Ang Liu and
                  Ming{-}Liang Wang and
                  Yan Lyu and
                  Shao{-}Qun Zhang},
  title        = {On the Impact of Weight Quantization on Deep Neural Network Uncertainty},
  booktitle    = {Proceedings of the 40th AAAI Conference on Artificial Intelligence},
  pages        = {23425--23432},
  year         = {2026},
}
```
