Vahid Zehtab

— De omnibus dubitandum est.

About Me

Hi there! I'm Vahid, a machine learning researcher and engineer based in Canada. My focus is generative modeling, stochastic and differential notation, and engineering my way from equations to code until, every now and then, things actually work.

Here is a brief overview of what I have been up to. If something overlaps with what you are thinking about, get in touch and we can bump heads over mildly unreasonable project plans.

Check out my work & Connect: Github • LinkedIn • Google Scholar • Email

Education
MSc. in Applied Computing - University of Toronto University of Toronto MSc. in Applied Computing

09/2022 - 12/2023

Toronto, ON, Canada

University of Toronto

MSc. in Applied Computing

Artificial Intelligence Concentration
Email

09/2022 - 12/2023 • Toronto, ON, Canada

BSc. in Computer Engineering - Sharif University of Technology Sharif University of Technology BSc. in Computer Engineering

09/2017 - 07/2022

Tehran, Iran

Sharif University of Technology

BSc. in Computer Engineering

Email

09/2017 - 07/2022 • Tehran, Iran

  • CGPA: 18.64/20 (3.8/4) — 141/140 credits
  • Thesis (Capstone project): Anomaly Detection via Explicit Density Estimation
  • Coursework: Specialized in AI, stochastic/probabilistic modeling, algorithms, and systems-oriented computer engineering. See the full course list.
  • Teaching Assistantships: TAed 12 courses across 19 offerings, including Head TAing Artificial Intelligence and several graduate-course placements as instructor, assignment designer, and TA mentor. See the full teaching record.
  • Extracurricular Activities: Led and built scientific/infrastructure tracks for major student-run AI and data-science events, including Datadays 2021, Iran's largest data science competition. See the full activities list.
Professional Experience
Senior Machine Learning Scientist - Stability AI Leading image foundation efforts, a.k.a. the battle of scales: shoving billions of images into billions of parameters with minimum GPU dollars until our diffusion models can perfectly draw that astronaut on Mars. Stability AI Senior Machine Learning Scientist Leading image foundation efforts, a.k.a. the battle of scales: shoving billions of images into billions of parameters with minimum GPU dollars until our diffusion models can perfectly draw that astronaut on Mars.

03/2025 - Present

Toronto, ON, Canada

Stability AI

Senior Machine Learning Scientist

03/2025 - Present • Toronto, ON, Canada

  • Leading the training of the upcoming text-to-image efficient Stable Diffusion model family.
  • Improved inpainting prompt adherence by 33%+ in production cases by designing a new production inpainting workflow from surveyed SOTA training-free diffusion guidance methods, noise calibration, and ODE desensitization techniques.
  • Enabled efficient T2I training on arbitrary image shapes within each batch by building a scalable mixed-aspect-ratio training stack with custom batching and efficient padding CUDA kernels, improving continuous shape/aspect-ratio generalization without significant compute overhead during training.
  • Designed a WebDataset-compatible data loading and filtration stack for massively parallel S3 streaming, achieving 20x higher throughput than SOTA WebDataset implementations and sustaining 1M+ images/sec.
  • Built a reproducible training framework by extending Hydra with distributed-backend abstractions for pure PyTorch, Lightning Fabric, Accelerate, & Ray, native experiment/checkpoint/artifact/data tracking, and custom FSDP/mixed-precision tooling for fault-tolerant scaling across hundreds of GPUs.
  • Co-built an internal 1B+ high-resolution, densely captioned text-to-image dataset, coordinating data/legal workflows and building vLLM VLM captioning plus BigQuery quality monitoring for long structured captions at scale.
  • Secured millions in compute budget by curating and presenting a de-risked, controllability-focused foundation-model training roadmap to Stability's Board; scaled the Image Foundation Models team through 20+ interviews and three senior/mid-level MLE hires.
  • Led the technical work on several multi-million-dollar diffusion foundation-model partnerships beyond image generation, including sonar signal cleaning via diffusion denoising and RAW video editing for custom non-RGB color filter array camera sensors.
Machine Learning Research Scientist - Huawei Technologies (Noah's Ark Lab) Inverse display optimization with optics, eye-aberration simulation & latent diffusion, because maybe your monitor just needs to meet your visual system halfway. Huawei Technologies (Noah's Ark Lab) Machine Learning Research Scientist Inverse display optimization with optics, eye-aberration simulation & latent diffusion, because maybe your monitor just needs to meet your visual system halfway.

06/2024 - 03/2025

Toronto, ON, Canada

Huawei Technologies (Noah's Ark Lab)

Machine Learning Research Scientist

06/2024 - 03/2025 • Toronto, ON, Canada

  • Implemented a differentiable optical simulator for display-to-eye image formation under visual aberrations, modeling how image details are perceived after propagation through an imperfect viewing system.
  • Led a team of 4 to invent an adaptive display image optimizer combining latent diffusion models and computational photography for real-time prescription-specific image pre-correction, improving perceived vision quality by +10 dB PSNR. Filed a US patent.
  • Co-developed a smart-glasses camera-control prototype for content-aware image refocusing through gesture-detected commands, moving toward conversationally interactive autofocus for wearable devices.
Machine Learning R&D Consultant - Tides Medical Histopathology segmentation foundation models & annotation workflows for tissue-sample pain points that deserve less squinting. Tides Medical Machine Learning R&D Consultant Histopathology segmentation foundation models & annotation workflows for tissue-sample pain points that deserve less squinting.

01/2024 - 12/2024

Toronto, ON, Canada

Tides Medical

Machine Learning R&D Consultant

01/2024 - 12/2024 • Toronto, ON, Canada

  • Collaborated cross-functionally with biomedical researchers to translate histopathology research objectives into concrete machine-learning, data-analysis, visualization, and annotation workflows.
  • Designed specialized annotation tooling and fine-tuned a segmentation foundation model for histopathology samples, accelerating sample evaluation for histologists while reaching over 99% prediction accuracy.
Machine Learning Researcher - Samsung Research America (SAIC) Normalizing flows & generative models for camera color processing and computational photography, until 100x smaller starts sounding normal. Samsung Research America (SAIC) Machine Learning Researcher Normalizing flows & generative models for camera color processing and computational photography, until 100x smaller starts sounding normal.

05/2023 - 12/2023

Toronto, ON, Canada

Samsung Research America (SAIC)

Machine Learning Researcher

Intern advised by: M.S. Brown, M.A. Brubaker, D.B. Lindell

05/2023 - 12/2023 • Toronto, ON, Canada

  • Devised a neural compression framework with Residual Normalizing Flows for efficient representation of color transformations in modern camera Image Processing Pipelines (ISPs), surpassing the SOTA by over 100x in compression efficacy while preserving color fidelity. Filed a US patent; published at AAAI 2025.
  • Assisted the camera quality team with RAW quad Bayer sensor noise modeling and denoising, supporting inverse vision-processing experiments for computational photography pipelines.
Machine Learning Engineer - Advanced Technology Lab (ZLAB), Fanap Co. Deep image super-resolution, CV tooling & PyTorch-native frameworks for maps that wanted more pixels than satellites gave them. Advanced Technology Lab (ZLAB), Fanap Co. Machine Learning Engineer Deep image super-resolution, CV tooling & PyTorch-native frameworks for maps that wanted more pixels than satellites gave them.

07/2020 - 07/2021

Tehran, Iran

Advanced Technology Lab (ZLAB), Fanap Co.

Machine Learning Engineer

07/2020 - 07/2021 • Tehran, Iran

  • Developed a novel satellite-image super-resolution GAN for urban maps, achieving up to 8x upsampling and deployment on a navigation platform.
  • Created a computer-vision deep-learning prototyping framework that enables direct Neural Architecture Search (NAS) over pure PyTorch models.
Software Engineer - Cafebazaar Distributed computing & statistical load testing for clusters that do not get to blink. Cafebazaar Software Engineer Distributed computing & statistical load testing for clusters that do not get to blink.

07/2019 - 09/2019

Tehran, Iran

Cafebazaar

Software Engineer

07/2019 - 09/2019 • Tehran, Iran

  • Built a distributed load-tester that statistically modeled user behavior across gRPC and REST API calls, scaling internal infrastructure tests to 1M+ requests per second.
  • Implemented a Kubernetes node-failure simulator and designed a cluster-monitoring system for stress-testing production infrastructure.
Computer Vision Engineering Intern - Pido On-device OCR, where Iranian debit cards meet very real mobile-camera constraints. Pido Computer Vision Engineering Intern On-device OCR, where Iranian debit cards meet very real mobile-camera constraints.

06/2018 - 06/2019

Tehran, Iran

Pido

Computer Vision Engineering Intern

06/2018 - 06/2019 • Tehran, Iran

  • Engineered an on-device OCR engine in C++ based on Google's Tesseract OCR Engine to extract information from Iranian debit card scans under mobile capture constraints.
Research Experience
Graduate Researcher - Vector Institute Normalizing flows for causal discovery, healthcare ML & research tooling so experiments have fewer sharp edges. Vector Institute Graduate Researcher Normalizing flows for causal discovery, healthcare ML & research tooling so experiments have fewer sharp edges.

10/2022 - 12/2023

Toronto, ON, Canada

Vector Institute

Graduate Researcher

Intern advised by R.G. Krishnan at University of Toronto's Machine Learning & Computational Healthcare Lab

10/2022 - 12/2023 • Toronto, ON, Canada

  • Co-developed a differentiable Bayesian causal structure discovery algorithm with Autoregressive Normalizing Flows that explicitly enforces acyclicity in causal DAGs, achieving the current SOTA for real-world datasets by a +2 SHD margin. NeurIPS submission.
  • Developed DyPy, introducing Pythonic syntactic sugar for dynamic code manipulation and cutting the source code required for complex deep-learning research projects roughly in half.
  • Developed Lightning-Toolbox to implement modularized, lazily evaluated loss functions and streamline version control for deep-learning objective functions across fast-moving research experiments.
Research Assistant - École polytechnique fédérale de Lausanne (EPFL) Interpretable planning for autonomous driving & some deep learning theory, with Transformers before they were cool. École polytechnique fédérale de Lausanne (EPFL) Research Assistant Interpretable planning for autonomous driving & some deep learning theory, with Transformers before they were cool.

06/2020 - 09/2021

Lausanne, Switzerland

École polytechnique fédérale de Lausanne (EPFL)

Research Assistant

Intern advised by A. Alahi at VITA lab
Research collaborator with LTS4 & MLO students

06/2020 - 09/2021 • Lausanne, Switzerland

  • Studied the theoretical foundations of deep learning and computed the inductive biases of Vision Transformers using JAX and Flax, helping identify the tasks for which Transformers are better suited than traditional Convolutional Neural Networks.
  • Developed an interpretable Transformer that employs Scalable Vector Graphics (SVGs) for vehicle-trajectory prediction through natural-language-processing style sequence modeling, unifying multiple perception data representations and improving the SOTA by 0.1 ADE on Argoverse. CVPR submission.
Research Assistant - Sharif University of Technology Deep generative models for representation learning, density estimation & image anomaly detection. Sharif University of Technology Research Assistant Deep generative models for representation learning, density estimation & image anomaly detection.

12/2019 - 07/2022

Tehran, Iran

Sharif University of Technology

Research Assistant

Advised by M.H. Rohban at Robust and Interpretable Machine Learning (RIML) lab

12/2019 - 07/2022 • Tehran, Iran

  • Designed an end-to-end training procedure for visual anomaly detection via explicit deep density estimation, reaching performance comparable to SOTA Autoencoder-based methods.
  • Assessed the potential of repurposed adversarial training and robust reconstructive models for visual anomaly detection tasks.
  • Conducted research on Energy-Based Models and other deep generative density-estimation methods, culminating in TorchDE, a unified PyTorch library for deep density estimation.
Publications
Efficient Neural Network Encoding for 3D Color Lookup Tables (AAAI 2025)
Zehtab, V., Brown, M.S., Brubaker, M.A., Lindell, D.B. - *Equal contribution as first Author

Paper Code arXiv Patent

Order-based Structure Learning with Normalizing Flows
Kamkari, H.*, Balazadeh, V.*, Zehtab, V., Krishnan, R.G. - *Equal contribution as first Author

arXiv Code

SVG-Net: An SVG-Based Trajectory Prediction Model
Bahari, M., Zehtab, V.*, Khorasani, S.*, Ayromlou, S.*, Saadatnejad, S., Alahi, A. - *Equal contribution as second author

Website Code arXiv