A Rate-Distortion Perspective on the Emergence of Number Sense in Unsupervised Generative Models

This is the official repository of the paper "A Rate-Distortion Perspective on the Emergence of Number Sense in Unsupervised Generative Models" by Leo D'Amato, Davide Nuzzi, Alberto Testolin, Ivilin Peev Stoianov, Marco Zorzi and Giovanni Pezzulo.

• Paper:

Overview

The simulations in this repository investigate the emergence of number sense in unsupervised generative models, under encoding capacity constraints, through the lens of rate-distortion theory (RDT). We trained $\beta$-Variational Autoencoders ($\beta$-VAEs) with varying capacities on a dataset of visual scenes containing different numerosities of items. We then analyzed the learned representations and behavioral performance of these models on numerosity estimation, discrimination, and generalization tasks.

Key Research Questions Addressed:

• Can unsupervised generative models learn to represent numerosity? How do they compare to human psychophysics?
• How does encoding capacity, as formalized by rate-distortion theory, affect the emergence and quality of numerosity representations?
• How do unsupervised generative models perform compared to supervised models trained on the specific task of numerosity estimation?

Repository Structure:

.
├── configs/            # Configuration files for different models and experiments
├── data/               # Datasets  
├── logs/               # Logs and models' checkpoints 
├── notebooks/          # Python notebooks with main analyses and results
├── src/                # Source code for the models and training scripts
├── requirements.txt    # Python dependencies
└── run.py              # Main script

Installation

Ensure you have Python 3.10+ installed. Then, in a new virtual environment, install the required dependencies:

pip install -r requirements.txt

Note: the Python version used in the experiments is 3.10.12.

Usage

To train one of the models presented in the paper use the run.py script. This script requires a configuration file that specifies the model, the dataset, the training parameters, and the evaluation parameters.

Available configuration files are:

• bvae.json → β-VAE models
• bvae_sup.json → β-VAE-sup models

After modifying the configuration file to adjust hyperparameters such as model capacity, dataset settings, etc., you can train and evaluate the β-VAE via:

python run.py -c configs/bvae.json

Notebooks

• data_gen.ipynb: Scripts for generating the training and test datasets used in the experiments.
• data_entropy.ipynb: Analysis of data entropy.
• DeWind_dataset.ipynb: Processing and analysis of the DeWind dataset.
• human_comparison.ipynb: Comparison of model results with human psychophysical data.
• fig4paper.ipynb: Code for generating figures used in the paper.
• gmm4paper.ipynb: Code to reproduce the results in the "generative capability" section of the paper.
• tab4paper.ipynb: Code for generating tables used in the paper.
• inference.ipynb: Running inference and extracting representations from trained models.
• model_summary.ipynb: Summary of model architecture.

Old notebooks (for reference):

• data_converter.ipynb: Utilities for converting data formats.
• gamma_tuning.ipynb: Experiments for tuning the gamma hyperparameter.
• h5_file_manager.ipynb: Utilities for managing HDF5 data files.

Other commands

To run the processing of raw matlab files of the DeWind dataset:

python -m src.utils.process_raw_data

About the data format

Data is stored as torch tensors in .pt files directly on disk.

Tensors representing 256x256 images are saved as 2D tensors with shape (N, 256*256), where N is the number of samples. The value range of the image tensors is [0, 255] and thus the tensors are stored as torch.uint8.

Labels are saved as tensors of shape (N,).

Embeddings are saved as tensors of shape (N, latent_dim), where latent_dim is the dimensionality of the latent space of the model.