Energy-Tweedie: Score meets Score, Energy meets Energy

Andrej Leban · University of Michigan · NeurIPS 2026

Paper · Code · Project page · Checkpoints

The paper lifts the Tweedie formula to a distributional identity valid for Gibbs noise distributions, with several applications such as a posterior-samples-to-score route, estimating unknown noise parameters from corrupted data, and providing the score-based perspective on recent "distributional" diffusion methods.

This repository hosts the released PyTorch checkpoints and saved configurations for the paper's Eight-Gaussians and MNIST experiments. The files match the public code repository and the MNIST model release.

Checkpoints

Directory	Models	Experiments
score_and_diffusion/	Isotropic Gaussian, isotropic generalized-Gaussian, and five-parameter posterior models	Score estimation, Richardson extrapolation, and diffusion-style sampling on Eight-Gaussians
score_comparison/	ET-DSM and Hyvärinen direct-score models	Three-estimator score comparison on Eight-Gaussians
parameter_estimation/	Posterior and score models	Noise-parameter recovery on Eight-Gaussians
mnist/posterior/	Five-parameter MNIST posterior	Five sampling paths on MNIST

Each checkpoint has an adjacent config.yaml or config.json. The reproduction notebooks use these saved configurations; the training scripts provide editable defaults for retraining.

The root config.json indexes each checkpoint and its saved configuration.

Download

Follow the code installation instructions, then run from the code repository root. The Hub root corresponds to models/ in the code repository.

Download all released checkpoints and configurations:

python -m pip install huggingface_hub
hf download andleb/et --local-dir models

To download the MNIST checkpoint, its configuration, and the index:

hf download andleb/et config.json mnist/posterior/final.pt mnist/posterior/config.yaml --local-dir models

Then generate the MNIST samples with the provided runner:

bash exp/sample_diffusion_mnist.sh

The MNIST notebook loads these samples and computes the comparisons presented in the paper. The Eight-Gaussians notebooks evaluate the corresponding checkpoints.

Load the MNIST model in Python

After installing the code repository and huggingface_hub, run from the code repository root:

import json

import torch
from huggingface_hub import hf_hub_download
from omegaconf import OmegaConf

from utils.mnist import build_mnist_engressor

index_path = hf_hub_download(repo_id="andleb/et", filename="config.json")
with open(index_path) as handle:
    model_files = json.load(handle)["checkpoints"]["mnist/posterior"]

config_path = hf_hub_download(
    repo_id="andleb/et",
    filename=model_files["config"],
)
checkpoint_path = hf_hub_download(
    repo_id="andleb/et",
    filename=model_files["weights"],
)
config = OmegaConf.load(config_path)
checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=True)
engressor = build_mnist_engressor(config, device="cpu")
engressor.model.load_state_dict(checkpoint["model_state_dict"])
engressor.model.eval()

The MNIST sampling code applies the noise-parameter conditioning and Energy-Tweedie sampler.

License and citation

Released under the BSD 3-Clause license. The code repository includes third-party attribution.

If you use this work, please cite the paper:

@inproceedings{leban2026energytweedie,
  title         = {Energy-Tweedie: Score meets Score, Energy meets Energy},
  author        = {Andrej Leban},
  booktitle     = {Advances in Neural Information Processing Systems},
  year          = {2026},
  eprint        = {2512.23818},
  archivePrefix = {arXiv},
  primaryClass  = {stat.ML},
  url           = {https://arxiv.org/abs/2512.23818}
}