Source-linked AI summary
DANTINOX: A Unified Framework for Multi-Paradigm Language Modeling
Marco Simoni, Aleksandar Fontana, Giulio Rossolini, Andrea Saracino
TL;DR
Language-generation paradigms are difficult to compare fairly because separate codebases introduce differing implementation details. DANTINOX addresses this with a unified JAX/Flax Transformer backbone and configuration-based API, enabling controlled training, streaming inference, and benchmarking across paradigms. The framework also incorporates efficiency profiling, while acknowledging safety and environmental considerations for deployment.
Problem
Separate codebases make comparisons among autoregressive, masked-diffusion, and flow-matching models difficult because implementation details can confound paradigm-level differences.
Method
DANTINOX uses one modular JAX/Flax Transformer backbone and configuration-driven interfaces for three generation paradigms, attention mechanisms, hardware setups, training, inference, and benchmarking.
Results
DANTINOX provides controlled cross-paradigm evaluation through shared training infrastructure and benchmark suites covering implementation checks, generation quality, and inference efficiency.
Takeaways & Limitations
The framework makes multi-paradigm ablations and streaming inference accessible through a single API rather than separate codebases.
Takeaways & Limitations
Generative infrastructure can be misused to produce biased or harmful text, so developers should evaluate models for safety and ethical alignment before deployment.
Abstract
from arXiv · showhide
Language generation research increasingly spans three paradigms: autoregressive decoding, discrete masked diffusion, and continuous flow-matching. Comparing them is difficult because each lives in a separate codebase, so measured differences often reflect implementation details rather than the paradigms themselves. We present DantinoX, an open-source JAX/Flax library in which a single modular Transformer backbone serves all three paradigms. Switching the generation paradigm, attention mechanism, or hardware topology requires only a configuration change, while the backbone architecture, tokenizer, initialization strategy, and training infrastructure remain consistent. This enables controlled cross-paradigm comparisons within one API for training, streaming inference, and benchmarking.
1 Introduction
DANTINOX addresses the difficulty of fairly comparing autoregressive, masked-diffusion, and flow-matching language models across incompatible software ecosystems. It provides an open-source unified substrate for controlled training, evaluation, and serving.
- Motivation: Separate codebases make paradigm comparisons difficult because measured differences can reflect tokenizers, initialization, or training-loop details rather than generation paradigms.The paper identifies a shared substrate as the missing infrastructure for controlled comparisons.
- Contribution: DANTINOX provides an open-source JAX/Flax library with one modular Transformer backbone supporting autoregressive, masked-diffusion, and flow-matching paradigms.The same framework covers training, evaluation, serving, and benchmarking.
- Contribution: Switching paradigms, attention mechanisms, or hardware setups requires configuration changes rather than new model code or repositories.The framework is designed to support configuration-based ablations and reusable pipelines.
- Evaluation: The evaluation includes implementation cross-checks, generation-quality measurements, and inference-efficiency analysis across paradigm-attention combinations.The reported quality metrics include MAUVE, perplexity, diversity, and conditional BLEU.
2 Related Work
Existing language-modeling frameworks tend to optimize either large-scale autoregressive infrastructure or research modularity, leaving multi-paradigm support fragmented. DANTINOX combines these capabilities on one configurable backbone and API.
- Existing frameworks: Production-oriented frameworks provide optimized large-scale training but generally assume left-to-right causal generation.Research-oriented frameworks instead emphasize modularity and reproducibility over raw scale.
- DANTINOX: DANTINOX combines autoregressive, discrete-diffusion, and flow-matching paradigms on a single JAX/Flax backbone.It also integrates attention variants, LoRA fine-tuning, multi-GPU scaling, and benchmarking infrastructure.
- DANTINOX: Changing the paradigm, attention mechanism, or hardware topology requires only a configuration change rather than a new codebase.This lets researchers and practitioners use the same API across alternatives that existing tools often provide separately.
3 DANTINOX Architecture
DANTINOX separates a modular Transformer backbone from generation methods, allowing autoregressive, discrete-diffusion, and continuous-flow paradigms to share configurable architecture and infrastructure. Each paradigm changes its masking, objective, decoding process, and inference strategy while remaining accessible through unified training and benchmarking interfaces.
- Architecture: DANTINOX separates the model backbone from the generation method, so users select the paradigm through configuration while interchangeable components define the core model.The backbone exposes choices such as attention, feed-forward layers, and normalization.
- Backbone: The shared backbone is a modular pre-norm Transformer with configurable blocks, attention variants including MHA, GQA, and MLA, and optional efficiency mechanisms.Attention can be causal or bidirectional depending on the paradigm.
- Generation Paradigms: Autoregressive modeling uses causal attention and next-token cross-entropy, with KV-cache decoding and standard sampling strategies during inference.Supported strategies include greedy search, temperature scaling, top-k, and top-p sampling.
- Generation Paradigms: Discrete diffusion uses bidirectional attention, scheduled token masking, weighted loss over masked positions, and S reverse-diffusion steps.It also supports block generation, multiple decoding strategies, and DualCache acceleration.
- Generation Paradigms: Continuous flow-matching adopts ELF in a frozen encoder’s continuous embedding space and discretizes outputs only at the final decoding stage.This avoids per-step token-level supervision.
- Training and Testing: Unified fit(), configuration-driven tokenization, paradigm-aware Generator routing, and profiling utilities cover training, inference, and benchmarking.The framework supports tokenizers, optimizer integrations, FLOP analysis, runtime profiling, and benchmark sweeps.
4 Evaluation
DANTINOX is evaluated through independent implementation checks, controlled generation-quality comparisons, and cross-paradigm inference profiling. The results support close external agreement and reveal a batch-dependent efficiency crossover between diffusion and autoregressive decoding.
- 4 Evaluation: The evaluation uses the same public pipeline to check implementations, compare nine paradigm–attention combinations, and profile latency, throughput, and energy.The controlled comparisons hold data, tokenizers, optimizers, and computational budgets identical across paradigms.
- 4.1 External Validation: Independent dllm and xLM reproductions closely track DANTINOX, with final losses within 1% for diffusion and 8% for AR.The remaining AR difference is consistent with retained defaults such as RMSNorm and weight tying versus xLM’s LayerNorm and untied embeddings.
- 4.1 External Validation: Flipping use_swiglu to match xLM recovers 0.14 nats of validation loss without code modifications.This ablation illustrates how configuration changes can reproduce architectural variants within DANTINOX.
- 4.3 Inference Efficiency Across Paradigms: At B ≈32, AR overtakes diffusion, reaching the highest throughput and lowest energy at B = 256 on the Large backbone.At low batch sizes, diffusion is more efficient for interactive serving, but KV-cache amortisation improves AR efficiency as concurrency grows.
- 4.3 Inference Efficiency Across Paradigms: 56 ms for diffusion versus 542 ms for token-by-token AR decoding, with diffusion using roughly 7× less energy per token for one request.Diffusion completes within the 200 ms interactive SLO, whereas AR does not.
5 Conclusion
The conclusion presents DANTINOX as an open-source JAX/Flax framework unifying autoregressive, discrete masked-diffusion, and continuous flow-matching language modeling. Its configuration-driven API supports accessible comparisons, training, streaming inference, education, and reusable deployment pipelines.
- 5 Conclusion: DANTINOX unifies autoregressive decoding, discrete masked diffusion, and continuous flow-matching in one modular Transformer backbone.Users can switch paradigms, attention mechanisms, and hardware setups through configuration rather than code edits.
- 5 Conclusion: The framework provides a single API for multi-paradigm training and streaming inference, simplifying ablations and reducing the need to build custom infrastructure.The stated audiences include researchers, educators, and practitioners.
- 5 Conclusion: DANTINOX is released as an open-source foundation for building and evaluating next-generation language models.The conclusion describes the framework as zero-boilerplate and reusable across its target communities.
Broader Impact Statement
The broader-impact statement frames DANTINOX as infrastructure for more accessible and controlled multi-paradigm generative-model research. It also acknowledges resource costs and potential misuse, emphasizing profiling and safety evaluation before deployment.
- Broader Impact Statement: DANTINOX aims to facilitate fair comparisons by isolating algorithmic differences from implementation details across generative paradigms.The shared infrastructure is also intended for educators and practitioners seeking reusable pipelines.
- Broader Impact Statement: The framework integrates hardware profiling and energy tracking to support more resource-aware language-model development.This responds to the environmental and computational costs associated with training language models.
- Broader Impact Statement: Developers are encouraged to evaluate models for safety and ethical alignment before practical deployment because generative infrastructure can be misused.The statement specifically notes the risk of producing biased or harmful text.
Generation Quality and Inference Profiling
DANTINOX exposes unified pipelines for controlled quality evaluation and inference profiling across paradigms and attention variants. The extended benchmarks show distinct scaling behavior: attention variants trade throughput against memory, while diffusion and AR respond differently to generation length and batch size.
- Generation Quality and Inference Profiling: The pipeline sweeps a 3 × 3 grid of paradigms and attention mechanisms while keeping model configuration and training loops decoupled from the generative formulation.This enables controlled comparisons using generic text-scoring utilities without paradigm-specific boilerplate.
- Generation Quality and Inference Profiling: GQA matches MHA throughput across compression ratios, while MLA trades ≈20% decode throughput for its compressed cache.MHA and GQA scale ideally to B=128; MLA scales linearly at an approximately 20% lower slope.
- Generation Quality and Inference Profiling: MHA’s KV cache grows ≈4× faster than GQA-1/4’s and ≈6× faster than MLA’s as context length increases.Decode throughput remains context-insensitive in the reported range.
- Generation Quality and Inference Profiling: AR latency grows linearly with generated length, whereas diffusion finishes in fixed S passes and keeps latency flat while throughput grows linearly.In the medium-backbone batch-size sweep, diffusion’s S=32 parallel passes beat AR’s 128 sequential steps at every shown batch size.
- Generation Quality and Inference Profiling: Figure 9 profiles latency, throughput, and memory/MFU against generation length and batch size for each paradigm.The figure provides the axes for comparing scaling behavior under those two workload dimensions.