# PUFFIN: Protein Unit Discovery with Functional Supervision

arXiv: https://arxiv.org/abs/2604.14796 · HTML: https://mlsanity.com/paper/2604.14796/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 8 of its 31 references that are in the map.
- First posted on arXiv 16 Apr 2026.

## Builds on (most cited first, 8 of 8)

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Graph Attention Networks](https://mlsanity.com/paper/1710.10903/) (2017)
- [Fast Graph Representation Learning with PyTorch Geometric](https://mlsanity.com/paper/1903.02428/) (2019)
- [InterPLM: Discovering Interpretable Features in Protein Language Models via Sparse Autoencoders](https://mlsanity.com/paper/2412.12101/) (2024)
- [Evaluating representation learning on the protein structure universe](https://mlsanity.com/paper/2406.13864/) (2024)
- [Endowing Protein Language Models with Structural Knowledge](https://mlsanity.com/paper/2401.14819/) (2024)
- [Linguistic Laws Meet Protein Sequences: A Comparative Analysis of Subword Tokenization Methods](https://mlsanity.com/paper/2411.17669/) (2024)
- [Protein Structure Tokenization via Geometric Byte Pair Encoding](https://mlsanity.com/paper/2511.11758/) (2025)

## Abstract (arXiv)

Proteins carry out biological functions through the coordinated action of groups of residues organized into structural arrangements. These arrangements, which we refer to as protein units, exist at an intermediate scale, being larger than individual residues yet smaller than entire proteins. A deeper understanding of protein function can be achieved by identifying these units and their associations with function. However, existing approaches either focus on residue-level signals, rely on curated annotations, or segment protein structures without incorporating functional information, thereby limiting interpretable analysis of structure-function relationships. We introduce PUFFIN, a data-driven framework for discovering protein units by jointly learning structural partitioning and functional supervision. PUFFIN represents proteins as residue-level structure graphs and applies a graph neural network with a structure-aware pooling mechanism that partitions each protein into multi-residue units, with functional supervision that shapes the partition. We show that the learned units are structurally coherent, exhibit organized associations with molecular function, and show meaningful correspondence with curated InterPro annotations. Together, these results demonstrate that PUFFIN provides an interpretable framework for analyzing structure-function relationships using learned protein units and their statistical function associations. We made our source code available at https://github.com/boun-tabi-lifelu/puffin.