← Back to projects

OCT Latent Diffusion 3D — Notes

3D latent diffusion pipeline for volumetric OCT medical image generation with VAE compression and UNet denoising.

PyTorch Diffusion Models 3D UNet Medical Imaging VAE Generative Models
Repo

Problem

Generating realistic volumetric OCT medical images is challenging due to the extremely high dimensionality of 3D data. Training diffusion models directly on full volumes is computationally expensive and memory intensive.

A more scalable approach is to first compress the volume into a latent representation using a Variational Autoencoder (VAE) and then train the diffusion model in that latent space.

What I Built

Implemented a 3D latent diffusion pipeline that learns to generate OCT volumes by operating in a compressed latent space.

The system uses a convolutional VAE to encode OCT volumes into latent tensors, and a 3D UNet diffusion model trained with v-prediction to iteratively denoise latent samples.

The generated latent volumes are decoded back into image space and evaluated across axial, coronal, and sagittal planes.

Architecture

OCT Latent Diffusion 3D pipeline diagram
OCT latent diffusion pipeline: VAE compression, 3D UNet diffusion in latent space, and volumetric reconstruction with multi-plane visualization.
  • VAE encoder compresses OCT volumes into latent space
  • Diffusion model learns latent distribution using a 3D UNet
  • v-prediction objective stabilizes training
  • Latent samples decoded back to volumetric images
  • Multi-plane visualization across axial, coronal, and sagittal views

Results

OCT latent diffusion 3D generated volume demo
Latent diffusion enabled stable 3D generation while significantly reducing training cost compared to full-resolution volumetric diffusion models.
OCT latent diffusion 3D additional generation demo
Example generated OCT volume slices visualized across multiple planes.

Key Insights

Why latent diffusion helps

Operating in the VAE latent space dramatically reduces the dimensionality of the diffusion process, making volumetric training feasible while preserving important anatomical structure.

Why 3D diffusion matters

Using a 3D UNet allows the model to capture spatial relationships between slices, improving consistency across axial, coronal, and sagittal views compared to independent 2D generation.

System Components

  • VAE encoder-decoder for OCT volume compression
  • 3D UNet diffusion model
  • v-prediction training objective
  • Latent space sampling and denoising
  • Volume reconstruction and multi-plane visualization