lfads-torch: A modular and extensible implementation of latent factor analysis via dynamical systems (2309.01230v1)

Abstract: Latent factor analysis via dynamical systems (LFADS) is an RNN-based variational sequential autoencoder that achieves state-of-the-art performance in denoising high-dimensional neural activity for downstream applications in science and engineering. Recently introduced variants and extensions continue to demonstrate the applicability of the architecture to a wide variety of problems in neuroscience. Since the development of the original implementation of LFADS, new technologies have emerged that use dynamic computation graphs, minimize boilerplate code, compose model configuration files, and simplify large-scale training. Building on these modern Python libraries, we introduce lfads-torch -- a new open-source implementation of LFADS that unifies existing variants and is designed to be easier to understand, configure, and extend. Documentation, source code, and issue tracking are available at https://github.com/arsedler9/lfads-torch .

• The paper presents a PyTorch-based LFADS implementation that leverages modern deep learning features to streamline neural data modeling.
• It utilizes dynamic computation graphs, PyTorch Lightning, and Hydra for simplified debugging, modular design, and configuration management.
• Its design facilitates robust hyperparameter tuning and performance benchmarking, improving reproducibility and usability in neuroscience research.

Overview of "lfads-torch: A Modular and Extensible Implementation of Latent Factor Analysis via Dynamical Systems"

The paper presents a new implementation of the Latent Factor Analysis via Dynamical Systems (LFADS), named lfads-torch, which leverages contemporary Python libraries to enhance usability, configurability, and extensibility. LFADS is a model that has demonstrated leading performance in denoising high-dimensional neural data, which is crucial for interpreting neural function and facilitating various downstream applications. The legacy implementations of LFADS, developed prior to significant advancements in machine learning libraries, posed several barriers to iterative development and wide adoption. By transitioning to a modern framework, lfads-torch offers a more intuitive interface, improved modularity, and reduced complexity for model extension and manipulation, therefore facilitating the modeling and analysis of neural data for the neuroscience community.

Technical Advancements and Implementation Features

Lfads-torch exploits several modern deep learning frameworks to streamline development and operational workflows. The primary advances made by lfads-torch in comparison to previous implementations include:

1. Dynamic Computation Graphs and Eager Execution: Leveraging PyTorch's dynamic computation graphs and eager execution allows users to benefit from an interactive debugging and development experience. This eliminates the need to comprehend complex aspects of placeholders and session management required by static computation graphs.
2. Reduction in Engineering Boilerplate: Through PyTorch Lightning, lfads-torch reduces boilerplate code required for dataset handling, optimization, logging, and more. Lightning promotes code compartmentalization, making it easier to understand and modify specific components.
3. Configurable Architecture via Hydra: The use of Hydra allows lfads-torch users to define and modify models through composable and reusable configuration files rather than altering the codebase directly, significantly simplifying the process of model tuning and experimentation.
4. Facilitated Hyperparameter Search: Employing Ray's Tune, lfads-torch simplifies large-scale hyperparameter tuning, including Population-Based Training, thereby enabling the automation of model optimization tasks.

Modules and Functional Components

Lfads-torch introduces several modules that support innovation in model development:

• Data Augmentation via AugmentationStack:

This module enables the application of a variety of data transformations seamlessly in the training pipeline, supporting methods like coordinated dropout and spatiotemporal super-resolution—techniques crucial for improving model resilience and performance on neural datasets.

• Modular Priors and Observation Models:

The implementation accommodates specialized priors and reconstruction models beyond the default multivariate normal and Poisson/ Gaussian distributions. By supporting heavy-tailed and zero-inflated distributions, the model adapts to diverse neural data modalities, ensuring robust inference capabilities.

Validation and Benchmarking

The paper validates lfads-torch by conducting paired comparisons against the previous LFADS implementation (lfads-tf2) across multiple Neural Latents Benchmark (NLB) datasets. Across these datasets, the implementation demonstrated comparable, if not improved, model performance, evidenced by metrics such as co-bps, vel R2, and fp-bps aligning closely with or exceeding those of the prior TensorFlow-based implementation.

Implications and Future Directions

Lfads-torch represents a significant step towards democratizing the use of advanced neural population models by lowering technical barriers and enhancing model configurability. As more complex and diverse neural datasets become available, implementations like lfads-torch provide a robust foundation for exploring innovative modeling approaches and developing new neuroscience methodologies. Future developments may focus on further expanding the range of data modalities supported, integrating additional neural models, and enhancing scalability to accommodate increasingly large and complex neural datasets.

In summary, lfads-torch not only builds on the pre-existing capabilities of LFADS but also advances them by incorporating modern software engineering practices that can significantly enhance both research productivity and reproducibility in neuroscientific contexts.