ddmtorch: A Differentiable, MPS-Accelerated Toolbox for Drift-Diffusion Model Estimation on Apple Silicon
Abstract The drift-diffusion model (DDM) explains choice probabilities and response-time (RT) dis- tributions through a small set of interpretable parameters, but fitting it requires repeated eval- uation of the Wiener first-passage time (WFPT) density. Existing tools are predominantly CPU-bound, rely on compiled code, or are embedded in Bayesian sampling frameworks. We present ddmtorch, an open-source PyTorch (Paszke et al., 2019) implementation that vectorizes the WFPT likelihood across trials and subjects, supports automatic differentiation, and runs on CPU and Apple Silicon Metal Performance Shaders (MPS). It provides batched maximum- likelihood estimation, a hybrid large-time/small-time density following Navarro and Fuss (2009), synthetic data generators, and a fixed-proportion contaminant mixture. On an Apple M5 Mac- Book Air, MPS evaluates a forward–backward pass over one million trials in 0.62 s, versus 1.22 s for CPU float32 and 3.21 s for CPU float64; a 50,000-trial, 20-subject MLE fit takes 2.32 s ver- sus 4.10 s and 7.30 s. In 100 synthetic datasets with 500 trials each, parameter correlations ranged from 0.981 to 0.998. On a public motion-discrimination dataset, ddmtorch and PyDDM recovered nearly identical parameters; on a 14-subject dataset, ddmtorch MLE and a PyMC 5 HDDM-style hierarchical fit correlated at r = 0.996–0.9995, with the expected shrinkage. The WFPT math is inherited from the literature; the contribution is a differentiable, MPS- ready implementation and a transparent evaluation. We document limitations, including the absence of inter-trial variability and hierarchical priors. Because MPS does not support float64, precision-sensitive workflows should use the CPU float64 path.
Authors
- ChiaAn Lin (ORCID: https://orcid.org/0009-0008-2034-1781)
Institutions
- National Taiwan University (TW)
Publication Details
- Journal
- Research Square
- Published
- 2026-11-02
- DOI
- https://doi.org/10.21203/rs.3.rs-11540857/v1
- Primary Topic
- Decision-Making and Behavioral Economics
- Type
- preprint