---
title: 'Flux1Joint: Joint Density Transformer'
url: https://www.emergentmind.com/topics/flux1joint
type: topic
---

# Flux1Joint: Joint Density Transformer

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 [2605.27499].

## 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 [2605.27499].

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 [2605.27499].

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,
$$
\delta=A+M,\qquad
h^{(\ell+1)}=h^{(\ell)}+g(t)\odot \delta.
$$
Because $g(t)$ 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 [2605.27499].

## 3. Joint density estimation and masked generative losses

The formal objective is to learn a joint model $q_\phi(z)\approx p(\theta,x)$ so that conditionals such as $q_\phi(\theta\mid x_{\mathrm{obs}})$ or $q_\phi(x\mid \theta_{\mathrm{obs}})$ can be recovered by masking. To do so, Flux1Joint introduces a random binary mask $m\in\{0,1\}^{d_\theta+d_x}$ and forms a masked noisy state at time $t\in[0,1]$,
$$
z_t^m=m\odot z_1+(1-m)\odot \psi_t(z_0\mid z_1),
$$
with $z_1\sim p(\theta,x)$, $z_0\sim p_0$, and affine scheduler $\psi_t(z_0\mid z_1)=\sigma_t z_0+\alpha_t z_1$ [2605.27499].

For conditional flow matching, the network $v_\phi(z_t^m,t,m)$ is trained against the per-example conditional velocity
$$
u_t(z_t,z_1)=\dot\sigma_t/\sigma_t\,(z_t-\alpha_t z_1)+\dot\alpha_t z_1.
$$
Under the optimal-transport scheduler $\alpha_t=t$ and $\sigma_t=1-t$—the CondOT path—this becomes the constant displacement $u_t=z_1-z_0$. The corresponding masked joint loss is
$$
\mathcal{L}_{\mathrm{CFM}}(\phi)=
\mathbb{E}_{m,t,z_0,z_1}
\Big\|
(1-m)\odot\bigl(v_\phi(z_t^m,t,m)-(z_1-z_0)\bigr)
\Big\|^2.
$$
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 [2605.27499].

## 4. Optimization, inference, and software realization

The flow-matching training loop samples minibatches of joint examples $z_1^k=(\theta^k,x^k)$, binary masks $m^k$, Gaussian noise draws $z_0^k\sim\mathcal{N}(0,I)$, and times $t^k\sim\mathrm{Uniform}[0,1]$. For each sample, it computes
$$
z_t^{k,m}=m^k\odot z_1^k+(1-m^k)\odot\bigl((1-t^k)z_0^k+t^k z_1^k\bigr),
$$
predicts $v^k=v_\phi(z_t^{k,m},t^k,m^k)$, sets the target displacement $\Delta^k=z_1^k-z_0^k$, and minimizes
$$
\mathrm{Loss}_k=\|(1-m^k)\odot(v^k-\Delta^k)\|^2
$$
with AdamW and an EMA copy of the parameters. In the JAX/Flax implementation, this loop is wrapped in `jax.jit` and `jax.vmap` [2605.27499].

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 [2605.27499].

## 5. Default configuration and relation to Flux1

Flux1Joint generalizes Flux1 by merging $\theta$ and $x$ 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 $\ell=16$ single-stream blocks and $H=4$ attention heads. The embedding dimensions are reported as $D_v=10$–$20$ for value embeddings, $D_{\mathrm{id}}=10$ for ID embeddings, and $D_c=5$ for condition embeddings, with merge mode given by concatenation of $\{v_i,e_i,c_i\}$ [2605.27499].

The timestep-conditioning machinery is adaLN-Zero with per-block shift $\alpha(t)$, scale $\beta(t)$, and gate $g(t)$ 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 $10^5$ simulations, reported as means over 10 test observations [2605.27499]:

| Task | Best C2ST |
|---|---:|
| Two Moons | $0.524\pm0.009$ |
| Gaussian Linear | $0.500\pm0.003$ |
| Gaussian Mixture | $0.513\pm0.007$ |
| SLCP | $0.566\pm0.039$ |
| Bernoulli GLM | $0.555\pm0.016$ |

The broader summary given in the paper is that GenSBI attains near-ideal mean C2ST scores in the range $0.50$–$0.56$, with $0.50$ 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 [2605.27499].

Source: https://www.emergentmind.com/topics/flux1joint