PAIR-Former: Budgeted Relational MIL for miRNA Target Prediction

TL;DR

The paper addresses functional miRNA–mRNA target prediction where each transcript yields a large, heavy-tailed pool of candidate target sites (CTSs) but only a bag-level label is observed. It formalizes this as Budgeted Relational MIL (BR-MIL) and introduces PAIR-Former, which uses a cheap full-pool scan, a CPU-based budgeted selector (STSelector), and a Set Transformer aggregator on a selected subset of CTSs. The authors provide theoretical results linking the budget $K$ to approximation error and generalization, and demonstrate empirically that PAIR-Former outperforms max-pooling baselines under a practical budget (K*=64) on miRAW, with a tunable accuracy–compute trade-off as $K$ varies. It also discusses the broader implications of budgeted MIL for efficiency in large candidate pools and presents a three-stage training regimen including teacher–student distillation. Overall, the work advances budget-aware relational reasoning in weakly supervised biological prediction tasks and offers a scalable, practical framework for heavy-tailed CTS data.

Abstract

Functional miRNA--mRNA targeting is a large-bag prediction problem: each transcript yields a heavy-tailed pool of candidate target sites (CTSs), yet only a pair-level label is observed. We formalize this regime as \emph{Budgeted Relational Multi-Instance Learning (BR-MIL)}, where at most $K$ instances per bag may receive expensive encoding and relational processing under a hard compute budget. We propose \textbf{PAIR-Former} (Pool-Aware Instance-Relational Transformer), a BR-MIL pipeline that performs a cheap full-pool scan, selects up to $K$ diverse CTSs on CPU, and applies a permutation-invariant Set Transformer aggregator on the selected tokens. On miRAW, PAIR-Former outperforms strong pooling baselines at a practical operating budget ($K^\star{=}64$) while providing a controllable accuracy--compute trade-off as $K$ varies. We further provide theory linking budgeted selection to (i) approximation error decreasing with $K$ and (ii) generalization terms governed by $K$ in the expensive relational component.

Figures (5)

• Figure 1: Overview of PAIR-Former under BR-MIL.Top: Inference pipeline. Candidate CTSs are generated on the $3^{\prime}\mathrm{UTR}$ using a TargetNet-compatible ESA scan/filter ($s_i^{\mathrm{esa}}\ge 6$), producing a variable-size heavy-tailed bag with $n$ candidates. A cheap student encoder $\tilde{e}_{\tilde{\theta}}$ scans all $n$ instances to obtain cheap logits $\tilde{z}_i$ and embeddings $\tilde{h}_i$. Under a strict per-bag budget $K=\min(\texttt{kmax},n)$, STSelector selects a subset $S$ ($|S|\le K$) by combining Top-$K_1$ exploitation (by $\tilde{z}_i$) and diversity over transcript position $p_i$ and embedding redundancy. Only selected sites are processed by the expensive encoder $e_\theta$ to form CTS tokens (e.g., concatenating $h_i$, $z_i$, $s_i^{\mathrm{esa}}$, and $p_i$), padded/masked to kmax, and aggregated by a permutation-invariant Set Transformer (SAB stack + PMA) into a pair logit $z_{\text{pair}}$ and prediction $\hat{Y}$. Bottom: Three-stage training: (1) train the expensive CTS encoder $e_\theta$ with window-level supervision; (2) distill the cheap encoder $\tilde{e}_{\tilde{\theta}}$ from a teacher; (3) train the pair-level aggregator $f_\phi$ using the budgeted forward, with an encoder-freeze warmup followed by joint fine-tuning.
• Figure 2: Performance vs. budget $K$ (evidence for Theorem \ref{['thm:approx']}). We vary the expensive-token budget $K\in\{8,16,32,64,128,256,512\}$ on the miRAW half-split test partition. truncate@Kmax trains once at $K_{\max}{=}512$ and evaluates smaller budgets by masking all but the first $K$ selected tokens (parameters fixed). retrain@K retrains a budget-matched model at each $K$. Points show mean$\pm$std over $R$ runs. Left: PR-AUC. Right: $\mathrm{F1}@0.5$.
• Figure 3: Online inference cost vs. budget $K$. End-to-end latency, throughput, and peak VRAM for BR-MIL_online and Naive_online across $K$, with TargetNet_like_online as a budget-independent reference. Bottom: representative profiled components.
• Figure 4: Stage-wise latency breakdown at $K{=}64$ (log scale). Across pipelines, the shared CPU gather stage dominates wall-time; BR-MIL selection/aggregation contribute a small overhead. See Appendix \ref{['app:runtime_extra']} for stage definitions and profiling details.
• Figure 5: Robustness to visible candidate pool size $n$ at fixed budget $K^\star{=}64$ (evidence for Theorem \ref{['thm:gen']}). We cap the selector's visible pool to the top-$n$ candidates ranked by the cheap logit $\tilde{z}_i$, with $n\in\{64,128,256,512,1024,2048\}$, while keeping the expensive-token budget and Set Transformer input length fixed at $K^\star$. We report test (a) PR-AUC and (b) $\mathrm{F1}@0.5$. Performance drops when $n\approx K^\star$ but saturates once $n$ is a few multiples of $K^\star$ (e.g., $n\ge256$), suggesting that varying $n$ mainly affects selection quality rather than the capacity of the expensive relational component.

Theorems & Definitions (15)

• Theorem 5.1: Approximation error vs. budget $K$
• Theorem 5.2: Generalization governed by budget $K$
• Definition 1.1: Budgeted Relational MIL (BR-MIL)
• Lemma 2.4: Gradient-based influence weights imply \ref{['eq:A2_weighted_lip']}
• proof
• Lemma 2.7: Best-$K$ mass equals the sorted top-$K$ sum
• proof
• Theorem 2.8: Formal approximation bound under budgeted masking
• proof
• Theorem 2.14: Generalization governed by budget $K$; formal version