Perturbation-Aware Differential Transformer

• PAD-Transformer is a perturbation-conditioned backbone that integrates graph-masked gene encoding and differential attention for accurate single-cell perturbation prediction.
• It employs dual token streams for control and evolving states, repeatedly injecting perturbation and time embeddings to capture nuanced distribution-level dynamics.
• Its subtraction-based differential attention mechanism suppresses noise by canceling irrelevant gene correlations, leading to improved performance metrics in scDFM evaluations.

Perturbation-Aware Differential Transformer (PAD-Transformer) is a perturbation-conditioned transformer backbone introduced within the scDFM framework for single-cell perturbation prediction. In that setting, the model receives a control-cell transcriptomic state, a perturbation condition, an intermediate state along a conditional flow trajectory, a time variable, and a gene-gene graph prior, and it uses graph-masked gene encoding together with differential attention to parameterize a conditional velocity field for distribution-level generation (Yu et al., 6 Feb 2026). The term also has a broader architectural interpretation: earlier work on Differential Transformer established subtraction-based attention as a mechanism for cancelling irrelevant context (Ye et al., 2024), mechanistic analysis clarified its links to signed relevance and head diversification (Kong et al., 22 May 2025), and channel-decoding work provided a closely related perturbation-sensitive, graph-structured differential-attention message-passing design even though it did not use the PAD-Transformer name (Lau et al., 19 Sep 2025).

1. Origin, naming, and problem setting

PAD-Transformer appears explicitly in the systems-biology context of scDFM, where the task is to predict the post-perturbation gene-expression distribution $p_\theta(x \mid c_x, c_p)$ from a control state $c_x$ and a perturbation condition $c_p$ (Yu et al., 6 Feb 2026). The motivation is specific to single-cell perturbation prediction: single-cell RNA-seq measurements are described as noisy, sparse, and zero-inflated, perturbation effects are nonlinear and context-dependent, and the supervision is unpaired because the same cell cannot be observed before and after perturbation. The stated consequence is that models assuming cell-level correspondences can miss distribution-level shifts such as altered variance, multimodal responses, and subpopulation changes.

In that literature, “perturbation-aware” does not denote certified robustness or adversarial defense. It refers to explicit conditioning on perturbation identity together with architectural mechanisms intended to remain informative under sparse and noisy measurements. PAD-Transformer is therefore best understood as a perturbation-conditioned representational backbone, not as a standalone generative model and not as a generic robustness framework.

The name also sits within a broader differential-attention lineage. Differential Transformer introduced the subtraction of two softmax attention maps to amplify relevant context while cancelling noise (Ye et al., 2024). The later analysis in DEX retained this subtraction-centered view but argued that the main gains arise from negative attention, reduced head redundancy, and improved learning dynamics rather than from sparsity alone (Kong et al., 22 May 2025). In a different domain, the Differential-Attention Message Passing Transformer for channel decoding implemented graph-masked differential attention, soft syndrome reasoning, and perturbation-sensitive message passing, making it a strong conceptual precursor even though the paper explicitly does not use the phrase “Perturbation-Aware Differential Transformer” (Lau et al., 19 Sep 2025).

2. Architectural composition in scDFM

PAD-Transformer operates on a gene-token representation of a cell. For a selected subset of genes $\mathcal{S}\subseteq \mathcal{G}$, the control and current trajectory states are embedded as

$$h_c = E_v\!\big(c_x^{(\mathcal{S})}\big) + E_g(\mathcal{S}), \qquad h_t^0 = E_v\!\big(x_t^{(\mathcal{S})}\big) + E_g(\mathcal{S}).$$

Here, $E_v$ maps scalar gene-expression values to $d$-dimensional embeddings, while $E_g$ provides contextualized gene identity embeddings (Yu et al., 6 Feb 2026). The architecture therefore maintains two aligned token streams: a control context sequence and an evolving perturbed-state sequence.

Gene structure is incorporated through a co-expression graph. The edge weight between genes $i$ and $j$ is defined by the absolute Pearson correlation,

$c_x$0

followed by KNN sparsification to form a sparse adjacency matrix $c_x$1 (Yu et al., 6 Feb 2026). This graph is used as a sparse attention mask in the gene encoder. The paper is explicit that the graph enters primarily through masked attention; it does not define graph positional encodings, graph bias terms added to attention logits, or explicit graph message-passing layers beyond this mask.

Conditioning is layered rather than front-loaded. PAD-Transformer uses a perturbation embedding $c_x$2 and a time embedding

$c_x$3

with the timestep embedding providing adaLN-Zero modulation for each self- and cross-differential attention layer (Yu et al., 6 Feb 2026). At every layer $c_x$4, perturbation information is re-injected through

$c_x$5

This repeated injection is central to the model’s perturbation-aware characterization: perturbation identity is not merely an input token but a persistent conditioning signal.

The core update in each layer has two stages. First, self-differential attention refines the evolving perturbed representation: $c_x$6 Second, cross-differential attention uses the control representation as reference: $c_x$7 After $c_x$8 layers, the decoder consumes the final latent together with the perturbation embedding: $c_x$9 In the main text this is presented as the predicted perturbed state, while the appendix clarifies that the operational role is to decode the conditional velocity $c_p$0 required by the flow-matching objective (Yu et al., 6 Feb 2026).

3. Differential attention as the defining operator

The defining attention mechanism is subtraction-based. PAD-Transformer computes

$c_p$1

then forms

$c_p$2

with $c_p$3, $c_p$4, and $c_p$5 (Yu et al., 6 Feb 2026). Relative to standard transformer attention, which uses a single nonnegative attention map, this produces a signed weighting scheme in which one softmax branch serves as a suppressive counterpart to the other.

The immediate purpose in scDFM is to reduce over-attending to irrelevant genes in sparse, noisy single-cell data. The broader differential-attention literature supplies a more general interpretation. Differential Transformer described the subtraction of two attention maps as a form of noise cancellation that amplifies relevant context while suppressing shared nuisance structure (Ye et al., 2024). The mechanistic study in DEX then argued that the empirical benefits are better explained by three factors: enhanced expressivity via negative attention, reduced redundancy among attention heads, and improved learning dynamics (Kong et al., 22 May 2025). That analysis matters for PAD-Transformer because it suggests that subtraction is not merely a sparsification heuristic. It enables signed relevance modeling, so a head can represent that a context should be actively discounted rather than simply weakly weighted.

This also clarifies a frequent misconception. PAD-Transformer does not compute an explicit algebraic difference between a control embedding and a perturbed embedding. The “differential” element lies in the internal attention operator, not in a direct feature subtraction such as $c_p$6. The control state enters as the key-value reference for cross-differential attention, and perturbation-specific deviations are learned implicitly through that interaction (Yu et al., 6 Feb 2026).

4. Function inside conditional flow matching

PAD-Transformer is embedded in a larger conditional generative framework rather than used as an isolated predictor. scDFM defines a time-dependent ODE

$c_p$7

where $c_p$8 is the current state, $c_p$9 is the control expression, and $\mathcal{S}\subseteq \mathcal{G}$0 is the perturbation condition (Yu et al., 6 Feb 2026). The model uses a linear interpolation path

$\mathcal{S}\subseteq \mathcal{G}$1

samples an intermediate point $\mathcal{S}\subseteq \mathcal{G}$2, and trains the backbone to predict the reference flow along that path.

The conditional flow matching objective is

$\mathcal{S}\subseteq \mathcal{G}$3

To connect local dynamics to endpoint quality, scDFM also uses the one-step approximation

$\mathcal{S}\subseteq \mathcal{G}$4

Distribution-level supervision is added through MMD. With the multi-kernel mixture

$\mathcal{S}\subseteq \mathcal{G}$5

the total objective is

$\mathcal{S}\subseteq \mathcal{G}$6

This division of labor is explicit: conditional flow matching supplies local trajectory supervision, while MMD enforces global alignment between generated and real perturbed populations (Yu et al., 6 Feb 2026).

A second misconception follows from this embedding. PAD-Transformer does not directly generate the final perturbed sample in the reported formulation; it parameterizes the conditional vector field used by the ODE solver. At inference time the model initializes $\mathcal{S}\subseteq \mathcal{G}$7, computes the graph-aware tokenization for the selected gene subset, and integrates forward from $\mathcal{S}\subseteq \mathcal{G}$8 to $\mathcal{S}\subseteq \mathcal{G}$9, using Euler by default and optionally Heun (Yu et al., 6 Feb 2026).

5. Robustness mechanisms, training protocol, and empirical behavior

The robustness of PAD-Transformer to sparsity and noise is attributed to several coordinated mechanisms. Differential attention is intended to suppress noisy or non-responsive gene relations. Cross-differential attention uses the control representation $$h_c = E_v\!\big(c_x^{(\mathcal{S})}\big) + E_g(\mathcal{S}), \qquad h_t^0 = E_v\!\big(x_t^{(\mathcal{S})}\big) + E_g(\mathcal{S}).$$0 as a reference, which is especially useful when absolute expression values are noisy but relative deviations are biologically informative. Gene-graph masking restricts the gene encoder to biologically plausible neighbors, reducing the space of possible spurious correlations. MMD then adds batch-level distributional supervision, which is important when individual cells are noisy observations (Yu et al., 6 Feb 2026).

The reported implementation details are concrete. On Norman, scDFM uses Adam, an initial learning rate of $$h_c = E_v\!\big(c_x^{(\mathcal{S})}\big) + E_g(\mathcal{S}), \qquad h_t^0 = E_v\!\big(x_t^{(\mathcal{S})}\big) + E_g(\mathcal{S}).$$1, cosine decay to $$h_c = E_v\!\big(c_x^{(\mathcal{S})}\big) + E_g(\mathcal{S}), \qquad h_t^0 = E_v\!\big(x_t^{(\mathcal{S})}\big) + E_g(\mathcal{S}).$$2, batch size $$h_c = E_v\!\big(c_x^{(\mathcal{S})}\big) + E_g(\mathcal{S}), \qquad h_t^0 = E_v\!\big(x_t^{(\mathcal{S})}\big) + E_g(\mathcal{S}).$$3, training for $$h_c = E_v\!\big(c_x^{(\mathcal{S})}\big) + E_g(\mathcal{S}), \qquad h_t^0 = E_v\!\big(x_t^{(\mathcal{S})}\big) + E_g(\mathcal{S}).$$4 steps, MMD weight $$h_c = E_v\!\big(c_x^{(\mathcal{S})}\big) + E_g(\mathcal{S}), \qquad h_t^0 = E_v\!\big(x_t^{(\mathcal{S})}\big) + E_g(\mathcal{S}).$$5, a kNN graph with $$h_c = E_v\!\big(c_x^{(\mathcal{S})}\big) + E_g(\mathcal{S}), \qquad h_t^0 = E_v\!\big(x_t^{(\mathcal{S})}\big) + E_g(\mathcal{S}).$$6, hidden size $$h_c = E_v\!\big(c_x^{(\mathcal{S})}\big) + E_g(\mathcal{S}), \qquad h_t^0 = E_v\!\big(x_t^{(\mathcal{S})}\big) + E_g(\mathcal{S}).$$7, $$h_c = E_v\!\big(c_x^{(\mathcal{S})}\big) + E_g(\mathcal{S}), \qquad h_t^0 = E_v\!\big(x_t^{(\mathcal{S})}\big) + E_g(\mathcal{S}).$$8 PAD-Transformer layers, $$h_c = E_v\!\big(c_x^{(\mathcal{S})}\big) + E_g(\mathcal{S}), \qquad h_t^0 = E_v\!\big(x_t^{(\mathcal{S})}\big) + E_g(\mathcal{S}).$$9 attention heads, and dropout $E_v$0 on attention and MLP blocks. Inference uses Euler rollout with $E_v$1 steps over $E_v$2, so $E_v$3. ComboSciPlex uses the same optimizer and backbone hyperparameters, but drug perturbations are represented with a dedicated embedding table rather than tying perturbation embeddings to gene identity (Yu et al., 6 Feb 2026).

Empirically, the strongest quantitative results are reported for the full scDFM system rather than PAD-Transformer in isolation. On the Norman additive split, compared with CellFlow, MSE improves from $E_v$4 to $E_v$5, MAE from $E_v$6 to $E_v$7, DS from $E_v$8 to $E_v$9, and Pearson $d$0 from $d$1 to $d$2. On the Norman holdout split, scDFM reports for single perturbations L2 $d$3, MSE $d$4, MAE $d$5, DE-Spearman $d$6, DS $d$7, and Pearson $d$8; for double perturbations it reports L2 $d$9, MSE $E_g$0, MAE $E_g$1, DE-Spearman $E_g$2, DS $E_g$3, and Pearson $E_g$4. On ComboSciPlex it reports L2 $E_g$5, MSE $E_g$6, MAE $E_g$7, DE-Spearman $E_g$8, Pearson $E_g$9, and DS $i$0 (Yu et al., 6 Feb 2026).

The ablation statements are qualitatively specific even where exact numbers are absent. Removing the gene-gene mask reduces correlation with ground truth, removing the Differential Transformer backbone also reduces correlation, and removing MMD causes the sharpest decline together with visible manifold mismatch in UMAP visualizations. This indicates that PAD-Transformer contributes through both its graph-constrained encoding and its differential-attention backbone, while the surrounding distributional objective remains essential for population-level fidelity (Yu et al., 6 Feb 2026).

6. Related lineages, scope conditions, and limitations

PAD-Transformer has a narrow literal meaning and a broader architectural meaning. Literally, it is the single-cell perturbation backbone in scDFM (Yu et al., 6 Feb 2026). More broadly, it belongs to a family of perturbation-sensitive differential-attention architectures in which a subtractive attention operator is combined with explicit structure and noise-aware conditioning. The channel-decoding DiffMPT model is especially close to this broader reading. It uses noisy-input decoding over AWGN, absolute LLRs, soft syndromes, BP-inspired alternating variable/check updates, Tanner-graph masks, differential attention, and a differentiable syndrome-validity loss, while explicitly reformulating decoding as prediction of multiplicative noise rather than direct bit prediction (Lau et al., 19 Sep 2025). That work suggests a general architectural pattern: perturbation-sensitive inputs, structured dual tokenization, graph-constrained differential attention, and a task-specific global consistency objective.

At the same time, PAD-Transformer should not be conflated with all uses of the adjective “perturbation-aware.” The transform-dependent adversarial-attack literature studies perturbations indexed by image transformations rather than biological interventions (Tan et al., 2024). The differential-privacy literature uses “perturbation” to denote noise injected into optimization and motivates robustness to parameter-space perturbations rather than differential attention (Wang et al., 2024). These lines are conceptually adjacent but architecturally distinct.

The limitations of PAD-Transformer as currently instantiated are explicit. Its graph prior is Pearson-correlation-based and therefore captures mainly linear co-expression rather than causal or nonlinear regulatory structure. The interpolation path in log-expression space is acknowledged as biologically simplistic. The reported setup predicts only selected genes rather than a full vocabulary, and a full-vocabulary imputation head is mentioned as possible but not used. The architecture has not yet been validated on larger multi-context datasets such as ARC-state. Computationally, masked attention in the gene encoder and differential attention in PAD-Transformer each scale as $i$1 per layer, and dynamic MMD adds $i$2 batchwise pairwise computation, which is why the model trains on sampled gene subsets $i$3 rather than all genes simultaneously (Yu et al., 6 Feb 2026).

A final scope condition concerns robustness claims. None of the cited PAD-relevant works establishes adversarial or certified robustness for this architecture. In scDFM, perturbation-awareness means conditioning on biological interventions under sparse and noisy measurements. In Differential Transformer, it means subtraction-based suppression of irrelevant context. In DiffMPT, it means noise-aware decoding with graph-structured differential message passing. A plausible implication is that these strands define a common design principle—explicit nuisance-sensitive conditioning plus subtractive attention—but the current literature still treats that principle through domain-specific instantiations rather than a single unified formal framework (Ye et al., 2024).