A feed-forward block that knows what changed, what it already carries, and where it is in depth.
A conventional FFN transforms the representation it receives.SAFFN adds explicit state and change signals to that transformation.It reads the current hidden stream, the current mixer output, the previous mixer output, a persistent FFN state, and a learned physical-depth embedding — then decides what to write, what to retain, and how strongly to read the state back into the model stream.
[ WHAT A STANDARD FFN DOES ]
The original Transformer defines its feed-forward sublayer as a position-wise network: the same two-layer transformation is applied separately to each position. Context is already embedded in the incoming representation by the mixer, but the FFN itself receives no separate previous-mixer signal, persistent FFN state, or physical-depth input.
Excellent at expanding and recombining features. Its behavior at a layer is driven by the current input and that layer's learned weights.
SAFFN keeps the nonlinear feature role of an FFN while making the update explicitly conditional on evolving model state.
Attention, ESA or Bolt can already place contextual information insidex. SAFFN's distinction is that state, mixer transition and depth are exposed as separate signals instead of asking the dense FFN to infer all of them implicitly from the current representation.
[ SIGNAL PATH ]
The reference implementation is explicit about where information comes from. The block does not simply concatenate everything: different projections and gates decide how context, transition and prior state affect the next state.
The normalized representation being refined.
The current ESA/Bolt/context update.
A reference for what the mixer produced one physical layer earlier.
Feature state carried through physical depth.
A learned signal telling the block where it sits in the stack.
Measure how much the mixer signal changed.
Build new state content and decide how much enters.
Control how much previous state survives.
Turn the updated state into the FFN contribution.
[ THE RECURRENT UPDATE ]
SAFFN computes the difference between current and previous mixer outputs, summarizes its magnitude, and feeds that signal into retention and read control. The exact behavior is learned — the architecture provides a mechanism, not a fixed rule that every change is “noise” or “important.”
The candidate combines the hidden stream, transition-adjusted mixer context, previous state and a learned depth term, then passes throughtanh.
A learned sigmoid gate sees the hidden stream, transition-adjusted mixer signal, previous state and depth.
new information controlRetention starts from a learned depth-dependent bias and is modulated by the magnitude of mixer change.
memory continuity controlThe updated state is multiplied by a SiLU value path and a read gate before projection back to model width.
output signal control
sₗ = (1 − wₗ) ⊙ (rₗ ⊙ sₗ₋₁) + wₗ ⊙ ĉₗREADOUTuₗ = Wₒ[sₗ ⊙ SiLU(vₗ) ⊙ qₗ]Notation is simplified for presentation. The MLBricks implementation separately projects the hidden stream, mixer/context, previous state and physical-depth embedding; transition coefficients and delta scales are learned parameters.
[ WHY THE SIGNAL CAN BE RICHER ]
SAFFN is designed to expose information that a conventional position-wise FFN would otherwise have to infer indirectly. That can make the optimization problem more expressive when a task benefits from persistent feature state.
The current-vs-previous mixer delta gives the FFN a direct measure of representation movement across physical depth.
CHANGE AWARENESSUseful intermediate features can be retained instead of being reconstructed only from the latest hidden vector at every layer.
CONTINUITYDifferent mechanisms control what enters state, what survives, and what contributes back to the residual stream.
CONTROLLED FLOWA learned depth embedding lets early and late blocks develop different state-update behavior even when the functional template is shared.
DEPTH SPECIALIZATION[ WHY IT FITS SOUP ]
SOUP combines configurable mixers with evolving state, Observer State Memory and learned fusion. A conventional FFN can still be used, but SAFFN is the natural state-aware complement because it consumes the mixer's current context, compares it with the previous layer's context, and carries its own feature state through the physical stack.
Produces the current context/update from the hidden stream.
Turns current context + previous context + persistent state into a controlled nonlinear refinement.
SOUP observes evolving state into memory and fuses current representation, state and retrieved memory.
ffn="saffn"is supported directly.During recurrent generation, the optimized SOUP path can pack the ESA projection and SAFFN input projections, cache depth terms and reuse fixed transition constants. That keeps the state-aware design while reducing one-token launch overhead.
[ QUALITY & LOSS — THE HYPOTHESIS ]
Mechanistically, SAFFN can reduce avoidable re-derivation by carrying feature state, distinguish stable context from changing context, and regulate information flow with learned gates. Those are plausible routes to better loss or generation quality on state-sensitive workloads. But architecture mechanisms must be separated from empirical claims.
A retained state gives later layers access to feature information without requiring the current hidden stream alone to reconstruct it.
The mixer delta tells the block whether the contextual signal is changing strongly or remaining relatively stable.
Write, retain and read gates can learn different information-flow regimes instead of always applying the same dense transformation pattern.
Depth conditioning lets earlier blocks and later blocks learn different state dynamics while preserving one architectural template.
final validation loss
42,838,108 parametersfinal validation loss
42,836,992 parametersSAFFN minus standard FFN
standard FFN was ~0.49% lower loss in this runtokens per model
same dataset / schedule / mixer familyThe controlled 6-layer run slightly favored the conventional FFN. SAFFN remains interesting because it exposes persistent state and transition-aware control that standard FFNs do not. Its quality advantage should be tested where those signals matter — especially state-centric architectures such as SOUP and longer, richer training regimes.
[ RESEARCH CONTEXT ]
These papers do not define SAFFN. They provide context for why memory-like FFNs, gating and explicit conditioning are useful directions to investigate.
Vaswani et al. define the Transformer FFN as two linear transforms with an activation, applied to each position separately and identically.
Attention Is All You Need ↗Analysis of Transformer FFNs found interpretable memory-like behavior and showed that their outputs are progressively refined through the model stack.
Read the paper ↗GLU variants in Transformer feed-forward sublayers showed that gated feature processing can outperform standard ReLU/GELU FFN variants in tested settings.
GLU Variants ↗FiLM demonstrates explicit feature conditioning; Highway Networks show learned gates regulating information flow. SAFFN explores related principles with mixer-transition and recurrent-state conditioning.
FiLM ↗Highway ↗Research context is conceptual, not a claim of derivation or equivalence. State-Aware FFN is an MLBricks architecture with its own recurrent state equations, mixer-delta conditioning, physical-depth embedding and native execution path.
[ IMPLEMENTATION ]
The exact PyTorch equations remain visible to autograd andtorch.compile. In eager no-grad inference, MLBricks can use its native implementation; withbackend="auto", that individual SAFFN element validates numerical parity, benchmarks native vs PyTorch once, freezes the valid winner, and then stops switching.
Autograd sees the original recurrent equations rather than an opaque native call.
Native inference packs three hidden-stream projections and two recurrent-state projections and caches depth terms.
Native must match the PyTorch reference before it can compete on latency; the winning route is frozen per element.
MLBricks direct tests verify that changing state or ESA/context inputs changes the output/state and that native/reference paths remain close.
StateAwareFFN and VirtualStateAwareFFN are checked for finite gradients through hidden and mixer conditioning during training.
SOUP can pack ESA QGV with SAFFN input projections and cache depth projections for recurrent generation.
[ INSTALLATION · MLBRICKS KIT 1.0.0B1 ]
mlbricks.This component ships inside the unifiedmlbricks-kitdistribution. Python imports continue to use themlbricksnamespace.
pip install mlbricks-kit==1.0.0b1[ QUICK START ]
StateAwareFFN is called with the hidden stream, current mixer update, previous mixer update and previous recurrent state. In a full MLBricks model, Bricks or SOUP manages those signals for you.
from mlbricks import StateAwareFFN
ffn = StateAwareFFN(
d_model=512,
state_dim=256,
depth_embedding_dim=64,
layer_index=0,
total_layers=6,
backend="auto",
)
state = ffn.initial_state(x)
update, next_state = ffn(
x,
mixer_update,
previous_mixer,
state,
)[ WITH SOUP ]
SOUP(dim=512, width=1024, depth=2, mixer=["esa", "bolt"], ffn="saffn", backend="auto")[ API · MLBRICKS KIT 1.0.0B1 ]
The direct component keeps a small recurrent FFN state and accepts the current and previous mixer signals explicitly.
StateAwareFFN(
d_model,
state_dim=256,
depth_embedding_dim=64,
layer_index=0,
total_layers=1,
use_native=None,
fused_cuda=True,
backend="auto",
)initial_state(x)Create the starting recurrent FFN state for an input batch.
ffn(x, esa_update, previous_esa, previous_state)Return the FFN update and next recurrent state.
set_backend(...)Request auto, native, or PyTorch execution.
resolved_backend()Inspect the current backend route.
[ STATE-AWARE FFN ]
SAFFN turns feed-forward processing into a state-conditioned operation: current representation, mixer transition, persistent state and depth all participate in the update. It is designed for architectures where information should evolve rather than be rebuilt from scratch at every physical layer.