mlr3torch: A Deep Learning Framework in R based on mlr3 and torch

Abstract: Deep learning (DL) has become a cornerstone of modern ML praxis. We introduce the R package mlr3torch, which is an extensible DL framework for the mlr3 ecosystem. It is built upon the torch package, and simplifies the definition, training, and evaluation of neural networks for both tabular data and generic tensors (e.g., images) for classification and regression. The package implements predefined architectures, and torch models can easily be converted to mlr3 learners. It also allows users to define neural networks as graphs. This representation is based on the graph language defined in mlr3pipelines and allows users to define the entire modeling workflow, including preprocessing, data augmentation, and network architecture, in a single graph. Through its integration into the mlr3 ecosystem, the package allows for convenient resampling, benchmarking, preprocessing, and more. We explain the package's design and features and show how to customize and extend it to new problems. Furthermore, we demonstrate the package's capabilities using three use cases, namely hyperparameter tuning, fine-tuning, and defining architectures for multimodal data. Finally, we present some runtime benchmarks.

• The paper introduces mlr3torch as an R-native deep learning framework that integrates mlr3 with torch for efficient pipeline design.
• It employs graph-based representations to model complex architectures, enabling hyperparameter tuning, automated benchmarking, and flexible model customization.
• Empirical evaluations show minimal runtime overhead, robust GPU acceleration, and significant extensibility for diverse deep learning tasks.

mlr3torch: A Deep Learning Framework in R Integrating mlr3 and torch

Introduction and Motivation

The paper presents "mlr3torch," a comprehensive deep learning (DL) extension for the R language that is seamlessly integrated into the mlr3 ecosystem and built atop the torch package (2604.18152). The motivation is to address current limitations in deep learning within R—notably, the dependency of several R packages on Python backends and the challenge of harmonizing deep learning with established R-based machine learning workflows for resampling, benchmarking, and hyperparameter tuning. mlr3torch provides native, extensible, and high-level abstractions facilitating the practical application and extension of DL techniques for both tabular and tensor data through a unified graph representation of machine learning pipelines and model architectures.

Positioning within Ecosystem

While Python dominates deep learning research and engineering—with PyTorch, TensorFlow, and JAX as mainstays—DL support in R has matured primarily through wrappers interfacing with these Python-based libraries. By leveraging the native torch R package, mlr3torch circumvents the dependency and performance overhead of Python bridges, enabling more direct access to GPU acceleration, automatic differentiation, and advanced neural architectures. Unlike earlier "reticulate"-driven packages (keras3, tensorflow, fastai), mlr3torch provides robust type safety, more efficient memory management, and extensibility without language interoperability friction.

Furthermore, mlr3torch introduces advantages over other high-level R DL frameworks such as luz (in terms of integration and pipeline complexity), cito, and brulee (in extensibility and architectural flexibility). Its graph-based approach, built on mlr3pipelines, allows generalization to multimodal, non-sequential, and hybrid ML pipelines, incorporating advanced tabular learning techniques and custom architecture assembly, and seamless benchmarking and tuning within the mlr3 interface.

Architecture and Core Abstractions

A central contribution is the introduction of graph representations for both pipeline and neural architectures—leveraging DAG semantics to adequately model contemporary DL structures such as transformer blocks, skip connections, and multimodal branching. Graphs encode all crucial transformations (preprocessing, data augmentation, model layers) as composable PipeOp objects, and these can be dynamically configured, tuned, and extended.

Figure 1: A simple transformer-based neural network architecture expressed as a DAG, representing tokenization, embedding, attention, residual connections, normalization, and feed-forward layers.

mlr3torch introduces a "lazy tensor" type as the core data abstraction, which supports memory-efficient handling of generic tensors and in-disk data access. In conjunction, all preprocessing and augmentation operations are implemented as graph nodes, executed lazily, enhancing performance and composability.

Furthermore, neural networks in mlr3torch are constructed either through high-level predefined modules (such as vision learners, tabular MLPs) or as custom, user-defined graphs combining arbitrary differentiable modules with explicit control over connections and data flows. The framework provides an abstraction for learners, supporting configuration via hyperparameter sets, loss specification, optimizer selection—including integration of ignite optimizers for competitive performance—and flexible callback injection for custom training logic.

Figure 2: Training and prediction phases in mlr3torch's generating graph infrastructure, visualizing how models are composed and traversed during batch execution.

Extensibility and Customization

One of mlr3torch’s pivotal contributions is enabling multi-level extensibility:

• Custom Losses and Optimizers: Arbitrary torch module losses and optimizers can be registered into mlr3torch with explicit ParamSet hyperparameter descriptions, supporting robust hyperparameter optimization and integration with benchmarking loops.
• Callbacks: Training hooks allow modification of the optimization process—such as custom regularization, gradient clipping, or dynamic unfreezing—controlled via modular, encapsulated R6 classes.
• Learner Specialization: Users can inherit and override training and prediction logic at will, or define custom data loading logic for task-specific cases (e.g., multimodal or non-tabular modalities).

The graph language extends to multimodal inputs, as illustrated in the construction of merged architectures for image and tabular feature processing, with explicit non-linearities and skip connections.

Figure 3: Residual layer implemented as a non-linear graph segment, demonstrating the combination of linear and non-linear branches via summation.

Use Cases and Empirical Results

Three principal use cases demonstrate mlr3torch’s versatility:

1. Neural Architecture Search (NAS) on Tabular Data: By constructing pipelines dynamically (with variable depth, activation, dropout, learning rates), the paper demonstrates seamless hyperparameter optimization over both architecture and optimization parameters. The configuration leverages mlr3tuning and internal early stopping using pipeline-internal validation.
2. Transfer Learning and Fine-tuning: Integration with torchvision allows for transfer learning tasks (e.g., fine-tuning ResNet-18 on Cats vs. Dogs). Data augmentation and preprocessing are encapsulated as graph nodes, while training leverages staged unfreezing via callbacks. The paper highlights rapid convergence—even with small data splits—when initializing from pretrained weights.
3. Multimodal Learning: Combining tabular and image data for melanoma detection, the paper showcases full-pipeline modeling (imputation, upsampling, augmentation, hybrid CNN-MLP architectures), managing class imbalance through loss weighting and resampling, and rigorous resampling with automatic respect of grouping variables.

   Figure 4: ROC curve for multimodal neural network evaluated via holdout resampling, demonstrating discriminative performance for melanoma classification.

Benchmarking and Runtime Analysis

A thorough runtime characterization is performed, comparing mlr3torch, native torch, and PyTorch on both CPU and GPU hardware. The critical insight is that, when using ignite optimizers, mlr3torch adds minimal overhead beyond pure torch; the dominant bottleneck is the LibTorch backend. Absolute and relative runtimes reveal that overhead is negligible for realistic model sizes and GPU scenarios, with PyTorch only marginally outperforming torch on large configurations. The analysis highlights residual overhead in simple SGD cases on CPU, which invites future optimization in R tensor management and garbage collection, but overall demonstrates practical parity for deep learning research scenarios.

Figure 5: Runtime breakdown for AdamW and SGD optimizers across varying CPU/GPU and latent sizes, contrasting mlr3torch, torch, and PyTorch.

Figure 6: Relative median runtime of mlr3torch and torch against PyTorch, confirming minimal overhead on large models and GPU settings.

Implications and Future Directions

mlr3torch advances the state of the art in R-based deep learning—not only matching Python-centric frameworks in usability for research and applied contexts but also providing tighter integration with modern benchmarking and tuning protocols, crucial for reproducibility and comparative method development. Its graph language and extensibility are particularly relevant for meta-learning, AutoML, NAS, and advanced pipeline orchestration.

The ability to define, tune, and evaluate complex DL architectures natively in R lowers the barrier for statistical/methodological researchers operating in the R ecosystem and catalyzes integration of DL with traditional statistical learning. Given its modularity, future expansions could cover unsupervised/semi-supervised learning, survival models, and broader integration with distributed computation or federated learning paradigms. Performance will trend positively as torch and LibTorch evolve, and fine-grained optimizations (especially in pipeline graph compilation and tensor memory management) present ongoing research opportunities.

Conclusion

mlr3torch delivers a highly modular, extensible, and performant deep learning framework for R, leveraging graph-based pipeline abstractions, native GPU acceleration, and seamless integration with the popular mlr3 machine learning infrastructure (2604.18152). It enables rapid prototyping of standard and custom deep learning architectures, robust benchmarking and tuning workflows, and supports advanced use cases in multimodal and transfer learning, while maintaining competitive runtime efficiency relative to established Python-based frameworks. The open architecture foreshadows straightforward extension to additional task types and modalities, strengthening R's position as a viable language for advanced machine learning research and deployment.