FLaG: Frequency-Domain Latent-attention Gated Pooling for Token Aggregation

Kewei Li; Rongying Zhang; Xueli Wang; Xiwen Gong; Zhongjian Wang; Qiuchen Zhao; Lan Huang; Ruochi Zhang; Fengfeng Zhou

Source-linked AI summary

FLaG: Frequency-Domain Latent-attention Gated Pooling for Token Aggregation

Kewei Li, Rongying Zhang, Xueli Wang, Xiwen Gong, Zhongjian Wang, Qiuchen Zhao, Lan Huang, Ruochi Zhang, Fengfeng Zhou

arXiv:2609.00831v1Sep 1, 2026cs.AIq-bio.BM

TL;DR

Token aggregation usually compresses encoder outputs in the original token space, motivating a more general aggregation interface. FLaG inserts Fourier-domain re-expression, latent-query summarization, and sample-conditioned gating before reconstruction and pooling, and evaluates the same module across protein, visual, and textual tasks. Across these settings, FLaG improves several benchmark outcomes but is not universally optimal, with benefits depending on the backbone and downstream task.

Problem
Most pooling methods operate in the original token space, limiting evidence for frequency-domain re-expression as a general aggregation stage before sample-level reduction.
Method
FLaG is a plug-in module that spectrally re-expresses encoder outputs, summarizes frequency tokens with latent queries, applies sample-conditioned channel gating, reconstructs tokens, and pools them.
Results
FLaG achieved the best aggregate AMP Spearman correlation, RMSE, and Recall@50 across four backbone-species settings, the highest CIFAR-10 top-1 accuracy, and best results on some language tasks, but not STSBenchmark.
Takeaways & Limitations
Spectral re-expression provides a transferable aggregation coordinate system, but FLaG’s benefits reflect interactions among spectral content, encoder depth, positional responses, and learned allocation.
Takeaways & Limitations
FLaG is not universally optimal; limitations include absent spectral localization, AMP-only mechanistic evaluation with ESM2-8M, and unresolved interpretations of non-DC components.

Abstract

from arXiv · show

Token aggregation converts token-level representations into fixed-dimensional sample representations, but most pooling methods operate only in the original token space. We introduce Frequency-Domain Latent-attention Gated Pooling (FLaG), a plug-in aggregation module that re-expresses encoder outputs in the Fourier domain before final pooling. FLaG represents the nonredundant rFFT spectrum through concatenated real and imaginary components, summarizes spectral tokens with learnable latent queries, derives a sample-conditioned channel gate, and reconstructs modulated token representations for downstream aggregation. We evaluate the same architecture across ESM2-based antimicrobial peptide (AMP) activity prediction, ResNet18 image classification on CIFAR-10 and CIFAR-100, and three RoBERTa-based language tasks. FLaG achieves the best macro-averaged Spearman correlation coefficient, RMSE, and Recall@50 across four AMP backbone-species settings and the highest top-1 accuracy on CIFAR 10. It also achieves the best mean results on five of seven language metrics, although mean pooling remains strongest on STSBenchmark. AMP-side mechanistic analyses reveal low-frequency prediction sensitivity across most encoder layers, with increased relative high-frequency sensitivity in the final layer, and pronounced peptide-specific positional responses. The residual gate broadly amplifies spectral channels while preserving the low-frequency-dominated energy profile, whereas latent cross-attention exhibits sample- and species-specific spectral allocation. Overall, FLaG provides a transferable frequency-domain aggregation bias across protein, visual, and textual representations, with benefits that depend on the backbone and downstream task. Supplementary materials, source code, and data are available at https://www.healthinformaticslab.org/supp/ and https://github.com/Kewei2023/AMPCliff/tree/FLaG.

1 Introduction

FLaG addresses the limitation that token aggregation typically operates in the original representation basis by inserting frequency-domain re-expression before pooling. It provides a plug-in mechanism evaluated across protein, visual, and textual tasks, with task- and backbone-dependent benefits.

Most pooling methods reduce encoder outputs directly in the original token coordinates without first re-expressing token-axis variation.
Pooling design can materially affect downstream performance across different representation architectures.
FLaG re-expresses encoder outputs spectrally, summarizes frequency tokens with learnable latent queries, gates channels conditionally, and reconstructs tokens for standard pooling.
FLaG operates after the encoder without modifying the backbone or requiring domain-specific external anchors.
The same core architecture is evaluated on ESM2-based AMP prediction, ResNet18 image classification, and RoBERTa-based language tasks.
The results characterize where frequency-domain aggregation is beneficial and show that its effect depends on the backbone and downstream objective.
AMP serves as a mechanistic case study examining positional sensitivity, layer-wise frequency dependence, spectral modulation, latent-query allocation, and physicochemical information.

2 Related Work

Related pooling methods differ in their reduction rules, information sources, and limitations, while frequency-domain approaches have mostly targeted modeling or filtering rather than general post-encoder pooling. FLaG positions spectral re-expression as an intermediate aggregation stage that avoids dependence on protein-specific auxiliary information.

Token Pooling Methods: Mean pooling uses uniform weighting, whereas attention pooling learns token importance but adds parameters and may be sensitive to irrelevant instances or dataset-specific responses.
Token Pooling Methods: Max pooling can preserve salient local responses but may be sensitive to task-irrelevant outlier activations, while last-token selection can be unsuitable for bidirectional encoders.
Token Pooling Methods: Multi-Layer Trainable Pooling combines hidden-layer representations, but adds a trainable layer-fusion module and has task-dependent benefits.
Token Pooling Methods: Latent-attention provides a flexible bottleneck through multiple learnable queries, but larger inputs may require a correspondingly larger latent set.
Token Pooling Methods: Latent-attention remains a token-domain strategy and does not explicitly decompose global trends, local motif variations, or frequency-domain perturbations.
Frequency-Domain Modeling: Frequency-domain research has addressed token mixing, filtering, representation analysis, and backbone modeling rather than a general pooling interface.
Frequency-Domain Modeling: FLaG instead applies spectral re-expression after the encoder and before final pooling, combining FFT re-expression, latent-query summarization, and channel-wise spectral gating.
Protein Pooling: Protein-specific pooling methods may rely on homologs, graph-derived signals, or biological priors, whereas FLaG uses an internal frequency-domain representation without these external anchors.

3 Methods

FLaG inserts a frequency-domain representation between encoder outputs and final pooling, allowing token variation to be summarized spectrally before producing a fixed-dimensional sample vector. It combines rFFT, latent attention, sample-conditioned channel gating, inverse reconstruction, and masked pooling.

Problem formulation: FLaG maps encoder token representations to a fixed-length sample vector while preventing padded positions from contributing.The optional binary mask identifies valid positions, and final pooling remains mask-aware.
Overview: FLaG re-expresses token-axis variation in spectral coordinates before final pooling rather than reducing representations directly in the token domain.This intermediate representation provides an alternative structural view before compression.
Frequency transformation: The rFFT retains nonredundant frequency bins, whose coefficients depend on all token positions and are represented by concatenated real and imaginary components.The zero-frequency coefficient represents the sequence-level mean component up to the normalization convention.
Latent-attention in frequency space: Learnable latent queries cross-attend to frequency tokens, compressing the spectrum into a smaller set of learned spectral summary slots without predefined frequency bands.The latent bottleneck uses L queries, typically with L < F, where F is the number of frequency tokens.
Latent-conditioned spectral gating: FLaG averages refined latent summaries and maps them through an MLP to a sample-conditioned channel gate shared across frequency bins.The gate recalibrates concatenated real and imaginary channels rather than selecting individual frequency bins.
Time-domain reconstruction and final pooling: The gated spectrum is recombined, inverse-transformed into real-valued token representations, and aggregated with masked max or mean pooling.Masked max pooling is the default final aggregation in the main FLaG experiments.

4.1 Experimental Setup

The study evaluates the same FLaG aggregation framework across protein, visual, and textual representations using matched backbone-specific protocols and multiple pooling baselines. Experiments use repeated random seeds, fixed core FLaG settings, and task-appropriate validation procedures.

Cross-domain evaluation: FLaG is evaluated on antimicrobial peptide activity prediction, image classification, and language representation tasks.The domains use ESM2, ResNet18, and RoBERTa-base backbones, respectively.
Antimicrobial peptide activity prediction: AMP experiments use ESM2-8M and ESM2-35M on E. coli and S. aureus activity-cliff-aware splits.Activity is represented as -log10(MIC[M]), with larger values indicating stronger antimicrobial activity.
Image classification: Image experiments reshape ResNet18 spatial feature maps into token sequences and evaluate CIFAR-10 and CIFAR-100 with top-1 and top-5 accuracy.The official 10,000-image test set is retained for final evaluation.
Language representation tasks: Language experiments use RoBERTa-base on IMDB, STSBenchmark, and SprintDuplicateQuestions under supervised downstream adaptation.This protocol accounts for trainable pooling modules rather than treating the comparison as a standard frozen MTEB leaderboard submission.
Baselines: FLaG is compared with six pooling baselines spanning fixed, attention-based, distributional, and multilayer aggregation strategies.All methods receive the same task-specific backbone representations and follow matched training protocols.
Training and model settings: Main experiments use 10 random seeds, eight latent queries, four attention heads, dropout 0.1, residual gating, and max pooling without domain-specific tuning.The pooling module is inserted immediately before the task-specific prediction head.

4.2 Main Results

Across protein, image, and language benchmarks, FLaG often improves aggregate or task-specific performance, but its advantage varies by dataset, metric, and adaptation setting. Component analyses likewise show that spectral gating is strong while latent-attention contributes unevenly across species.

Antimicrobial Peptide Activity Prediction: 0.577 macro-average Spearman, 0.552 macro-average RMSE, and 19.025 macro-average Recall@50 were FLaG’s best AMP aggregate results across four backbone-species settings.Relative to max pooling, FLaG increased Spearman by 0.003, reduced RMSE by 0.004, and increased Recall@50 by 0.518.
Image Classification: 96.01% top-1 accuracy on CIFAR-10 was FLaG’s highest image result, while 77.20% on CIFAR-100 ranked second.FLaG also achieved the highest aggregate accuracy across four reported metrics, at 91.55%, with the lowest average standard deviation of 0.11 percentage points.
Language Benchmarks: Mean pooling remained strongest on STSBenchmark, reaching 0.852±0.003 Spearman and 0.856±0.002 Pearson correlations.FLaG reached 0.837±0.004 Spearman and 0.832±0.005 Pearson, so its benefit did not extend uniformly to semantic textual similarity.
Language Benchmarks: FLaG ranked first on five of seven language metrics, with modest gains on IMDB and larger gains on SprintDuplicateQuestions.On SprintDuplicateQuestions, FLaG achieved 0.993±0.001 accuracy, 0.647±0.032 Average Precision, and 0.617±0.026 F1.
Component Ablation Study: FLaG’s component contributions were task-dependent: FFT+gate was strongest on E. coli, whereas full FLaG improved macro-averaged regression and correlation through stronger S. aureus performance.FFT+gate nevertheless achieved the highest macro-average Recall@50 at 18.45, compared with 18.30 for FLaG.

4.3 AMP-Side Mechanistic Analysis

AMP-side analyses show that FLaG’s predictive sensitivity varies across peptide positions and frequency bands, with layer- and species-dependent spectral allocation. Its gate preserves low-frequency dominance while latent attention reallocates spectral emphasis across samples.

Single-position hidden-state knockout: FLaG shows broader within-peptide positional-response heterogeneity than mean pooling and MLTP, while remaining comparable to other position-selective methods.This pattern appears for both E. coli and S. aureus, with especially pronounced upper-range responses for S. aureus.
Single-position hidden-state knockout: FLaG ranks first in Spearman correlation among ESM2-8M pooling methods for both E. coli and S. aureus.The figure’s circled ranks connect positional-response distributions with the corresponding test-set rankings.
Interpretive caveat: The knockout analyses measure model-level positional sensitivity rather than biological residue essentiality, and greater heterogeneity is not established as causal for performance.The perturbation changes contextual hidden representations rather than the peptide sequence directly.
Sequence-frequency band knockout: Shallow ESM2-8M layers are most sensitive to the lowest-frequency band, whereas the final layer shows greater relative sensitivity to the highest-frequency band.This indicates a depth-dependent redistribution from slowly varying toward rapidly varying sequence components.
Sequence-frequency band knockout: Representative peptides exhibit distinct layer-band sensitivity profiles, especially at the lowest and highest frequencies, while middle bands vary less with depth.Thus, the aggregate frequency shift does not describe every peptide identically.
Latent-query cross-attention: Latent attention reallocates E. coli emphasis from the DC bin toward the first non-DC bin, while S. aureus allocation is less consistent across peptides.The eight queries generally share sample-level profiles, supporting sample-conditioned readout rather than fixed query-to-frequency assignments.
Gate spectral-effect analysis: The residual gate broadly recalibrates spectral channels while preserving the dominance of the lowest-frequency band.No clear redistribution toward intermediate- or high-frequency components is observed after gating.
Physicochemical-property analysis: Global physicochemical properties are most strongly decodable from low-frequency representations, but their activity associations differ between bacterial species.Net charge is more strongly associated with E. coli activity, whereas mean hydrophobicity is more strongly associated with S. aureus activity; hydrophobic moment and helix-group fraction show no significant species-specific differences.

5 Conclusions

FLaG is a transferable spectral token-aggregation module whose benefits vary across backbones and tasks. Its AMP analyses identify depth- and sample-dependent spectral behavior, while the paper notes unresolved mechanistic and localization limitations.

Contribution: FLaG combines spectral re-expression, latent-query summarization, sample-conditioned channel gating, inverse reconstruction, and final token pooling.It is evaluated for AMP activity prediction, image classification, and language tasks.
Cross-task results: FLaG achieves the best aggregate AMP Spearman correlation, RMSE, and Recall@50 across four backbone-species settings, highest CIFAR-10 top-1 accuracy, and best results on IMDB and SprintDuplicateQuestions.Mean pooling remains strongest on STSBenchmark, and ablations show task-dependent benefits for latent spectral summarization.
Mechanistic findings: AMP analyses reveal heterogeneous positional responses, predominantly low-frequency sensitivity across most layers, increased relative high-frequency sensitivity in the final layer, and broad gate modulation.Latent cross-attention shows more frequency-specific and species-dependent allocation.
Limitations: FLaG is not universally optimal, and its current limitations include no spectral localization, AMP-only mechanistic evaluation with ESM2-8M, and unresolved interpretations of non-DC components.Its benefits reflect interactions among spectral content, encoder depth, positional responses, and learned allocation.