Flux1Joint: Joint Density Transformer

• Flux1Joint is a transformer model that jointly represents parameters and observations using a mask-driven conditioning mechanism to recover various conditionals.
• It employs single-stream transformer blocks with adaLN-Zero gating to improve training stability and support flexible inference modalities.
• The architecture integrates flow matching, score matching, and EDM diffusion strategies to achieve competitive results on standard simulation-based inference benchmarks.

Searching arXiv for the specified paper and closely related work on GenSBI/Flux1Joint. Flux1Joint is a transformer architecture for simulation-based inference introduced in GenSBI as a mechanism for joint density estimation in the concatenated variable space $z=(\theta,x)\in\mathbb{R}^{d_\theta+d_x}$. It extends Flux1’s gate-modulated transformer blocks to a single-stream formulation over the joint vector, and it is designed so that posteriors, likelihoods, or other conditionals can be recovered by masking rather than by retraining separate models. Within GenSBI, Flux1Joint is available under a unified interface that decouples the generative method, neural backbone, and inference mode, and it can be trained with flow matching, score matching, or EDM diffusion (Amerio, 26 May 2026).

1. Position within the GenSBI framework

Flux1Joint is one of the three transformer-based architectures provided by GenSBI, alongside SimFormer and Flux1. The framework is implemented entirely in JAX and is structured so that flow matching, score matching, and denoising diffusion can be swapped independently of the neural architecture. In that setting, Flux1Joint occupies the role of a joint-density model: it targets $q_\phi(z)\approx p(z)=p(\theta,x)$ rather than only a posterior or likelihood factorization (Amerio, 26 May 2026).

A distinguishing property is that Flux1Joint uses only single-stream transformer blocks, whereas Flux1 mixes double- and single-stream components. This suggests a deliberate architectural simplification aligned with the objective of treating parameters and observations as a single token sequence. A plausible implication is that the conditioning operation is moved from architectural separation into the masking mechanism itself, allowing a single trained model to support multiple conditioning patterns.

2. Sequence construction and transformer mechanics

Flux1Joint operates on the concatenated joint vector $z=(\theta,x)$ and treats it as a length-$(d_\theta+d_x)$ sequence. Each token carries three embeddings: a value embedding $v_i=h_{\mathrm{val}}(z_i)\in\mathbb{R}^{D_v}$, a learned ID embedding $e_i\in\mathbb{R}^{D_{\mathrm{id}}}$ determined by position in the joint vector, and a condition embedding $c_i=m_i\cdot e_{\mathrm{cond}}\in\mathbb{R}^{D_c}$, where the binary mask $m_i\in\{0,1\}$ indicates whether the coordinate is observed or must be inferred. These are concatenated into

$$h_i^{(0)}=[v_i;e_i;c_i]\in\mathbb{R}^{D_v+D_{\mathrm{id}}+D_c}.$$

This embedding scheme makes the masking pattern explicit at token level and integrates conditioning directly into the sequence representation (Amerio, 26 May 2026).

The transformer stack inherits Flux1’s adaLN-Zero gate-modulated blocks. For each block, the hidden state is layer-normalized, modulated by timestep-dependent scale and bias, and then processed by self-attention and an MLP in parallel:

$$\tilde h=\alpha(t)\cdot \hat h+\beta(t),\qquad A=\mathrm{SelfAttention}(\tilde h),\qquad M=\mathrm{MLP}(\tilde h).$$

The update is then gated and added residually,

$q_\phi(z)\approx p(z)=p(\theta,x)$0

Because $q_\phi(z)\approx p(z)=p(\theta,x)$1 is initialized near zero, each block initially behaves as the identity. The paper states that this improves training stability, which places Flux1Joint squarely within the family of gated residual transformer designs optimized for generative training regimes (Amerio, 26 May 2026).

3. Joint density estimation and masked generative losses

The formal objective is to learn a joint model $q_\phi(z)\approx p(z)=p(\theta,x)$2 so that conditionals such as $q_\phi(z)\approx p(z)=p(\theta,x)$3 or $q_\phi(z)\approx p(z)=p(\theta,x)$4 can be recovered by masking. To do so, Flux1Joint introduces a random binary mask $q_\phi(z)\approx p(z)=p(\theta,x)$5 and forms a masked noisy state at time $q_\phi(z)\approx p(z)=p(\theta,x)$6,

$q_\phi(z)\approx p(z)=p(\theta,x)$7

with $q_\phi(z)\approx p(z)=p(\theta,x)$8, $q_\phi(z)\approx p(z)=p(\theta,x)$9, and affine scheduler $z=(\theta,x)$0 (Amerio, 26 May 2026).

For conditional flow matching, the network $z=(\theta,x)$1 is trained against the per-example conditional velocity

$z=(\theta,x)$2

Under the optimal-transport scheduler $z=(\theta,x)$3 and $z=(\theta,x)$4—the CondOT path—this becomes the constant displacement $z=(\theta,x)$5. The corresponding masked joint loss is

$z=(\theta,x)$6

Only the unconditioned coordinates contribute to the objective. The same masking principle is extended to the joint versions of score matching and EDM diffusion by multiplying the residual on the unconditioned coordinates. Methodologically, this makes Flux1Joint a conditional-by-masking generative model rather than a model tied to a single inference direction (Amerio, 26 May 2026).

4. Optimization, inference, and software realization

The flow-matching training loop samples minibatches of joint examples $z=(\theta,x)$7, binary masks $z=(\theta,x)$8, Gaussian noise draws $z=(\theta,x)$9, and times $(d_\theta+d_x)$0. For each sample, it computes

$(d_\theta+d_x)$1

predicts $(d_\theta+d_x)$2, sets the target displacement $(d_\theta+d_x)$3, and minimizes

$(d_\theta+d_x)$4

with AdamW and an EMA copy of the parameters. In the JAX/Flax implementation, this loop is wrapped in jax.jit and jax.vmap (Amerio, 26 May 2026).

At the software level, Flux1Joint is implemented as a flax.nn.Module with signature apply(params, z_t, t, mask) -> predicted field. It is connected to GenSBI’s JointPipeline through a JointWrapper that manages the mask and extracts unconditioned outputs. The generative method is supplied via the strategy pattern, so FlowMatchingMethod, ScoreMatchingMethod, and DiffusionEDMMethod can be exchanged without changing the architecture. Numerical integration is handled through interchangeable diffrax solvers—FMODESolver, EDMSolver, SMODESolver, and SMSDESolver—and priors are represented as numpyro.distributions.Distribution, allowing exact log_prob through the probability-flow ODE, although the paper notes that this route is expensive. Training, inference, and diagnostics such as run_sbc, run_tarp, and LC2ST are fully batched and JIT-compiled, while checkpointing and EMA are handled by Orbax (Amerio, 26 May 2026).

5. Default configuration and relation to Flux1

Flux1Joint generalizes Flux1 by merging $(d_\theta+d_x)$5 and $(d_\theta+d_x)$6 into a single sequence and replacing the separate-stream design with a mask-driven conditioning mechanism. The paper’s typical default configuration for SBIBM runs uses $(d_\theta+d_x)$7 single-stream blocks and $(d_\theta+d_x)$8 attention heads. The embedding dimensions are reported as $(d_\theta+d_x)$9–$v_i=h_{\mathrm{val}}(z_i)\in\mathbb{R}^{D_v}$0 for value embeddings, $v_i=h_{\mathrm{val}}(z_i)\in\mathbb{R}^{D_v}$1 for ID embeddings, and $v_i=h_{\mathrm{val}}(z_i)\in\mathbb{R}^{D_v}$2 for condition embeddings, with merge mode given by concatenation of $v_i=h_{\mathrm{val}}(z_i)\in\mathbb{R}^{D_v}$3 (Amerio, 26 May 2026).

The timestep-conditioning machinery is adaLN-Zero with per-block shift $v_i=h_{\mathrm{val}}(z_i)\in\mathbb{R}^{D_v}$4, scale $v_i=h_{\mathrm{val}}(z_i)\in\mathbb{R}^{D_v}$5, and gate $v_i=h_{\mathrm{val}}(z_i)\in\mathbb{R}^{D_v}$6 initialized to zero. Architecturally, this means that Flux1Joint preserves Flux1’s gated block design while altering the data representation and conditioning semantics. The significance is not merely cosmetic: by training directly in the joint space with random masks, the model is intended to support post-training conditioning for posteriors, likelihoods, and marginals using the same learned backbone.

6. Empirical performance on SBIBM

On SBIBM tasks with flow matching and EMA checkpoints, Flux1Joint achieves the following best C2ST scores at $v_i=h_{\mathrm{val}}(z_i)\in\mathbb{R}^{D_v}$7 simulations, reported as means over 10 test observations (Amerio, 26 May 2026):

Task	Best C2ST
Two Moons	$v_i=h_{\mathrm{val}}(z_i)\in\mathbb{R}^{D_v}$8
Gaussian Linear	$v_i=h_{\mathrm{val}}(z_i)\in\mathbb{R}^{D_v}$9
Gaussian Mixture	$e_i\in\mathbb{R}^{D_{\mathrm{id}}}$0
SLCP	$e_i\in\mathbb{R}^{D_{\mathrm{id}}}$1
Bernoulli GLM	$e_i\in\mathbb{R}^{D_{\mathrm{id}}}$2

The broader summary given in the paper is that GenSBI attains near-ideal mean C2ST scores in the range $e_i\in\mathbb{R}^{D_{\mathrm{id}}}$3–$e_i\in\mathbb{R}^{D_{\mathrm{id}}}$4, with $e_i\in\mathbb{R}^{D_{\mathrm{id}}}$5 designated as ideal. For Flux1Joint specifically, TARP curves lie on the diagonal within Jeffreys 95% bands in all five tasks, which the paper interprets as evidence of well-calibrated joint-to-conditional posteriors. It also reports that Flux1Joint outperforms Flux1 and matches or beats SimFormer and OneFlowSBI on SLCP and the Gaussian tasks. These results position Flux1Joint as a joint-estimation architecture that is competitive not only in calibration diagnostics but also in two-sample discrimination metrics across standard SBI benchmarks (Amerio, 26 May 2026).