Sculpting Features from Noise: Reward-Guided Hierarchical Diffusion for Task-Optimal Feature Transformation

We thank you for acknowledging the novel and well-motivated formulation, strong performance, and thorough experimental evaluation of our paper.

Q1-more experiments in terms of varying sample sizes:

Since the training data consists of feature transformation sequence and corresponding score, only the feature dimensionality has an influence on the training efficiency (sequence length).

To verify this, we run an additional study in which we fix the feature dimensionality and vary the sample size only.

According to the results, DIFFT shows nearly identical training times across datasets with different sample sizes.

Q2-why the random method (RDG)’s performance is not bad in a specific dataset:

The random method is not our focus, but this is an interesting side observation. The random method didn’t perform well on many datasets. The specific ap-omentum-ovary dataset is high-dimensional (>10000 features), making the search space extremely large. While RDG randomly obtain a great performance, the two baselines MOAT and ELLM fall into a local optima in a large space.

Q3-explain the GCN-embedding design:

Feature transformation sequences are “order-free bags of features”; their semantics are governed by the pairwise (and higher-order) relations among columns rather than by token order. We therefore convert each transformed table into a feature–feature graph and let a simplified GCN produce a permutation-invariant conditioning vector of diffusion model.

Q4-memory consumption:

Our method is lightweight: on a single RTX 4090, training for both the diffusion model and the VAE stays below 10 GB of GPU memory, while inference requires only about 3 GB. We will add more details in the appendix further.

Dear Reviewer,

As the discussion phase is approaching its end, we kindly request the reviewer to let us know if the above clarifications and the previously added experiments have addressed the remaining questions. We would be happy to address any additional points the reviewer may have during the remaining time of the discussion phase.

We thank the reviewer for engaging with us in the discussion.

Best Regards,

Submission 16802 Authors

We thank you for taking the time to review our paper.

Q1-unified optimization objective:

1. Generally speaking, we don't need a unified optimization objective. According to multiple previous latent diffusion models, e.g. Stable Diffusion [1], DreamFusion [2] and PixArt$\alpha$ [3], the VAE’s goal is to learn a lossless, geometry-faithful latent space via standard reconstruction, while the diffusion component focuses on generative consistency.

2. When a single objective attempts to balance both, gradients from reconstruction and generative terms often conflict, leading to either blurry latents or unstable sampling, and even collapse. Empirically, separate optimization produces faster convergence, clearer division of labour, and simpler hyper-parameter tuning—hence we follow the established practice of first training the VAE to completion and then training the diffusion model in the fixed latent space.

Q2-reward guidance in training:

We clarify that 1) The reward signal does NOT reshape the latent embedding space, which is indeed learned by the VAE. 2) The reward signal is to guide optimal diffusion sampling.

We assume you are asking: how reward signal influences diffusion sampling. Below is our theoretical analysis, following RCGDM [4]:

Reward acts only at sampling time, tightening samples toward high-reward regions by adding $\nabla_z R_\psi(z)$ to the base score.

Let the learned conditional latent distribution be $p_0(z\mid c)$ with score
$s_0(z,t,c)=\nabla_z \log p_0(z_t\mid c)$. Introducing guidance strength $\lambda$, we effectively sample from the exponentially tilted distribution

$$p_{\lambda}(z\mid c)\propto p_{0}(z\mid c)\exp(\lambda\,R_{\psi}(z)),$$

which yields the score decomposition

$$\nabla_{z}\log p_{\lambda}(z\mid c) = \nabla_{z}\log p_{0}(z\mid c) + \lambda\,\nabla_{z}R_{\psi}(z),$$

i.e., $\tilde{s} = s_{0} + \lambda\,\nabla_{z}R_{\psi}$ in the DDIM/ODE reverse updates.

Define the energy

$$\mathcal{E}(z) = -\log p_{0}(z\mid c) - \lambda\,R_{\psi}(z).$$

Discrete reverse-diffusion steps are gradient descent on $-\nabla_{z}\mathcal{E}(z)$. If $\mathcal{E}$ is smooth and bounded below, the step size is small, and $\lambda$ is finite, then

$$\mathcal{E}(z_{k+1}) \le \mathcal{E}(z_{k}) - \eta\,\|\nabla_{z}\mathcal{E}(z_{k})\|^{2}$$

for some $\eta > 0$, so energy decreases monotonically and the trajectory converges to a stationary point—practically, high-reward/high-density regions of the manifold.

Following Reward-Directed Conditional Diffusion, the reward gap satisfies

$$R^{\star} - \mathbb{E}[R \mid p_{\lambda}] \le \varepsilon_{\text{eval}} + \varepsilon_{\text{diff}} + \text{DistroShift}(\lambda),$$

where $\varepsilon_{\text{eval}}$ is evaluator error, $\varepsilon_{\text{diff}}$ is diffusion approximation error, and $\text{DistroShift}(\lambda)$ quantifies the tilt-induced shift. When these are controlled, $\mathbb{E}[R \mid p_{\lambda}]$ exceeds unguided sampling. Under a linear-manifold assumption, sampling remains on the true data subspace, so reward redistributes probability within the manifold rather than distorting it. Thus, reward guidance is a posterior reweighting plus an energy-gradient flow mechanism at inference; training is unchanged and the procedure is predictable and convergent.

Q3-case study and quality analysis:

To provide the quality analysis of generated features, we conduct a case study on the svmguide3 dataset to (1) analyze the importance on downstream model of the generated features, (2) evaluate their robustness under a simulated distribution shift.

Analysis of Feature Importance: On the original table (21 features) the RandomForest reaches 0.778 F1 Score; After applying DIFFT we obtain 0.866 (+ 0.088). 7/10 most important positions are taken by DIFFT features. This observation shows that DIFFT creates significant features for the downstream model.

Evaluation under Distribution Shift: The following table summarizes the performance changes of models using original versus transformed dataset under scaling shifts of different severity. It analyzes the proportion of generated features within the Top-10 most important features for the model using the transformed dataset.

Q4-misunderstanding of the components:

1. Our diffusion and evaluator are not pretrained and have no prior knowledge transferred from other datasets. Instead, they are trained from scratch for task-optimal learning.

2. How to perform ablation studies: we remove/replace each of these individual components from the entire framework (Section 3.3). We blow present more ablation studies on six more datasets to show the importance of each individual component:

Notably, we add a new setting which sets the strength coefficient $\lambda$ to 10 for further validating the effectiveness of the reward mechanism. More detailed description of the variant models (NAR, AR, NoR, and CS) can be found in Section 3.1.

[1] Rombach, Robin, et al. "High-resolution image synthesis with latent diffusion models." Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2022.

[2] Poole, Ben, et al. "Dreamfusion: Text-to-3d using 2d diffusion." arXiv preprint arXiv:2209.14988 (2022).

[3] Chen, Junsong, et al. "Pixart-$\alpha$: Fast training of diffusion transformer for photorealistic text-to-image synthesis." arXiv preprint arXiv:2310.00426 (2023).

[4] Yuan, Hui, et al. "Reward-directed conditional diffusion: Provable distribution estimation and reward improvement." Advances in Neural Information Processing Systems 36 (2023): 60599-60635.

Dear Reviewer,

As the discussion phase is approaching its end, we kindly request the reviewer to let us know if the above clarifications and the previously added experiments have addressed the remaining questions. We would be happy to address any additional points the reviewer may have during the remaining time of the discussion phase. If you are satisfied, we kindly request you to consider updating the score to reflect the newly added results and discussion.

We thank the reviewer for engaging with us in the discussion.

Best Regards,

Submission 16802 Authors

Thank you for the follow-up. Below we give a clarification addressing the remaining points.

1. Training paradigm: pretraining vs. from-scratch

The two concerns raised by the reviewer—"too powerful due to heavy pretraining" vs. "too limited because trained from scratch"—are in tension. To fairly evaluate our methodological contribution (shifting from continuous search to reward-guided diffusion sampling), we follow the established standard in this line (e.g., MOAT) and train from scratch, enabling an apples-to-apples comparison of the paradigm itself. Large-scale pretraining is orthogonal and can further broaden support, but it is not a prerequisite for our contribution.

2. Losses and modular roles

Our VAE (Sec. 2.2) is trained with a composite variational objective (reconstruction + KL), whereas the LDM (Sec. 2.3) uses the standard denoising MSE; they do not "share ELBO" in practice. The synergy arises not from a shared loss but from a principled division of labor in a shared latent space: the VAE defines a stable, compact latent manifold (the space of generation), while the LDM learns a distribution over that manifold to produce plausible new points. This staged "representation-then-generation" approach is the validated paradigm for latent diffusion and supports both stability and efficiency.

3. Reward model and convergence

The reward is a parametric evaluator ($R_\psi$) trained by supervised ERM (Sec. 2.2, Eq. 4), not a hand-crafted heuristic.

3.1 Convergence sense

Building on the exponential-tilting view $p_\lambda \propto p_0 \exp(\lambda r)$ and the score decomposition $\tilde s = s_0 + \lambda \nabla r$, we formalize the update as preconditioned energy descent on

$$\mathcal{E}(z) = - \log p_0(z \mid c) - \lambda r(z).$$

The guided DDIM step can be written as

$$z_{k+1} = z_k - \eta_k H_{t_k} \nabla \mathcal{E}(z_k),$$

with a positive-definite preconditioner $H_{t_k}$. Under $L$-smoothness of $\mathcal{E}$, bounded $H_{t_k}$ ($m I \preceq H_{t_k} \preceq M I$), small step sizes $\eta_k < \tfrac{2m}{L M^2}$, bounded $\lambda$, and normalized/clipped $\nabla r$, the descent lemma gives

$$\mathcal{E}(z_{k+1}) \le \mathcal{E}(z_k) - \frac{\eta_k m}{2}\|\nabla \mathcal{E}(z_k)\|^2 + O(\eta_k^2),$$

so for sufficiently small steps $\mathcal{E}$ decreases monotonically, $\|\nabla \mathcal{E}(z_k)\|\to0$, and limit points are stationary in practice (high-density, high-utility regions of the prior manifold). Exponential tilting also yields a first-order improvement

$$E_{p_\lambda}[r] = E_{p_0}[r] + \lambda\operatorname{Var}_{p_0} \left[r\right] + O(\lambda^2)$$

implying that for small $\lambda$ the expected reward increases.

3.2 Density vs. reward

We do not maximize density and reward separately; we maximize the product-of-experts posterior

$$\log p_\lambda(z \mid c) = \log p_0(z \mid c) + \lambda r(z),$$

whose optimum satisfies $\nabla \log p_0(z^\*) + \lambda \nabla r(z^\*) = 0$ and typically lies in high-density, high-utility regions when $\lambda$ is in a stable range (see the $\lambda$-sweep below).

3.3 Which comes first? Density, then reward

We quantify the instantaneous dominance by

$$\rho_t = \frac{\|\lambda \nabla r(z_t)\|}{\|s_0(z_t, t, c)\|}.$$

Early in reverse diffusion (high noise), $\|s_0\|$ is large, while we L2-normalize/clip $\nabla r$ and then scale by a fixed $\lambda=100$. Thus $\rho_t \ll 1$ and updates are density-dominated, keeping trajectories on the prior manifold and preserving valid decodability. As noise decreases, $\|s_0\|$ shrinks and $\rho_t$ rises to $O(1)$, at which point reward-driven refinement takes over—a natural density-first, reward-later schedule even with constant $\lambda$.

3.4 Effect of guidance strength ($\lambda$-sweep)

We fix $\lambda=100$ in our experiments (no per-dataset tuning). To further investigate the effect of guidance strength, we run a $\lambda$-sweep, as shown in the table.

The sweep reveals a broad stable regime (approximately 50–150) rather than delicate tuning; the default $\lambda=100$ sits near the center and consistently improves over $\lambda=0$. Only extreme guidance ($\lambda=400$) induces the so-called mediocre–mediocre failure via off-manifold drift, which we explicitly avoid. In the revision we will add error bars over three seeds, expected evaluator reward $\mathbb{E}[R_\psi]$ versus $\lambda$, and valid-decode rates versus $\lambda$, which remain high in the 50–150 range and drop only at extreme values.

We thank you for acknowledging the novelty, excellent performance, robustness, and clear experiments of our paper.

Q1-beyond single run in experiments:

We agree and fully follow your suggestion to run more experiments. We repeated each experiment five times and report the mean ± standard deviation:

Repeated five-run averages reveal consistently low standard deviations across nearly all datasets (typically ≤ 0.02), indicating that our method’s performance is stable and not due to lucky single runs.

Q2-single-task and multi-task scenario are two different problem settings:

1. The current problem setting is focused on addressing a widely-used single-task scenario. The multi-task scenario is another problem setting that needs another solution, because multiple predictive targets lead to deviated optimal feature transformations. In the limitations and future work, we discussed the multi-task scenario can be solved by a task-agnostic foundation model. Our DIFFT exactly paves a way toward training a foundation generative feature transformation model.

2. In our experiments, we observed that our DIFFT not just identify task-specific optimal feature transformations to boost predictive accuracy (Table 1), but also significantly reduce time costs. For example, DIFFT completes a training epoch on the large-scale Ap-omentum-ovary dataset more than 25 times faster than MOAT [1] (19.64s vs. 524.4s).

Q3-multiple components but actual training costs are low:

1. Although we have multiple components (look sophisticated), indeed there is a widely-adopted latent diffusion model paradigm [2] (e.g., Stable Diffusion) which strictly requires two-stage training (VAE + diffusion).

2. We proposed the idea of semi-autoregressive decoding to reduce training costs.

3. Our experiments truly show that the total per-epoch training time for DIFFT is significantly reduced. For instance, DIFFT’s time cost is lower than for the end-to-end LSTM-based model in MOAT (Table 3).

Q4-what if feature dimensionality exceeds 50K (ultra-high dimensionality scenarios):

1. In our experiments, we tested DIFFT on datasets with up to 10,936 features (Ap-omentum-ovary) and demonstrated strong performance and efficiency.

2. Conventionally, feature transformation methods aim to enhance the expressive power of low-dimensional datasets by expanding the feature dimension. Consequently, these techniques are typically applied to datasets with a limited number of features (e.g., fewer than one hundred). In ultra-high-dimensional scenarios, however, applying feature transformation is generally not a suitable approach. Instead, feature selection is the more conventional strategy, aimed at reducing redundant features to simplify the problem.

3. We acknowledge that over 50K features will place increasing pressure on both memory and time for the VAE and diffusion model. However, DIFFT's architecture is built on a Transformer backbone that can be parallelized. So it is algorithmically more scalable than LSTM-based models in prior literature (e.g., MOAT).

Q5-learning task-agnostic representation space of feature transformation:

In this paper, we mainly focus on addressing task-optimal feature transformation problems, and the baselines we compared also follow the same problem.

If we want a task-agnostic representation space, using self-supervised loss and more feature transformation data can do the job. If we want the entire pipeline to be task-agnostic, we can use the strategy of foundation model. Our reward-guided diffusion framework can be extended to do that.

Q6-Same ablation experiments on more datasets:

We include results on six additional datasets to investigate the consistency of our ablation study as follows:

Notably, we add a new setting which sets the strength coefficient $\lambda$ to 10 for further validating the effectiveness of the reward mechanism. More detailed description of the variant models (NAR, AR, NoR, and CS) can be found in Section 3.1.

As shown above, DIFFT(with $\lambda$ = 100) achieves the best performance across all datasets. These results highlight the strength of our semi-autoregressive VAE design and reward-guided diffusion model.

[1] Wang, Dongjie, et al. "Reinforcement-enhanced autoregressive feature transformation: Gradient-steered search in continuous space for postfix expressions." Advances in Neural Information Processing Systems 36 (2023): 43563-43578.

[2] Rombach, Robin, et al. "High-resolution image synthesis with latent diffusion models." Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2022.

Dear Reviewer,

As the discussion phase is approaching its end, we kindly request the reviewer to let us know if the above clarifications and the previously added experiments have addressed the remaining questions. We would be happy to address any additional points the reviewer may have during the remaining time of the discussion phase.

We thank the reviewer for engaging with us in the discussion.

Best Regards,

Submission 16802 Authors

We thank you for acknowledging the novelty, technical contributions, and excellent performance of our paper.

Q1-list of operators in experiments:

During paper preparation, due to page limits, we focus on presenting technical insights. In experiments, we use the following operators: {+, -, *, /, sqrt, square, sin, cos, tanh, sigmoid, log, reciprocal, standard_scaler, minmax_scaler, quantile_transform}, following the setting of MOAT [1].

Q1-operator combination constraints in experiments:

The operators are either binary or unary. If a generated transformation sequence is valid, each operator must have the required number of operands, and can be parsed into a valid postfix expression tree. Otherwise, the transformation sequence is invalid and will be discarded.

We will add such experimental setup details in the appendix.

Q2-whether our framework is a meta-network:

Our goal is to learn a task/domain-optimal model that generates a feature transformation action sequence given a specific dataset-predictive task pair. However, meta-network is a generic network of generalizable knowledge learned over diverse predictive tasks/datasets. The two formulations are different. Besides, to the best of our knowledge, there is no related work adopting meta-network in feature transformation. During the different application scenarios, MAML and Propotypical Networks are not suitable to be baselines. We appreciate your suggestion and we will explore the opportunity to apply meta-learning in our future work.

Q3-raw data and our method’s training data are different concepts:

In terms of data details: raw data is the data of a predictive task; the training data fed to our model consists of different (feature transformation sequences and corresponding score) pairs, which are explored and collected by our data collection pipeline (Appendix C). Because the scale/space of feature crossing combinations is big, even raw data is not large-scale, our method’s training data on feature crossing combinations can be large. Therefore, the small scale of datasets will have limited influence on our training.

In terms of implementation details: we used a simplified GCN with symmetric normalization [2] (Appendix B). It can be seen as a fixed encoder. There is no need for additional training. We commit to add more details to the appendix in the revised version. More details can also be found in our publicly available code.

[1] Wang, Dongjie, et al. "Reinforcement-enhanced autoregressive feature transformation: Gradient-steered search in continuous space for postfix expressions." Advances in Neural Information Processing Systems 36 (2023): 43563-43578.

[2] Kipf, T. N. "Semi-Supervised Classification with Graph Convolutional Networks." arXiv preprint arXiv:1609.02907 (2016).

Dear Reviewer,

As the discussion phase is approaching its end, we kindly request the reviewer to let us know if the above clarifications have addressed the remaining questions. We would be happy to address any additional points the reviewer may have during the remaining time of the discussion phase.

We thank the reviewer for engaging with us in the discussion.

Best Regards,

Submission 16802 Authors