SAFT: Shift-Aware Feature Transformation

• Shift-Aware Feature Transformation (SAFT) is a framework that redefines feature engineering by embedding transformed features into a continuous latent space to counteract train–test distribution shifts.
• It replaces discrete search with a continuous optimization process using a graph-based encoder, an evaluator for performance scoring, and an LSTM decoder to generate symbolic feature-cross sequences.
• SAFT integrates three robustness mechanisms—embedding decorrelation with sample reweighting, flatness-aware generation via suboptimal averaging, and normalization-based distribution alignment—to improve generalization in tabular data.

Searching arXiv for the cited SAFT and closely related shift-aware feature-transformation papers. Shift-Aware Feature Transformation (SAFT) is a tabular-learning framework for Distribution Shift Tabular Learning (DSTL), the setting in which training and test sets share the same raw features but differ in their data distributions. SAFT reframes feature transformation from a discrete search problem into a continuous representation-generation paradigm: transformed feature sets are embedded into a latent space, scored by an evaluator, and decoded back into symbolic feature-cross sequences. Its design is explicitly shift-aware through three coupled mechanisms—shift-resistant representation via embedding decorrelation and sample reweighting, flatness-aware generation through suboptimal embedding averaging, and normalization-based alignment between training and test distributions—with the goal of improving robustness, effectiveness, and generalization under train–test shift in tabular data (Ying et al., 27 Aug 2025).

1. Problem formulation and conceptual scope

SAFT is motivated by the observation that tabular learning often depends heavily on feature engineering, with transformed features such as $f_1+f_2$, $f_1/f_3-f_2$, and $\sqrt{f_1}$ frequently determining downstream performance. In the authors’ DSTL formulation, the core difficulty is that a transformation appearing useful on the training set may encode spurious relationships that do not persist under distribution shift. The resulting transformed space can therefore fail to generalize (Ying et al., 27 Aug 2025).

The framework is specifically positioned against feature-transformation methods that are search-based and discrete. In that view, the space of transformations is combinatorial, expensive to explore, and difficult to regularize for robustness. SAFT replaces direct symbolic search with a continuous latent-space procedure: learn a representation of transformed feature sets, optimize that representation by gradient ascent, and decode optimized embeddings back into feature-cross sequences. This suggests a change in emphasis from enumerating candidate formulas to learning a geometry over feature transformations in which robustness can be imposed directly.

The paper’s notion of shift is operationalized through an experimental protocol that iteratively creates 80/20 train/test splits along each feature and applies a Kolmogorov–Smirnov test at 95% confidence, excluding datasets with no detected shift. The target setting is therefore train–test distribution mismatch in tabular covariates rather than a labeled-target domain-adaptation setup. The paper also argues that feature importance itself can drift: its appendix case study reports that features important on the training set differ from those important on the test set, indicating that transformations exploiting training correlations alone are brittle (Ying et al., 27 Aug 2025).

2. Representation-generation architecture

SAFT assumes access to many examples of transformed feature sets paired with their measured downstream performance. A feature cross denotes a transformed feature such as $f_1+f_2$ or $\sin(f_3)$; a feature cross sequence is a tokenized postfix-expression representation of one or more feature crosses; and a transformed feature set is the feature set produced by applying such rules to the original variables (Ying et al., 27 Aug 2025).

The architecture has three learned components. First, each transformed feature set is converted into a feature-feature similarity graph $\mathcal{G}$, whose nodes are features, node attributes are feature values, and edges reflect pairwise similarities such as cosine similarity. The encoder maps this graph to an embedding,

$E = \phi_\theta(\mathcal{G}).$

The text states that the encoder is inspired by GraphSAGE-style node sampling and multi-hop neighborhood aggregation to accommodate dynamic graph size and topology, although the main text does not provide a layer-by-layer aggregation equation (Ying et al., 27 Aug 2025).

Second, an evaluator predicts downstream performance from the embedding. Its stated loss is

$$\mathcal{L}_{est} = \frac{1}{N}\sum_{i=1}^N \left(p_i - \omega_{\theta}(E_i)\right)^2,$$

where $p_i$ is the ground-truth downstream performance of the $i$-th transformed feature set and $f_1/f_3-f_2$0 is the evaluator network. Third, a single-layer LSTM + softmax decoder reconstructs a feature-cross sequence from the embedding. For token $f_1/f_3-f_2$1 in sequence $f_1/f_3-f_2$2,

$f_1/f_3-f_2$3

and the reconstruction loss is

$f_1/f_3-f_2$4

Joint training uses

$f_1/f_3-f_2$5

The appendix also reports practical settings: graph edges are created with threshold at the 95th percentile of similarity values; node attributes are mapped to 64-dimensional embeddings; the encoder uses a 2-layer GNN and 2-layer projection head; the evaluator is a 2-layer feed-forward network with hidden dimension 200 per layer; the decoder is a 1-layer LSTM; training uses batch size 256, 500 epochs, and learning rate range 0.001–0.0005 (Ying et al., 27 Aug 2025).

3. Data collection and training workflow

Before latent-space learning begins, SAFT constructs a supervised knowledge base of feature-set/performance pairs using reinforcement learning. The appendix states that three RL agents are used: a head feature agent, an operation agent, and a tail feature agent. These agents iteratively build transformed features, with reward defined as the downstream performance improvement of the current explored feature set relative to the previous one. Exploration is $f_1/f_3-f_2$6-greedy and DQN-style, and the top 5,000 feature set–performance pairs are retained as training data (Ying et al., 27 Aug 2025).

The overall workflow can be divided into three phases. In the first phase, SAFT normalizes data, builds graphs from transformed feature sets, and jointly trains encoder, evaluator, and decoder. In the second phase, it selects the top-$f_1/f_3-f_2$7 embeddings among observed training transformations, optimizes them by gradient ascent in latent space, performs flatness-aware cyclic averaging, and decodes the resulting embeddings back into feature-cross sequences. In the third phase, it normalizes test data, applies the learned or generated feature-cross sequences, and denormalizes the transformed outputs (Ying et al., 27 Aug 2025).

The latent-space generation step is central to the framework’s continuous optimization perspective. Starting from an embedding $f_1/f_3-f_2$8, SAFT updates it by

$f_1/f_3-f_2$9

with $\sqrt{f_1}$0 a step size and $\sqrt{f_1}$1 the evaluator. The paper states that the process is initialized from the top-$\sqrt{f_1}$2 embeddings among observed training transformations, producing a set

$\sqrt{f_1}$3

Optimized embeddings are then decoded autoregressively until an $\sqrt{f_1}$4 token is emitted, with $\sqrt{f_1}$5 separating multiple feature crosses inside a sequence (Ying et al., 27 Aug 2025).

This training-and-generation loop is what distinguishes SAFT from conventional automated feature engineering. Rather than searching over raw operators during optimization, SAFT first amortizes knowledge of transformations into a representation space and then searches that learned space directly.

4. The three shift-aware robustness mechanisms

The first mechanism is shift-resistant representation via embedding decorrelation and sample reweighting. The paper argues that under distribution shift, a learned embedding space can entangle invariant signals with spurious signals. To reduce this contamination, SAFT introduces a bilevel sample-weighted decorrelation objective. The inner problem optimizes a sample-weight vector $\sqrt{f_1}$6 by minimizing weighted cross-dimension covariance,

$\sqrt{f_1}$7

subject to

$\sqrt{f_1}$8

The weighted partial cross-covariance is

$\sqrt{f_1}$9

where $f_1+f_2$0 and $f_1+f_2$1 are random Fourier feature functions. The resulting optimal weights are then used to reweight evaluator training: $f_1+f_2$2 and the final weighted objective is

$f_1+f_2$3

The intended effect is to suppress spurious correlations unstable across distributions while amplifying more stable structure (Ying et al., 27 Aug 2025).

The second mechanism is flatness-aware generation through suboptimal embedding averaging. After evaluator training, SAFT performs gradient ascent in embedding space, but does not simply keep the final iterate. Instead, it uses a cyclic learning-rate schedule, lets the search oscillate around an optimum, and averages suboptimal embeddings to approximate the center of a flat high-performance region. The averaging update in the pseudocode is

$f_1+f_2$4

The paper’s motivation is that, under shift, the train-optimal embedding and the test-optimal embedding may differ; a flatter optimum should degrade less under perturbation. This suggests that robustness is being imposed not only on the representation itself but also on the geometry of the latent optimization landscape (Ying et al., 27 Aug 2025).

The third mechanism is normalization-based alignment between training and test distributions. SAFT introduces normalization in preprocessing and denormalization in postprocessing, mentioning z-score normalization and reverse z-score in prose. The procedure is described as normalizing both training and testing data before feature transformation, training and generating transformations in that normalized space, applying the learned transformation sequence to normalized test data, and then denormalizing the transformed test data afterward. The paper does not provide an explicit normalization equation, nor does it fully specify whether test normalization uses test-time statistics, training statistics, or another estimate. Its appendix nevertheless reports that normalization better aligns marginal and some joint distributions between training and test sets on example datasets (Ying et al., 27 Aug 2025).

5. Empirical performance and ablation evidence

SAFT is evaluated on 16 public tabular datasets from UCI and OpenML, spanning both classification and regression, with Random Forests used as the main downstream learner to reduce confounding from model choice. For classification, the primary metric is F1-score; for regression, the primary metric is $f_1+f_2$5-RAE. Baselines include RDG, ERG, LDA, AFAT, NFS, TTG, GRFG, MOAT, NEAT, and ELLM-FT (Ying et al., 27 Aug 2025).

Representative regression results include Housing Boston at $f_1+f_2$6, exceeding the next-best $f_1+f_2$7; openml_586 at $f_1+f_2$8, exceeding NEAT $f_1+f_2$9 and MOAT $\sin(f_3)$0; openml_589 at $\sin(f_3)$1; openml_607 at $\sin(f_3)$2, ahead of ELLM-FT $\sin(f_3)$3; openml_616 at $\sin(f_3)$4; and openml_637 at $\sin(f_3)$5. Representative classification results include SpectF at $\sin(f_3)$6, ahead of ERG/ELLM-FT $\sin(f_3)$7; UCI Credit at $\sin(f_3)$8, ahead of RDG $\sin(f_3)$9; Wine Quality Red at $\mathcal{G}$0, ahead of MOAT $\mathcal{G}$1 and ELLM-FT $\mathcal{G}$2; Wine Quality White at $\mathcal{G}$3; PimaIndian at $\mathcal{G}$4, ahead of MOAT $\mathcal{G}$5; and German Credit at $\mathcal{G}$6, ahead of ELLM-FT $\mathcal{G}$7 (Ying et al., 27 Aug 2025).

The ablation study removes each of SAFT’s three robustness components: SAFT-f removes Flatness-Aware Gradient Ascent / weight averaging, SAFT-n removes normalization, and SAFT-w removes optimization of graph/sample weights. Across three classification and three regression datasets, all ablations underperform full SAFT. The reported interpretation is that learned graph/sample weights help eliminate spurious correlations, flatness-aware ascent reduces local brittleness, and normalization improves alignment of training/test statistics under OOD shift (Ying et al., 27 Aug 2025).

The paper also includes robustness analyses beyond the main table. On Wine Quality White, five different split schemes are tested, with SAFT remaining robust in Precision, Recall, and F1. On Wine Quality Red/White, the transformed features are evaluated with RF, SVM, KNN, DT, LASSO, and Ridge, with reported scores $\mathcal{G}$8, $\mathcal{G}$9, $E = \phi_\theta(\mathcal{G}).$0, $E = \phi_\theta(\mathcal{G}).$1, $E = \phi_\theta(\mathcal{G}).$2, and $E = \phi_\theta(\mathcal{G}).$3, respectively. A feature-importance case study on openml_616 reports that a generated feature space using only 7 features, including 2 new ones, improves ML performance by 29.24%. The appendix further notes that training time increases with number of features, that RL data collection is expensive though offline, and that inference time is relatively low after training converges (Ying et al., 27 Aug 2025).

6. Relation to adjacent work and acronym ambiguity

Within the supplied literature, SAFT has multiple meanings, and disambiguation is necessary. In the tabular-learning sense, Shift-Aware Feature Transformation denotes the DSTL framework described above (Ying et al., 27 Aug 2025). By contrast, SAFT: Structure-Aware Fine-Tuning of LLMs for AMR-to-Text Generation is a graph-to-LLM adaptation method that injects magnetic-Laplacian positional encodings into decoder-only LLMs without architectural change; it is explicitly about Structure-Aware Fine-Tuning, not shift-aware feature transformation (Kamel et al., 15 Jul 2025).

Two nearby bodies of work nonetheless illuminate the broader intellectual neighborhood of shift-aware transformation. “Sharpness & Shift-Aware Self-Supervised Learning” develops a shift-aware contrastive learning framework in which the crucial mismatch is between an ideal semantic positive-pair distribution and the practical augmentation-induced one. Its proposed Sharpness & Shift-Aware Contrastive Learning (SSA-CLR) uses a Fourier-based positive-sample transformation and a sharpness-aware objective, but it does not define a method literally called Shift-Aware Feature Transformation (Tran et al., 2023). In temporal tabular learning, “Feature-aware Modulation for Learning from Temporal Tabular Data” proposes a time-conditioned feature transformation

$E = \phi_\theta(\mathcal{G}).$4

with the stated goal of aligning feature semantics across time; the paper never uses the term SAFT, but it is naturally interpretable as a temporal-shift-aware feature warping mechanism (Cai et al., 3 Dec 2025).

Taken together, these distinctions matter because “SAFT” is not a stable acronym across arXiv. In the tabular-learning literature, it denotes a framework for robust symbolic feature generation under train–test distribution shift. In adjacent areas, the acronym instead names methods for structure-aware LLM adaptation, semantic-aware adversarial fine-tuning, sensitivity-aware wireless transmission, or even the Special Affine Fourier Transform. A plausible implication is that the tabular SAFT of (Ying et al., 27 Aug 2025) is best understood not as a generic label for any shift-aware method, but as a specific encoder–evaluator–decoder framework that unifies robust representation learning, flatness-aware embedding optimization, and normalization-based alignment for symbolic feature transformation under distribution shift.