Beyond Tokens: Probing Higher-Order Epistasis in Learned Protein Representations

Maryam Rahimimovassagh; Ivan Garibay; Niloofar Yousefi

q-bio.QMAug 24, 2026

Beyond Tokens: Probing Higher-Order Epistasis in Learned Protein Representations

Authors: Maryam Rahimimovassagh, Ivan Garibay, Niloofar Yousefi

Organizations: University of Central Florida

Abstract

Protein fitness landscapes contain nonlinear interactions in which mutation effects depend on other residues. We introduce ORBIT, an Order-Resolved Benchmarking of Interaction Transformations framework that separates interaction presence, representation accessibility, and functional recovery. ORBIT first validates Walsh-based diagnostics on synthetic landscapes with known interaction order, then analyzes the experimentally measured GB1 fitness landscape under the FLIP 2-vs-rest setting. We compare ridge regression, a standard MLP, independent tokens, nonlinear independent tokens, and Residual Interaction Tokenization (RIT). Across 20 paired training seeds, the primary two-hidden-layer comparison found no significant architecture differences in FLIP test R^2, third- or fourth-order functional recovery, or final-layer third- or fourth-order accessibility. However, RIT significantly increased pairwise accessibility at the token stage relative to both independent-token controls (Delta A_tok,2 = 0.2468, d_z = 1.67, Holm-adjusted p = 1.14 x 10^-5), without a detectable downstream higher-order advantage. A pre-specified depth/capacity analysis showed that deeper MLPs improved FLIP prediction, third-order functional recovery, and final-layer third-order accessibility; fourth-order accessibility also improved relative to the shallow MLP but remained below zero in absolute held-out R^2. ORBIT therefore reveals representation-level changes hidden by conventional prediction metrics and distinguishes early interaction-aware encoding from higher-order structure constructed by downstream nonlinear capacity.

Figures & tables

Figure 1: ORBIT workflow. ORBIT validates order-resolved diagnostics on ORBIT-Synth and then applies them to GB1 to compare functional Walsh recovery, representation energy, and linear accessibility across five models. The AA-identity-strict probe protocol uses disjoint non-WT amino-acid identities across probe partitions, and a separate pre-specified depth/capacity analysis extends the neural backbones.

Model	RMSE $\downarrow$	$R^{2}$ $\uparrow$	Spearman $\rho$ $\uparrow$
Ridge	1.1619 ± 0.0000	0.1084 ± 0.0000	0.5617 ± 0.0000
MLP	1.1377 ± 0.0986	0.1392 ± 0.1531	0.6667 ± 0.0208
Independent tokens	1.0601 ± 0.0604	0.2556 ± 0.0862	0.6054 ± 0.0365
Nonlinear tokens	1.0898 ± 0.0876	0.2109 ± 0.1298	0.7066 ± 0.0210
RIT	1.0746 ± 0.0737	0.2340 ± 0.1085	0.6037 ± 0.0344

Table 1: Prediction performance on the official FLIP two_vs_rest test set (Hamming distance 3–4). Neural-model values are mean

\pm

standard deviation across 20 seeds (42–61); Ridge is deterministic under the frozen configuration.

Figure 2: Order-resolved functional Walsh recovery on the AA-identity-strict probe test. Pooled functional recovery

R_{k}^{2}

is shown for interaction orders

k=1,\ldots,4

. Neural-model points indicate means across 20 stochastic realizations (seeds 42–61), with error bars showing

\pm 1

standard deviation. Ridge is deterministic. Recovery decreases as interaction order increases, with substantially greater uncertainty at third and fourth order.

Panel of Figure 3 — Figure 3: Order-specific structure across learned GB1 representations. (A) Normalized representation Walsh energy $E_{\ell,k}$ quantifies the fraction of nonconstant representation variation associated with interaction order $k$ . (B) Accessibility $A_{\ell,k}$ measures held-out linear-probe recovery of experimentally measured order- $k$ Walsh coefficients from each representation stage. Energy and accessibility capture distinct properties: the presence of order-specific variation does not imply that experimentally measured epistatic structure is linearly accessible.

Panel of Figure 4 — Figure 4: Depth/capacity sensitivity of GB1 prediction and third-order functional recovery. Neural backbones with two, three, and four hidden layers are evaluated using the same 20 paired training seeds. (A) Official FLIP test $R^{2}$ . (B) Strict-probe third-order functional recovery $R_{3}^{2}$ . The standard MLP shows the clearest depth-associated gains in both prediction and third-order recovery. Because deeper networks also contain additional parameters, these results are interpreted as depth/capacity sensitivity rather than a parameter-matched causal effect of depth.

Endpoint	Comparison	$\Delta$	95% CI	$d_{z}$	$p_{\mathrm{Holm}}$
FLIP test $R^{2}$	RIT $-$ Independent	-0.022	[-0.066, 0.019]	-0.222	0.687
FLIP test $R^{2}$	RIT $-$ Nonlinear	0.023	[-0.049, 0.097]	0.136	0.687
Functional $R_{3}^{2}$	RIT $-$ Independent	-0.007	[-0.046, 0.031]	-0.082	1.000
Functional $R_{3}^{2}$	RIT $-$ Nonlinear	-0.021	[-0.079, 0.038]	-0.156	1.000
Functional $R_{4}^{2}$	RIT $-$ Independent	0.036	[-0.007, 0.077]	0.360	0.741
Functional $R_{4}^{2}$	RIT $-$ Nonlinear	0.042	[-0.019, 0.101]	0.301	0.983

Table 2: Primary two-hidden-layer architecture analysis before the pre-specified depth/capacity extension. Paired statistical comparisons are based on 20 stochastic training seeds. Differences are defined as Model A

-

Model B. Confidence intervals are paired-bootstrap 95% intervals, and

d_{z}

denotes the standardized paired effect size.

p_{\mathrm{Holm}}

values in this table correspond to the original primary-analysis multiplicity families. The subsequently frozen depth/capacity extension and the final 57-test multiplicity analysis are reported in Appendix D .

Appendix figures & tables8 assets

Supplementary material from the paper’s appendix.

Appendix

Model	$R_{1}^{2}$	$R_{2}^{2}$	$R_{3}^{2}$	$R_{4}^{2}$
Ridge	0.4126 ± 0.0000	-0.0012 ± 0.0000	-0.0147 ± 0.0000	-0.0218 ± 0.0000
MLP	0.5397 ± 0.0764	0.4969 ± 0.0657	0.1230 ± 0.0858	-0.0423 ± 0.0787
Independent tokens	0.5498 ± 0.1187	0.4903 ± 0.1406	0.2117 ± 0.0864	0.0190 ± 0.0965
Nonlinear tokens	0.6356 ± 0.1011	0.6769 ± 0.0972	0.2257 ± 0.1218	0.0129 ± 0.1091
RIT	0.5516 ± 0.1160	0.4951 ± 0.1287	0.2043 ± 0.0955	0.0545 ± 0.1037

Appendix

Table 3: Order-resolved functional Walsh recovery on the 613 AA-identity-strict probe-test cubes. Values are pooled

R^{2}

across coefficients of each interaction order. Neural-model values are mean

\pm

standard deviation across 20 seeds (42–61); Ridge is deterministic under the frozen configuration.

Model	Layer	$A_{\ell,1}$	$A_{\ell,2}$	$A_{\ell,3}$	$A_{\ell,4}$
MLP	$H_{1}$	0.7016 ± 0.0348	0.4721 ± 0.0675	0.0413 ± 0.0808	-0.3206 ± 0.1733
MLP	$H_{2}$	0.7481 ± 0.0469	0.5534 ± 0.0633	0.1774 ± 0.0715	-0.2187 ± 0.2091
Independent tokens	$Z_{\mathrm{token}}$	0.3908 ± 0.0798	-0.0539 ± 0.0000	-0.0355 ± 0.0000	-0.1758 ± 0.0000
Independent tokens	$H_{1}$	0.6091 ± 0.0516	0.4336 ± 0.0974	0.0976 ± 0.1133	-0.3316 ± 0.1537
Independent tokens	$H_{2}$	0.7491 ± 0.0715	0.5791 ± 0.1244	0.2245 ± 0.1058	-0.1726 ± 0.1777
Nonlinear tokens	$Z_{\mathrm{token}}$	0.5127 ± 0.0483	-0.0539 ± 0.0000	-0.0355 ± 0.0000	-0.1758 ± 0.0000

Appendix

Table 4: Linear accessibility of experimentally derived order-specific Walsh coefficients from learned representations. Probe alpha is selected on probe validation, the probe is refit on probe train plus validation, and evaluated once on the 613 AA-identity-strict probe-test cubes. Values are mean

\pm

standard deviation across 20 seeds.

Model	Layer	$E_{\ell,1}$	$E_{\ell,2}$	$E_{\ell,3}$	$E_{\ell,4}$
MLP	$H_{1}$	0.937370 $\pm$ 0.006139	0.057134 $\pm$ 0.005665	0.004840 $\pm$ 0.000519	0.000656 $\pm$ 0.000097
MLP	$H_{2}$	0.878389 $\pm$ 0.034318	0.111567 $\pm$ 0.029501	0.009065 $\pm$ 0.004603	0.000979 $\pm$ 0.000480
Independent tokens	$Z_{\mathrm{token}}$	1.000000 $\pm$ 0.000000	0.000000 $\pm$ 0.000000	0.000000 $\pm$ 0.000000	0.000000 $\pm$ 0.000000
Independent tokens	$H_{1}$	0.922745 $\pm$ 0.008103	0.070863 $\pm$ 0.007632	0.005628 $\pm$ 0.000612	0.000764 $\pm$ 0.000126
Independent tokens	$H_{2}$	0.856304 $\pm$ 0.038567	0.128730 $\pm$ 0.033405	0.013499 $\pm$ 0.004856	0.001467 $\pm$ 0.000551
Nonlinear tokens	$Z_{\mathrm{token}}$	1.000000 $\pm$ 0.000000	0.000000 $\pm$ 0.000000	0.000000 $\pm$ 0.000000	0.000000 $\pm$ 0.000000

Appendix

Table 5: Normalized representation Walsh energy on the 613 strict probe-test cubes. The order-0 component is excluded from the normalization denominator.

E_{k}

denotes total energy at order

k

, not average coefficient magnitude. Values are mean

\pm

standard deviation across 20 seeds.

Model	Backbone depth	Official FLIP test $R^{2}$
MLP	2	0.1392 $\pm$ 0.1531
MLP	3	0.2729 $\pm$ 0.0729
MLP	4	0.2953 $\pm$ 0.0435
Independent tokens	2	0.2556 $\pm$ 0.0862
Independent tokens	3	0.3110 $\pm$ 0.0520
Independent tokens	4	0.3249 $\pm$ 0.0429

Appendix

Table 6: Prediction performance in the ORBIT-GB1 deeper-backbone ablation. Values are mean

\pm

sample standard deviation across 20 paired training seeds (42–61). Depths 3 and 4 add 32-wide hidden layers, so the comparison measures depth/capacity sensitivity rather than a parameter-matched causal effect of depth.

Model	Backbone depth	$R_{1}^{2}$	$R_{2}^{2}$	$R_{3}^{2}$	$R_{4}^{2}$
MLP	2	0.5397 $\pm$ 0.0764	0.4969 $\pm$ 0.0657	0.1230 $\pm$ 0.0858	-0.0423 $\pm$ 0.0787
MLP	3	0.7003 $\pm$ 0.0623	0.6243 $\pm$ 0.0506	0.2963 $\pm$ 0.0727	-0.0263 $\pm$ 0.0758
MLP	4	0.7081 $\pm$ 0.0492	0.6503 $\pm$ 0.0443	0.3399 $\pm$ 0.0487	-0.0042 $\pm$ 0.0872
Independent tokens	2	0.5498 $\pm$ 0.1187	0.4903 $\pm$ 0.1406	0.2117 $\pm$ 0.0864	0.0190 $\pm$ 0.0965
Independent tokens	3	0.6251 $\pm$ 0.1331	0.5445 $\pm$ 0.1128	0.2493 $\pm$ 0.1074	-0.0023 $\pm$ 0.1561
Independent tokens	4	0.6589 $\pm$ 0.1099	0.5757 $\pm$ 0.1097	0.2618 $\pm$ 0.1426	0.0064 $\pm$ 0.1675

Appendix

Table 7: Order-resolved functional Walsh recovery on the 613 AA-identity-strict probe-test cubes in the deeper-backbone ablation. Each value is mean

\pm

sample standard deviation across 20 paired training seeds (42–61).

Model	Depth	Stage	$A_{\ell,1}$	$A_{\ell,2}$	$A_{\ell,3}$	$A_{\ell,4}$
MLP	2	$H_{1}$	0.7016 $\pm$ 0.0348	0.4721 $\pm$ 0.0675	0.0413 $\pm$ 0.0808	-0.3206 $\pm$ 0.1733
MLP	2	$H_{2}$	0.7481 $\pm$ 0.0469	0.5534 $\pm$ 0.0633	0.1774 $\pm$ 0.0715	-0.2187 $\pm$ 0.2091
MLP	3	$H_{1}$	0.6383 $\pm$ 0.0531	0.4138 $\pm$ 0.1491	0.0047 $\pm$ 0.1240	-0.4583 $\pm$ 0.5004
MLP	3	$H_{2}$	0.7632 $\pm$ 0.0453	0.5663 $\pm$ 0.0737	0.1977 $\pm$ 0.0762	-0.1943 $\pm$ 0.1115
MLP	3	$H_{3}$	0.7998 $\pm$ 0.0337	0.6669 $\pm$ 0.0408	0.2831 $\pm$ 0.0604	-0.0535 $\pm$ 0.0639
MLP	4	$H_{1}$	0.6419 $\pm$ 0.0559	0.4104 $\pm$ 0.0622	0.0193 $\pm$ 0.0636	-0.2850 $\pm$ 0.1378

Appendix

Table 8: Layer- and order-resolved linear accessibility in the deeper-backbone ablation. The same AA-identity-strict probe protocol is used at every stage. Values are mean

\pm

sample standard deviation across 20 paired training seeds (42–61).

Model	Depth	Stage	$E_{\ell,1}$	$E_{\ell,2}$	$E_{\ell,3}$	$E_{\ell,4}$
MLP	2	$H_{1}$	0.937370 $\pm$ 0.006139	0.057134 $\pm$ 0.005665	0.004840 $\pm$ 0.000519	0.000656 $\pm$ 0.000097
MLP	2	$H_{2}$	0.878389 $\pm$ 0.034318	0.111567 $\pm$ 0.029501	0.009065 $\pm$ 0.004603	0.000979 $\pm$ 0.000480
MLP	3	$H_{1}$	0.944213 $\pm$ 0.004841	0.050727 $\pm$ 0.004442	0.004466 $\pm$ 0.000424	0.000595 $\pm$ 0.000095
MLP	3	$H_{2}$	0.918602 $\pm$ 0.019442	0.076230 $\pm$ 0.018444	0.004638 $\pm$ 0.001154	0.000530 $\pm$ 0.000154
MLP	3	$H_{3}$	0.806304 $\pm$ 0.025654	0.169747 $\pm$ 0.021498	0.021968 $\pm$ 0.004587	0.001981 $\pm$ 0.000457
MLP	4	$H_{1}$	0.945460 $\pm$ 0.004794	0.049532 $\pm$ 0.004442	0.004416 $\pm$ 0.000406	0.000593 $\pm$ 0.000079

Appendix

Table 9: Normalized representation Walsh energy in the deeper-backbone ablation. The order-0 component is excluded from the normalization denominator;

E_{\ell,k}

is a descriptive representation-energy metric, not a biochemical interaction-strength estimate. Values are mean

\pm

sample standard deviation across 20 paired training seeds (42–61).

Endpoint	Comparison	$\Delta$	95% CI	$d_{z}$	$p_{\mathrm{raw}}$	$p_{\mathrm{Holm}}$
Prediction (11 tests)
FLIP test $R^{2}$	Nonlinear $-$ Independent	-0.045	[-0.106, 0.014]	-0.320	0.172	1.000
FLIP test $R^{2}$	RIT $-$ Independent	-0.022	[-0.066, 0.019]	-0.222	0.344	1.000
FLIP test $R^{2}$	RIT $-$ Nonlinear	0.023	[-0.049, 0.097]	0.136	0.548	1.000
FLIP test $R^{2}$	MLP: $d_{3}-d_{2}$	0.134	[0.071, 0.201]	0.876	$5.68\times 10^{-4}$	0.006
FLIP test $R^{2}$	MLP: $d_{4}-d_{2}$	0.156	[0.097, 0.221]	1.076	$6.1\times 10^{-5}$	$6.71\times 10^{-4}$

Appendix

Table 10: Complete pre-specified confirmatory statistical results. Differences are defined as Model A

-

Model B for architecture contrasts and depth A

-

depth B for depth contrasts. Confidence intervals are paired-bootstrap 95% intervals;

d_{z}

is the standardized paired effect size.

p_{\mathrm{raw}}

is the two-sided exact paired sign-flip permutation

p

-value and

p_{\mathrm{Holm}}

is adjusted within the corresponding pre-specified family. Bold adjusted

p

-values satisfy

p_{\mathrm{Holm}}\leq 0.05

Protein Representation Learning

Protein representation learning builds computational features that capture sequence, structure, function, and molecular relationships.

See all 69 papers

ORBIT-FMIB: Tracking Order-Resolved Epistatic Information Through ESM-2
Sep 30, 2026Maryam Rahimimovassagh, Ivan Garibay, Niloofar YousefiProtein Language Models Protein Representation Learning
University of Central Florida
Interpreting Protein Language Model Embeddings via Orthogonal Projection for Protein Fitness Prediction
Aug 26, 2026Paulo Yanez Sarmiento, Pia Francesca Rissom, Manuel Pfeuffer +6Protein Language Models Protein Representation Learning
Hasso Plattner Institute, University of Potsdam · Bioinformatics and Machine Learning, Broad Institute of MIT and Harvard · School of Business and Economics, Humboldt University of Berlin +1
Structure-Aware Masking for Protein Representation Learning
May 15, 2026Thomas Walton, Ayan Goel, Amirali AghazadehProtein Fitness Prediction Masked Language Modeling
Georgia Institute of Technology

Sep 30, 2026J/K move · Enter open · S save

ORBIT-FMIB: Tracking Order-Resolved Epistatic Information Through ESM-2

Maryam Rahimimovassagh, Ivan Garibay, Niloofar Yousefi

University of Central Florida

Protein Language Models Protein Representation Learning

Open paper AskPDF arXiv

Protein foundation models support mutation-effect and structural prediction, but predictive performance alone does not reveal which forms of biological interaction information remain accessible through model depth. We ask whether ESM-2 retains higher-order epistatic information as strongly as first- and second-order information across its representation hierarchy, introducing ORBIT-FMIB, a diagnostic framework combining Walsh-based interaction decomposition with subset-conditioned neural dependence estimation. The method is validated on synthetic landscapes with known interaction structure before being applied to the dense four-site GB1 fitness landscape using frozen ESM-2 representations. An initial production run suggested ESM-2 retains higher-order epistatic information less well than lower-order information (

Δ_{\mathrm{HO-LO}}=-0.107

). An independent replication of the complete measurement grid, under matched GPU hardware and identical critic seeds, substantially reduced this contrast (

Δ_{\mathrm{HO-LO}}=-0.017

), and its sign was unstable across otherwise-defensible evaluation-pairing choices applied to the same trained critics (

-0.011

+0.015

). We therefore do not currently have robust evidence that ESM-2 selectively loses higher-order epistatic information, nor that retention is equal across orders; the directional question remains open. The measurement protocol itself, including its documented removal of a positional-subset shortcut in pooled critics, remains validated and is unaffected by this finding. ORBIT-FMIB is offered as a diagnostic framework for probing interaction structure in protein foundation models; this study's own replication result illustrates why such probing requires adequately-powered reproducibility checks before its output is treated as a biological finding.

Abstract

Protein Representation Learning

Explore similar work

ORBIT-FMIB: Tracking Order-Resolved Epistatic Information Through ESM-2

Interpreting Protein Language Model Embeddings via Orthogonal Projection for Protein Fitness Prediction

Structure-Aware Masking for Protein Representation Learning