Normalizing Flow for 2D Density Estimation

A from-scratch implementation of a Normalizing Flow model, based on Real NVP (Real Non-Volume Preserving) transformations, to learn and sample from a 2D 'two-moons' distribution. This project was developed for the M.S. Generative Models course at Sharif University of Technology.

Features

• Real NVP Implementation: Core logic built using affine coupling layers.
• Modular Flow Class: A nn.Module that stacks multiple transformations for complex density estimation.
• Maximum Likelihood Training: Optimized by minimizing the negative log-likelihood (NLL) of the data.
• Density Visualization: Generates and saves the learned probability density at each epoch.
• GIF Generation: Automatically creates a GIF visualizing the evolution of the learned density over training.
• Data Sampling: Capable of sampling new data points from the learned distribution.

Core Concepts & Techniques

• Normalizing Flows: A class of generative models that learn a data distribution $p_X(x)$ by transforming a simple base distribution $p_Z(z)$ (e.g., a Gaussian) through an invertible, differentiable mapping $f: \mathcal{Z} \to \mathcal{X}$.

• Change of Variables Formula: The core principle used for training. The log-likelihood of a data point $x$ is given by the formula:

  $$\log p_X(x) = \log p_Z(f(x)) + \log \left| \det J_f(x) \right|$$

  where $z = f(x)$ is the forward transformation (data to latent) and $J_f(x)$ is its Jacobian.

• Coupling Layers (Real NVP): An efficient invertible transformation $y = f(x)$ where the input $x$ is split into two parts ($x_A, x_B$). One part is left unchanged ($y_A = x_A$), while the other is transformed using a simple affine function (scaling and translation) whose parameters ($s, t$) are generated by a neural network that only takes the unchanged part as input:

  $$y_B = x_B \odot \exp(s(x_A)) + t(x_A)$$

  This design results in a triangular Jacobian matrix, making the determinant (and its log) trivial to compute: it is simply the sum of the scaling factors, $\sum s(x_A)$.

• Maximum Likelihood Estimation (MLE): The model is trained by adjusting the parameters of the $s$ and $t$ networks to maximize the log-likelihood of the training data, which is equivalent to minimizing the negative log-likelihood (NLL) loss.

---

How It Works

This project learns a complex distribution (two moons) by mapping it to a simple one (a 2D standard Gaussian).

1. Core Logic & Architecture

The repository implements a Flow model composed of a series of CouplingTransform layers. Each CouplingTransform layer splits the 2D input $x = (x_1, x_2)$ based on an alternating binary mask.

1. Transform 1 (Mask [1, 0]):

   • $x_1$ is fixed.
   • $s_1, t_1 = \text{NN}_1(x_1)$
   • $z_1 = x_1$
   • $z_2 = x_2 \odot \exp(s_1) + t_1$

2. Transform 2 (Mask [0, 1]):

   • $z_2$ is fixed.
   • $s_2, t_2 = \text{NN}_2(z_2)$
   • $z'_1 = z_1 \odot \exp(s_2) + t_2$
   • $z'_2 = z_2$

This sequence of invertible transformations ($x \to z \to z' \dots$) maps the complex 'two-moons' data to a 2D standard Gaussian. Training is done by maximizing $\log p_X(x)$, and sampling is done by drawing $z$ from the Gaussian and applying the inverse of all transforms: $x = f_1^{-1}(f_2^{-1}(\dots f_N^{-1}(z)))$.

2. Key Files & Functions

• main.py: The main executable script. It handles argument parsing, sets up logging, loads the data, initializes the Flow model, and runs the training loop. After training, it generates all final visualizations.
• src/model.py: Defines the Flow class. This class stacks multiple CouplingTransform layers. It implements:
  • forward(x): Maps data to latent space ($x \to z$) and returns $z$ and the total $\log \det J$.
  • inverse(z): Maps latent space to data space ($z \to x$).
  • log_prob(x): Applies the Change of Variables formula to compute the log-likelihood of the data.
  • sample(n): Generates $n$ new samples by drawing from the base distribution and applying the inverse method.
• src/transforms.py: Defines the CouplingTransform class. This class implements the core Real NVP logic, including the $s$ and $t$ neural networks, the forward pass, the inverse pass, and the log-det-Jacobian calculation.
• src/utils.py: Contains all helper functions for loading the make_moons dataset (load_data), creating the alternating masks (create_alternating_masks), setting up logging (setup_logging), and handling all visualizations (visualize_density, plot_comparison, create_gif).

---

Project Structure

pytorch-normalizing-flows-2d/
├── .gitignore             # Ignores Python caches, envs, and output files
├── LICENSE                # MIT License file
├── README.md              # This repository README
├── main.py                # Main script to train the model and generate outputs
├── requirements.txt       # Required Python packages
├── run_flow_model.ipynb   # Jupyter Notebook to run the training and view results
└── src/
    ├── __init__.py        # Makes 'src' a Python package
    ├── model.py           # Contains the main 'Flow' model class
    ├── transforms.py      # Defines the 'CouplingTransform' layer
    └── utils.py           # Helper functions for data, logging, and visualization

How to Use

1. Clone the Repository:

   git clone https://github.com/msmrexe/pytorch-normalizing-flows-2d.git
   cd pytorch-normalizing-flows-2d

2. Setup the Environment: It's recommended to use a virtual environment.

   python -m venv venv
   source venv/bin/activate  # On Windows: venv\Scripts\activate
   pip install -r requirements.txt

3. Run the Training Script: You can run the model directly from the command line.

   python main.py --epochs 150 --lr 0.001 --output_dir outputs

   • --epochs: Number of training epochs (e.g., 150).
   • --lr: Learning rate (e.g., 0.001).
   • --output_dir: Directory to save final plots and GIF.
   • --skip_frames: (Optional) Use this flag to disable saving frames every epoch, which speeds up training.

4. Example Output / Check Results: The script will log progress to the console and to logs/flow.log. After training, you will find the following files in the outputs/ directory:

   • density_evolution.gif: An animation of the learned density.
   • data_comparison.png: A scatter plot comparing original and generated data.
   • final_density.png: The final learned probability density.

   Alternatively, you can run the run_flow_model.ipynb notebook in Jupyter to execute the training and see the outputs displayed inline.

---

Author

Feel free to connect or reach out if you have any questions!

• Maryam Rezaee
• GitHub: @msmrexe
• Email: ms.maryamrezaee@gmail.com

---

License

This project is licensed under the MIT License. See the LICENSE file for full details.