Pico-Banana

Overview

Pico Banana is a lightweight, high-fidelity implementation of a Denoising Diffusion Probabilistic Model (DDPM), designed to demonstrate the generative capabilities of diffusion architectures on the FashionMNIST dataset. This project serves as a compact yet robust reference for understanding the core mechanisms behind modern generative AI, specifically focusing on the U-Net architecture, sinusoidal time embeddings, and noise scheduling strategies.

The model is engineered to synthesize $32 \times 32$ grayscale images that statistically align with the FashionMNIST distribution, effectively capturing the structural diversity of clothing items ranging from trousers and coats to bags and footwear.

Technical Architecture

The core of Pico Banana is a custom-built U-Net parameterized to balance computational efficiency with generative quality.

U-Net

The neural network employs a symmetric encoder-decoder structure:

• Encoder (Downsampling): Progressively increases channel depth (64 $\to$ 128 $\to$ 256) while reducing spatial resolution. Each stage utilizes residual blocks with Group Normalization (8 groups) and SiLU (Swish) activation functions to ensure gradient flow and non-linearity.
• Bottleneck: A central stage featuring a Self-Attention block. This mechanism allows the model to capture long-range dependencies across the feature map, crucial for maintaining global coherence in the generated images.
• Decoder (Upsampling): Reconstructs the spatial resolution using bilinear interpolation combined with skip connections from the encoder path. This preserves fine-grained details lost during downsampling.
• Time Embeddings: Sinusoidal position embeddings are injected into every residual block, providing the network with temporal context $t$ essential for the reverse diffusion process.

Diffusion Process

• Noise Schedule: Implements a cosine beta schedule. This approach offers a theoretically improved noise progression compared to linear schedules, ensuring better signal preservation near $t=0$ and $t=T$, which contributes to higher sample quality.
• Training: The model is trained to predict the noise component $\epsilon$ added to a clean image $x_0$ at timestep $t$. The objective function minimizes the Mean Squared Error (MSE) between the predicted noise and the actual Gaussian noise.
• Sampling: Generation employs the standard DDPM iterative denoising process. An Exponential Moving Average (EMA) of the model weights (decay rate 0.9999) is maintained throughout training. Inference using the EMA weights significantly reduces visual artifacts and improves the robustness of the generated samples.

Implementation Details

The project is implemented in PyTorch and leverages the following key components:

• Optimizer: AdamW with a learning rate of $5 \times 10^{-5}$ and Cosine Annealing Learning Rate Scheduler for stable convergence.
• Data Pipeline: Standard torchvision transforms including resizing, tensor conversion, and normalization to $[-1, 1]$.
• Hardware Acceleration: Automatic detection of CUDA-enabled GPUs for accelerated training and inference.

Usage

To train the model and generate samples:

python pico_banana.py

The script will:

1. Download the FashionMNIST dataset.
2. Initialize the U-Net and Diffusion process.
3. Train for the specified number of epochs.
4. Periodically save model checkpoints (.pth) and generated sample grids (.png) to the results/ directory.

Requirements

• Python 3.8+
• PyTorch
• torchvision
• NumPy
• Matplotlib