HyFormer: Hybrid Transformer for Molecule Modeling

• HyFormer is a hybrid Transformer model that integrates joint molecule generation and property prediction using alternating attention masks and multi-task pretraining.
• It employs a unified Transformer backbone with mode switching via task tokens and alternating causal and bidirectional attention masks for efficient training.
• Empirical results show improved OOD predictions, enhanced representation learning, and superior performance in drug design tasks such as antimicrobial peptide discovery.

HyFormer refers to a family of architectures and models unified by the use of hybrid or synergistic combinations of Transformer-based mechanisms with auxiliary inductive biases or architectural elements, tailored for joint tasks, cross-modality interactions, or multimodal fusion. Several distinct models bearing the HyFormer name have emerged across disciplines, including molecular modeling, recommender systems, and biomedical image analysis. The following overview focuses on the original Transformer-based joint model for molecule generation and property prediction, as introduced in "Synergistic Benefits of Joint Molecule Generation and Property Prediction" (Izdebski et al., 23 Apr 2025), and contextualizes its impact and technical details within the broader HyFormer paradigm.

1. Motivation and Problem Statement

Traditional deep learning strategies in molecular science have bifurcated into generative models (for molecule synthesis and structure proposal) and predictive models (for property estimation from molecular representations). Conventional approaches treat these as separate tasks, missing inter-task synergies, and often suffer from inefficient transfer or compromised generalization—particularly in out-of-distribution (OOD) settings. HyFormer addresses this by learning the joint distribution $p(x, y)$, where $x$ is the data sample (e.g., molecule as SMILES) and $y$ is its property vector or scalar label. This enables a unified architecture achieving simultaneous optimization for both generation and prediction, with empirical evidence of cross-task benefits such as improved conditional sampling, enhanced OOD prediction, and enriched representation learning (Izdebski et al., 23 Apr 2025).

2. Architectural Design and Alternating Attention Mechanism

Input Representation

HyFormer utilizes character-level tokenization for molecular SMILES strings, extending methods from Schwaller et al., and, for peptide applications, ESM-2-based sequence tokenization. Each input is prepended with a task token: [LM] for unconditional generation, [MLM] for masked reconstruction, and [PRED] for property prediction.

Transformer Backbone

Model variants typical of HyFormer are summarized as follows:

Variant	#Params	Embedding dim	Hidden dim	#Layers	#Heads
Small	8.7 M	256	1024	8	8
Base	50 M	512	2048	8	8

A single shared backbone is used, with mode switching via the prepended task token.

Alternating Attention Mask

Self-attention masks alternate between causal (autoregressive LM) and bidirectional (MLM or prediction) at each forward pass.

$$\mathrm{Attention}(Q,K,V) = \mathrm{softmax}\left(\frac{QK^T}{\sqrt{d}} + M\right)V$$

where $M$ is varied:

• Causal mask $M_{\to}$: $M_{ij} = 0$ for $i\geq j$, $-\infty$ for $i < j$ (autoregressive generation)
• Bidirectional mask $M_{\leftrightarrow} = 0$ (full attention)

Routing is determined by the sampled task token, with representations dispatched to either the autoregressive head or predictive head.

3. Joint Pre-Training Scheme

HyFormer's unified pre-training objective mixes three loss functions:

• Language Modeling (LM):

$$\mathcal{L}_{\mathrm{LM}}(\theta) = -\mathbb{E}_{x\sim p(x)} \sum_{t=1}^T \log p_\theta(x_t | x_{<t})$$

• Masked LM (MLM) Reconstruction:

$$\mathcal{L}_{\mathrm{MLM}}(\theta) = -\mathbb{E}_{x\sim p(x)}\, \mathbb{E}_{\mathcal{M}}\left [ \sum_{i\in\mathcal{M}} \log p_\theta(x_i | x_{\backslash\mathcal{M}}) \right ]$$

• Property Prediction (PRED):

$$\mathcal{L}_{\mathrm{PRED}}(\theta) = -\mathbb{E}_{(x,y)\sim p(x,y)} \log p_\theta(y | x)$$

(or $\|\!f(x) - y\!\|^2$ for regression)

The total objective:

$$\mathcal{L}_{\mathrm{pre}} = \mathcal{L}_{\mathrm{LM}} + \mu\,\mathcal{L}_{\mathrm{MLM}} + \eta\,\mathcal{L}_{\mathrm{PRED}}$$

where $(\mu, \eta)$ are tuned. Typical scheduling draws [LM] 80% of the time, [MLM] or [PRED] 10% each.

Fine-tuning mixes generative and predictive loss with an adjustable $\lambda$:

$$\mathcal{L}_{\mathrm{fine}} = \mathcal{L}_{\mathrm{LM}} + \lambda\,\mathcal{L}_{\mathrm{PRED}}$$

4. Training Setup and Hyperparameters

• Optimizer: AdamW ($\beta_1=0.9$, $\beta_2=0.95$)
• Pre-training: batch size 512, learning rate $6\times10^{-4}$, weight decay 0.1, cosine annealing with 2 500-step warmup
• Steps: 50 K (small) to 200 K (base)
• Fine-tuning: grid search over batch sizes $\{16,64,128,256\}$, learning rates in $[10^{-5},10^{-3}]$, weight decay $[10^{-2},3\times10^{-1}]$, early stopping after 5 epochs without improvement

5. Synergistic Empirical Benefits

Unconditional Molecule Generation (GuacaMol)

On the GuacaMol dataset, HyFormer (8.7 M params) outperforms SMILES-based baselines for FCD and KL:

Model	FCD (↑)	KL (↑)	Validity (↑)
MolGPT	0.907	0.992	0.981
Hyformer	0.922	0.996	0.970

Molecular Property Prediction (MoleculeNet, Scaffold Split)

Pre-trained on 19M samples, HyFormer (Base) attains top results:

Dataset	Metric	Best Predictive	Graph2Seq (joint)	HyFormer (joint)
ESOL	RMSE↓	0.788 (Uni-Mol)	0.860	0.774
BBBP	AUC↑	72.9 (Uni-Mol)	72.8	75.9
ClinTox	AUC↑	91.9 (Uni-Mol)	–	99.2

HyFormer surpasses joint Graph2Seq in 8/10 MoleculeNet tasks and is competitive with SOTA predictive-only models.

Out-of-Distribution (OOD) Predictions

On Lo-Hi Hit ID tasks (Tanimoto similarity $<$ 0.4 to train), HyFormer achieves highest AUPRC on 3/4 targets.

Representation Learning

Frozen HyFormer embeddings rank 1st or 2nd by linear and kNN probes in 8/10 tasks, supporting richer, more linearly separable representation spaces.

6. Application to Drug Design: Antimicrobial Peptide Discovery

Data and Training

• 1.1M peptides (HydrAMP + AMPSphere), 39 physical descriptors as pre-training labels
• Fine-tuned as binary AMP classifier ("HyFormer AMP") and MIC regressor ("HyFormer MIC")

Conditional Sampling

Best-of-$K$ sampling for 50 K candidates, acceptance if classifier logit $\geq 0.5$ (AMP) or MIC $\leq 0.3$.

Metrics and Comparative Results

Model	Perp↓	Entr↑	JS-3↑	JS-6↑	$P_{\rm amp}$	$P_{\rm mic}$
AMP-Diffusion	12.84	3.17	0.99	0.028	0.81	0.50
HyFormer AMP	1.59	3.88	0.99	0.026	0.84	–
HyFormer MIC	1.62	4.72	0.83	0.030	–	0.71

Generated amino-acid and physicochemical distributions for sampled peptides match or exceed "true AMP" signals, with attention maps highlighting known antimicrobial residue motifs.

7. Limitations and Prospects for Extension

HyFormer currently ingests 1D sequences and precomputed descriptors, whereas integration of 3D structural information may further enhance performance. Throughput is bounded by model size and the adopted sampling (best-of-$K$), suggesting future work in more efficient conditional sampling. Optimal joint loss weights and sampling probabilities require task-dependent tuning; adaptive or learned scheduling mechanisms constitute a future research direction.

---

HyFormer demonstrates that alternating-mask Transformers with unified joint training not only match or surpass baseline generative and predictive models within molecule modeling but also exhibit substantial gains under OOD scenarios, richer representational capacity, and tangible impact in practical drug-design tasks such as antimicrobial peptide discovery (Izdebski et al., 23 Apr 2025).