PromptGen in a class-embedding space (Figure REF )    Previous works [1]}, [2]} have shown that class conditioning boosts generative models' performances on ImageNet [3]}. Specifically, class-conditioned generative models map a latent code \(\mathbf {z}\)  and a class embedding \(\mathbf {y}\)  to \(\mathbf {x}= G(\mathbf {z}, \mathbf {y})\) .
To extend PromptGen to these models, we train an INN \(h_{\mathbf {\theta }}\)  to map \(\mathbf {\xi } \sim \mathcal {N}(\mathbf {\mu }, \mathbf {\sigma }^2\mathbf {I})\)  to \(\mathbf {y}= h_{\mathbf {\theta }}(\mathbf {\xi })\) , where \(\mathbf {\mu }\)  and \(\mathbf {\sigma }\)  are the mean and standard deviation of \(G\) 's class embeddings.
The motivation for defining the distribution of \(\mathbf {\xi }\)  as \(\mathcal {N}(\mathbf {\mu }, \mathbf {\sigma }^2\mathbf {I})\)  but not \(\mathcal {N}(\mathbf {0}, \mathbf {I})\)  is that we want the learned INN \(h_{\mathbf {\theta }}\)  to be volume-preserving, which is easier to train.
The training objective is to minimize \(\mathbb {D}_{\mathsf {KL}}(p_{\mathbf {\theta }}(\mathbf {z}, \mathbf {y})\Vert p(\mathbf {z}, \mathbf {y}|\mathcal {C}))\) , which is equivalent to (full derivations in Appendix REF )
\(\begin{split}&\mathop {\arg \min }_{\mathbf {\theta }} \mathbb {E}_{\mathbf {\epsilon }\sim \mathcal {N}(\mathbf {0}, \mathbf {I}), \mathbf {\xi } \sim \mathcal {N}(\mathbf {\mu }, \mathbf {\sigma }^2\mathbf {I}), \mathbf {z}= f_{\mathbf {\theta }}(\mathbf {\epsilon }), \mathbf {y}= h_{\mathbf {\theta }}(\mathbf {\xi }), \mathbf {x}= G(\mathbf {z}, \mathbf {y})}\Big [- \log |\det (\frac{\partial f_{\mathbf {\theta }}}{\partial \mathbf {\epsilon }})| - \log p_{\mathbf {z}}(\mathbf {z}) \\&\quad \quad \quad \quad \quad \quad \quad \quad \quad \quad \quad \quad \quad \quad \quad \quad - \log |\det (\frac{\partial h_{\mathbf {\theta }}}{\partial \mathbf {\xi }})| - \log p_{\mathbf {y}}(\mathbf {y}) + E_{\mathcal {C}}\big (\mathbf {x}\big )\Big ],\end{split}\) 
