tribev2-rs: A Rust-Implemented Inference Engine for Multimodal fMRI Brain Encoding Models

tribev2-rs is a pure Rust implementation of the TRIBE v2 brain encoding model inference engine, supporting text/audio/video multimodal inputs and capable of predicting cerebral cortex activity. This project addresses the performance bottlenecks, memory management issues, and deployment complexities of the original Python implementation. Leveraging Rust's zero-cost abstractions, memory safety, and concurrent performance, it achieves high-performance inference. It is open-source and provides a complete toolchain, supporting fields such as computational neuroscience and brain-computer interfaces.

Functional Magnetic Resonance Imaging (fMRI) non-invasively records brain activity via BOLD signals, but the complexity and high dimensionality of the data pose challenges. Brain encoding models aim to establish a mapping from external stimuli to brain activity. Traditional models are mostly unimodal, while the human brain integrates multiple modalities. TRIBE v2 (developed by Meta) is a deep multimodal brain encoding foundation model that can process text/audio/video inputs, predict neural activity at approximately 20484 cortical vertices in the fsaverage5 space, and simulate multisensory integration mechanisms.

Reasons for Rust Rewrite: Python has performance bottlenecks, memory management issues, and deployment complexities, while Rust offers zero-cost abstractions, memory safety, and concurrent performance.

Model Architecture:

1. Multimodal Feature Extraction: Extracts features using LLaMA3.2 (text), V-JEPA2 (video), and Wav2Vec-BERT (audio), then projects them to a unified dimension for aggregation;
2. Transformer Encoder: 8 layers, 8 attention heads, ScaleNorm normalization, and RoPE;
3. Low-Rank Prediction Head: Maps to the cortical surface and controls the number of parameters;
4. Temporal Smoothing Module: Uses depthwise separable convolution to simulate the delay effect of BOLD signals.

Engineering Innovations:

• Segmented Inference: Handles long sequence inputs while maintaining temporal continuity;
• Event Pipeline: Automates conversion from raw media to input (WhisperX speech recognition, ffmpeg audio extraction);
• Brain Surface Visualization: SVG rendering with multi-view, color mapping, and RGB overlay;
• FreeSurfer Compatibility: Supports mainstream neuroimaging formats.

Performance Optimization: Reduced inference time from 27.6ms to 16.8ms. Optimization steps include fixing architectural issues (e.g., non-causal attention), using f16 half-precision, Metal WMMA instructions, CubeCL fused kernels, etc., across Metal/Vulkan/DirectX12 backends.

tribev2-rs can be applied in:

• Computational Neuroscience: Verifying hypotheses about brain multimodal integration;
• Brain-Computer Interfaces: Improving the accuracy and real-time performance of neural signal decoding;
• AI Safety and Alignment: Understanding the correspondence between multimodal models and human brain representations;
• Clinical Neuroscience: Assisting in the diagnosis of neurological diseases and treatment evaluation.

tribev2-rs is open-sourced under the Apache-2.0 license, providing a complete inference engine, example code, benchmark tools, and visualization components. The project collaborates with the Rust ML ecosystem such as llama-cpp-rs and Burn, demonstrating Rust's performance and reliability advantages in the AI/ML field and promoting the maturity of the Rust ML toolchain.

tribev2-rs integrates cutting-edge models from computational neuroscience, the rigor of Rust systems programming, and the spirit of open-source collaboration, serving as a bridge between AI and human intelligence. It provides a solid starting point for researchers to understand brain multimodal processing and for engineers to seek high-performance neural computing solutions.

Project Link: https://github.com/eugenehp/tribev2-rs Original Model: https://github.com/facebookresearch/tribev2 Tech Stack: Rust · Burn ML Framework · llama-cpp · wgpu · Metal/CUDA/Vulkan