Cross+FFN: Integrating Attention and FFN

• The paper shows that standard FFN layers can be interpreted as specialized cross-attention mechanisms by algebraically collapsing retrieval operations into canonical FFN structures.
• It presents modular integration patterns—such as explicit cross-attention modules, pipelined designs, and dynamic compute allocation—that improve efficiency and model capacity.
• Advanced variants incorporating persistent activity and semantic clustering demonstrate practical gains in long-context reasoning and multimodal performance.

Cross+FFN design refers to architectural strategies that directly combine or draw equivalence between cross-attention mechanisms and the feed-forward network (FFN) modules within Transformer blocks, with the goals of improving interpretability, efficiency, knowledge integration, and long-context reasoning. Recent research demonstrates that standard FFN layers can be interpreted as specialized cases of cross-attention retrieval from a static or learned knowledge base, enabling modular systems, explicit knowledge reasoning, and principled avenues for scalable model improvements (Guo et al., 1 Jan 2025). Additionally, emerging Cross+FFN variants incorporate memory banks, clustering, and dynamic compute allocation to further enhance capacity and efficiency in both unimodal and multimodal settings (Li et al., 31 Jan 2025, Li et al., 18 Jun 2025).

1. Generalized Cross-Attention and FFN Correspondence

The foundational result is the rigorous derivation that a standard Transformer FFN is a closure (special case) of a generalized cross-attention mechanism. In this framework, post–self-attention activations $H_\ell \in \mathbb{R}^{n \times d}$ at layer $\ell$ attend to a knowledge base $E \in \mathbb{R}^{|E|\times d_E}$ via learnable projections $W_Q^\ell$, $W_K^\ell$, $W_V^\ell$ and a thresholding function $B_1^\ell(E)$. The knowledge retrieval per layer computes:

• $Q_\ell = H_\ell W_Q^\ell$
• $K_\ell = E W_K^\ell$
• $V_\ell = E W_V^\ell$
• $\ell$0 (sparse, thresholded selection)
• $\ell$1
• $\ell$2

By folding a static $\ell$3 into fixed projection matrices, the above collapses algebraically to the canonical FFN structure:

$\ell$4

where $\ell$5, $\ell$6, $\ell$7 are derived from joint-training or static knowledge base parameters, confirming that FFN can be interpreted as knowledge retrieval via cross-attention (Guo et al., 1 Jan 2025).

2. Modular Integration Patterns for Cross+FFN

Multiple integration patterns for Cross+FFN emerge, depending on the Transformer variant and downstream requirements:

• Explicit cross-attention module: Insert layer-wise (possibly multi-head) cross-attention to a shared or external knowledge base after self-attention, with residuals and normalization, optionally followed by or replacing the FFN.
• Combined design (pipelined): Cross-attention output is refined by a (smaller) FFN, yielding the sequence: Cross-Attention $\ell$8 Residual + LayerNorm $\ell$9 FFN $E \in \mathbb{R}^{|E|\times d_E}$0 Residual + LayerNorm.
• FFN as closure-only: Entirely replace the FFN by a folded cross-attention module for maximum interpretability and explicit knowledge retrieval.
• Dynamic/selective compute: In MLLMs or efficiency-focused models, selectively reduce compute by gating or reducing FFN and/or cross-attention only for certain tokens or layers, based on layer ranking or token importance heuristics (Li et al., 31 Jan 2025).

3. Extensions: Persistent Activity and Semantic Clustering

Advanced Cross+FFN variants incorporate additional mechanisms to address context length limitations and functional modularity:

• Persistent Activity (PA): Augments FFN layers with an Activation Memory Bank (AMB), storing and retrieving prior activations via cosine similarity, dynamically fusing retrieved “memories” with current inputs. Depending on retrieval strength, the update mixes past outputs with the current activation, and memory slots are updated using LRU or similarity-based rules. Empirically, this improves long-context retention without altering overall structure (Li et al., 18 Jun 2025).
• Cortical Expert (CE) Clustering: FFN weights $E \in \mathbb{R}^{|E|\times d_E}$1, $E \in \mathbb{R}^{|E|\times d_E}$2 are partitioned into $E \in \mathbb{R}^{|E|\times d_E}$3 balanced clusters (“experts”) using constrained K-means. Weights are permuted so that semantically similar rows/columns are contiguous, without conventional gating or MoE routing. This semantic modularity mitigates fragmentation and improves performance especially in long-context QA and summarization tasks (Li et al., 18 Jun 2025).

4. Practical Design Principles and Hyperparameter Choices

Empirical findings and architectural recommendations for Cross+FFN designs include:

Parameter	Typical Choices	Rationale/Notes
Hidden dim ($E \in \mathbb{R}^{|E|\times d_E}$4)	$E \in \mathbb{R}^{|E|\times d_E}$5–$E \in \mathbb{R}^{|E|\times d_E}$6	Standard; match baseline model width
Per-head size	$E \in \mathbb{R}^{|E|\times d_E}$7 or $E \in \mathbb{R}^{|E|\times d_E}$8	Set heads $E \in \mathbb{R}^{|E|\times d_E}$9 $W_Q^\ell$0 for parallelism
KB size ($W_Q^\ell$1)	$W_Q^\ell$2 (joint), $W_Q^\ell$3–$W_Q^\ell$4 (external)	Joint: matches $W_Q^\ell$5; external: retrieved subset $W_Q^\ell$6 (Guo et al., 1 Jan 2025)
Activation	ReLU (sparse), Sparsemax, entmax	Sparse selection for interpretability
Threshold MLP	$W_Q^\ell$7–$W_Q^\ell$8 layer, width $W_Q^\ell$9	Initial biases tuned s.t. $W_K^\ell$0–$W_K^\ell$1 active
Weight init	Xavier/Glorot uniform, $W_K^\ell$2	Standard initialization protocol
FFN refinement	GeLU or Swish	When retaining a refinement FFN block

Interpretability and sparsity are encouraged by $W_K^\ell$3 penalties on $W_K^\ell$4 and by sharing $W_K^\ell$5 across layers. Empirical work supports high selectivity and visible “knowledge injection” per token (Guo et al., 1 Jan 2025).

5. Cross+FFN for Multimodal and Efficient Architectures

Cross+FFN analysis reveals substantial, structured redundancy in certain decoder-only MLLM architectures, motivating selective reduction and dynamic computation (Li et al., 31 Jan 2025):

• Probe-Activated Dynamic FFN: For visual tokens, sample activations are used to determine which hidden dimensions are most active. Only the top-$W_K^\ell$6 hidden channels are used, reducing FLOPs by approximately $W_K^\ell$7.
• Layer Ranking Algorithm: Greedy or hybrid layer ranking identifies which FFN or attention layers are most redundant, enabling compute reductions in late layers with minimal performance degradation.
• Empirical finding: Reducing compute for $W_K^\ell$8 of layers preserves accuracy within $W_K^\ell$9, and combining Cross+FFN reduction with token-pruning yields $W_V^\ell$0 runtime reduction at little or no cost (Li et al., 31 Jan 2025).

6. Training and Evaluation Protocols

Robust Cross+FFN deployment follows a strict protocol:

• Joint training: All parameters—including KB embeddings, threshold MLPs, and projection matrices—are trained end-to-end, with regularization on sparsity/selectivity.
• External KB retrieval: For external knowledge bases (e.g., FAISS-indexed), retrieve a subset $W_V^\ell$1 per query; optionally employ differentiable or straight-through retrieval for end-to-end fine-tuning (Guo et al., 1 Jan 2025).
• Regularization and monitoring: Sparsity regularizers, visualizations of $W_V^\ell$2 patterns and KB “hot spots”, and layerwise probes of the “folded” $W_V^\ell$3, $W_V^\ell$4 matrices for interpretability.
• Evaluation metrics: Downstream accuracy, perplexity, retrieval precision/recall (if available), sparsity, selectivity, compute/memory scaling with $W_V^\ell$5 or $W_V^\ell$6.

7. Significance and Impact

Cross+FFN architecture establishes an explicit pathway to interpret, modularize, and scale knowledge and reasoning within Transformer models. The equivalence between FFN and cross-attention enables rich integration with external memory, interpretable knowledge retrieval, and tissue-like modularity that mirrors cortical organization, as evidenced by PA and CE enhancements (Li et al., 18 Jun 2025). In multimodal models, Cross+FFN analysis provides principled reduction schedules for visual token processing, bridging performance-efficiency tradeoffs between decoder-only and cross-attention-based MLLMs (Li et al., 31 Jan 2025). The combined theoretical and empirical paradigm thus opens new research avenues for modular, adaptive, and knowledge-centric neural systems.