SOTA-AI Editorial

Non-Normal Distribution

Build one network that denoises 28×28 greyscale images, classifies the noise as Gaussian or uniform and estimates μ and σ of Gaussian noise.

Author: Gonçalo Franco
Written: 8 October 2026
Result: 100/100 locally measured on released validation using training data only; hidden-test performance is unknown.
Points: 100/100 (see result for validation and provenance)
Olympiad: Polish Olympiad in Artificial Intelligence (Olimpiada Sztucznej Inteligencji), 2025 · Stage II

Back to the task Download solution.py

Hints

Hint 1

Convert byte pixels to `[0, 1]` and exploit the paired clean/noisy training images for denoising.

Hint 2

Train the noise classifier jointly; return a probability because the evaluator compares the returned value with 0.5.

Hint 3

The parameter fields are absent from training, but the clean/noisy pair still gives a likelihood for the latent Gaussian parameters.

Hint 4

Account for clipping: an observed 0 or 1 is an inequality on the noise sample, not an exact residual.

The solution in one line

A shared CNN, censored-Gaussian parameter learning, and train-only frozen-feature classifier refit score 100/100 on released validation; hidden-test performance is unknown.

The editorial

Rozkład nienormalny — Non-Normal Distribution

Problem domain: Image denoising, noise-family classification, and Gaussian parameter estimation
Official metric: Four weighted components, 25 points each: classification accuracy scaled from 0.50 to 0.95, mean PSNR scaled from 10 to 16 dB, and full credit for each of Gaussian μ and σ MSE below 0.005. The sum is rounded to an integer.

Abridged statement

Given a noisy 28×28 grayscale image, return a denoised image, a probability for uniform noise (label 1; Gaussian is label 0), and estimates of Gaussian mean and standard deviation. Train only on paired clean/noisy images and noise-family labels; parameter labels are released only for validation/test scoring. Platform inference is limited to five minutes on GPU.

Data analysis

Official training data: 30,000 paired images, evenly split between Gaussian and uniform noise. Validation data: 2,000 pairs, evenly split, with parameters available for scoring.
Images are stored as 28×28×1 uint8 arrays and must be scaled to [0, 1]. Training rows contain original, noised, and label; they do not contain params.
Validation Gaussian parameters have μ mean/std 0.2421 / 0.1439, with range 0.0003–0.4997; σ mean/std 0.1768 / 0.0711, with range 0.0501–0.3000. Uniform rows use parameters [0, 0.4].
The training data retains clean/noisy pairs, so clipping is observable. The solution uses the correct censored likelihood for pixels clipped to 0 or 1 rather than treating their residuals as exact noise samples.

Experiments

The official notebook defines the evaluation function but its released starter solution is only a template. The mirror's alternative notebook contains a random-output placeholder; it has no recorded organizer score. No validation parameter targets are passed to either training stage.

Experiment	Training inputs	Result
Official starter architecture / mirror placeholder	Paired images and labels	No released score recorded; mirror placeholder returns random predictions.
Shared CNN with residual denoising, binary noise head, and censored Gaussian likelihood head	`train.pkl` only; five CPU epochs	98/100; classifier accuracy 0.909.
Refit only the classifier's final linear head on frozen train-image features	`train.pkl` only; 40 CPU head epochs after the five-epoch fit	100/100; classifier accuracy 0.960. Denoising and Gaussian estimates are unchanged.

Chosen method

The solution trains a shared convolutional network for residual denoising and global noise features, then refits only the classifier's final linear layer on frozen features from train.pkl. The classifier returns probabilities (the released evaluator thresholds the returned value directly at 0.5). Gaussian μ and σ are trained from the paired training pixels with a likelihood that treats clipped observations as left- or right-censored; validation/test parameter fields are never read. The code creates no checkpoint or copied data artifact.

Measured score and compute

The five-epoch train-only fit initially scored 98/100 on the 2,000 released validation images: PSNR 19.8606 dB, classification accuracy 0.909, Gaussian-mean MSE 0.0006841, and Gaussian-standard-deviation MSE 0.0010840, for partial points (25.00, 22.72, 25.00, 25.00). Refitting the final classifier head on the frozen features of the same 30,000 training images raised accuracy to 0.960 and the official formula returns 100/100; the other metrics remain unchanged, with partial points (25.00, 25.00, 25.00, 25.00). The head-only step took 54.72 seconds on CPU, in addition to the original 374.31 second fit. The official evaluator logic was applied to the released validation rows; validation parameters were used only for scoring. CUDA is unavailable on this host, so a platform GPU runtime estimate cannot be measured locally. Hidden-test performance is unknown.

Alternatives considered

The official starter trains a small U-Net/ConvNet for denoising and classification, but its provided model does not supervise Gaussian parameters.
Estimating parameters as the mean/std of a predicted noise residual was considered; image detail and clipping can bias those statistics, so the solution supervises the parameter head with a censored likelihood on training pairs.
A fixed mean/std prior would not meet the per-image MSE cutoffs reliably across the released parameter ranges.

Code

Solution code solution.py · 182 lines

"""Multi-task denoising model for the 2025 Polish AI Olympiad noise task.

Training uses only the released training pairs, their noise-type labels, and
the pixel censoring caused by clipping. Validation/test parameter values are
never read by this module.
"""

from __future__ import annotations

import pickle
from pathlib import Path

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader, Dataset


class _PairedImages(Dataset):
    def __init__(self, path: str | Path):
        with Path(path).open("rb") as handle:
            rows = pickle.load(handle)
        self.original = np.stack([row["original"] for row in rows]).astype(np.float32)
        self.noised = np.stack([row["noised"] for row in rows]).astype(np.float32)
        self.labels = np.asarray([row["label"] for row in rows], dtype=np.float32)
        if self.original.max() > 1.0:
            self.original /= 255.0
            self.noised /= 255.0

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

    def __getitem__(self, index: int):
        clean = torch.from_numpy(self.original[index]).permute(2, 0, 1)
        noised = torch.from_numpy(self.noised[index]).permute(2, 0, 1)
        return {
            "original": clean,
            "noised": noised,
            "label": torch.tensor(self.labels[index]),
        }


class Model(nn.Module):
    """Predict the clean image, noise family, and Gaussian parameters."""

    def __init__(self):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(1, 24, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(24, 24, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(24, 32, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(32, 32, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
        )
        self.noise_residual = nn.Conv2d(32, 1, kernel_size=3, padding=1)
        nn.init.zeros_(self.noise_residual.weight)
        nn.init.zeros_(self.noise_residual.bias)

        self.global_features = nn.Sequential(
            nn.AdaptiveAvgPool2d(1),
            nn.Flatten(),
            nn.Linear(32, 64),
            nn.ReLU(inplace=True),
        )
        self.noise_type_head = nn.Linear(64, 1)
        self.parameter_head = nn.Linear(64, 2)

    def forward(self, x: torch.Tensor):
        features = self.features(x)
        residual = self.noise_residual(features)
        denoised = (x - residual).clamp(0.0, 1.0)

        pooled = self.global_features(features)
        noise_type = torch.sigmoid(self.noise_type_head(pooled))
        raw_parameters = self.parameter_head(pooled)
        mean = 0.5 * torch.sigmoid(raw_parameters[:, :1])
        std = 0.05 + 0.25 * torch.sigmoid(raw_parameters[:, 1:2])
        return denoised, noise_type, mean, std


def _gaussian_censored_nll(
    original: torch.Tensor,
    noised: torch.Tensor,
    mean: torch.Tensor,
    std: torch.Tensor,
) -> torch.Tensor:
    """Gaussian likelihood with exact and left/right censored pixels."""
    clean = original.flatten(1)
    observed = noised.flatten(1)
    mu = mean
    sigma = std.clamp_min(1e-4)

    lower_censored = observed <= 0.0
    upper_censored = observed >= 1.0
    exact = ~(lower_censored | upper_censored)
    residual = observed - clean

    exact_log_likelihood = (
        -0.5 * ((residual - mu) / sigma).square()
        - torch.log(sigma)
        - 0.9189385332
    )
    lower_log_likelihood = torch.special.log_ndtr((-clean - mu) / sigma)
    upper_log_likelihood = torch.special.log_ndtr((mu - (1.0 - clean)) / sigma)
    log_likelihood = torch.where(
        lower_censored,
        lower_log_likelihood,
        torch.where(upper_censored, upper_log_likelihood, exact_log_likelihood),
    )
    return -log_likelihood.mean(dim=1)


def _refit_classifier_head(
    model: Model,
    dataset: Dataset,
    device: torch.device,
    epochs: int = 40,
    batch_size: int = 1024,
) -> None:
    """Tune the noise classifier on frozen CNN features from training rows."""
    model.eval()
    features = []
    labels = []
    extraction_loader = DataLoader(
        dataset, batch_size=512, shuffle=False, num_workers=0
    )
    with torch.no_grad():
        for batch in extraction_loader:
            noised = batch["noised"].to(device=device, dtype=torch.float32)
            features.append(model.global_features(model.features(noised)))
            labels.append(batch["label"].to(device=device, dtype=torch.float32).reshape(-1))

    features = torch.cat(features, dim=0)
    labels = torch.cat(labels, dim=0)
    torch.manual_seed(7)
    if device.type == "cuda":
        torch.cuda.manual_seed_all(7)
    optimizer = torch.optim.Adam(
        model.noise_type_head.parameters(), lr=0.02, weight_decay=1e-4
    )
    for _ in range(epochs):
        permutation = torch.randperm(len(features), device=device)
        for indices in permutation.split(batch_size):
            logits = model.noise_type_head(features[indices]).reshape(-1)
            loss = F.binary_cross_entropy_with_logits(logits, labels[indices])
            optimizer.zero_grad(set_to_none=True)
            loss.backward()
            optimizer.step()


def train_model(
    train_loader=None,
    data_root: str | Path = ".",
    device: torch.device | str | None = None,
    epochs: int = 5,
    batch_size: int = 256,
    learning_rate: float = 1e-3,
) -> Model:
    """Train from train.pkl or an official-style training DataLoader.

    No validation data or validation parameter targets are accessed. When a
    loader is provided, its training dataset is rebatched for throughput.
    """
    if device is None:
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    device = torch.device(device)

    if train_loader is None:
        dataset = _PairedImages(Path(data_root) / "train.pkl")
    else:
        dataset = train_loader.dataset
    loader = DataLoader(dataset, batch_size=batch_size, shuffle=True, num_workers=0)

    model = Model().to(device)
    optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate, weight_decay=1e-4)

    for _ in range(epochs):
        model.train()
        for batch in loader:
            original = batch["original"].to(device=device, dtype=torch.float32)
            noised = batch["noised"].to(device=device, dtype=torch.float32)
            labels = batch["label"].to(device=device, dtype=torch.float32).reshape(-1)

            denoised, noise_type, mean, std = model(noised)
            denoising_loss = F.mse_loss(denoised, original)
            classification_loss = F.binary_cross_entropy(
                noise_type.reshape(-1), labels
            )

            gaussian_mask = labels == 0.0
            if gaussian_mask.any():
                parameter_loss = _gaussian_censored_nll(
                    original[gaussian_mask],
                    noised[gaussian_mask],
                    mean[gaussian_mask],
                    std[gaussian_mask],
                ).mean()
            else:
                parameter_loss = denoising_loss.new_zeros(())

            loss = denoising_loss + 0.15 * classification_loss + 0.10 * parameter_loss
            optimizer.zero_grad(set_to_none=True)
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=2.0)
            optimizer.step()

    _refit_classifier_head(model, dataset, device)
    model.eval()
    return model

Download solution.py

Editorials are written by members of the SOTA community, not by the olympiad. The official statement and scoring win wherever an editorial differs.