More depth by duplicating full blocks.
LAYER 01LAYER 02LAYER 03
Each extra physical block carries another full set of projections and parameters. That can improve capacity, but stored weights and model memory grow with depth.
VIRTUAL DEPTH · STATE-AWARE COMPUTE
Virtual State-Aware
FFN
ActiveReuse the refiner. Change the state. Change the pass.
A StateAwareFFN extension that performs the physical state update once, computes model-width conditioning once, then repeatedly refines the evolving state through ashared virtual-refiner core. Each refinement has its own learned pass identity and gate, so repeated computation is not just the exact same transformation seeing the exact same representation.
[ WHY VIRTUAL DEPTH ]
Physical depth stores another block.
Virtual depth can spend compute instead.
When device storage or weight memory is the constraint, adding another complete FFN block can be expensive. VirtualStateAwareFFN explores a different trade: keep one physical state-aware block, then perform additional lightweight state refinements around a shared refiner.
More state refinement around one physical block.
PHYSICAL SAFFNPASS 01PASS 02
The base SAFFN owns the physical state update. A smaller shared refiner then revisits the evolving state multiple times with pass-specific identity.
THE TRADEParameters ↓ does not mean computation ↓.
Virtual depth trades stored parameter capacity for repeated computation. It can be attractive on memory-constrained hardware when compute is available, but latency and activation memory still need workload-specific measurement.
[ HOW THE CONTROLLER CHANGES EACH PASS ]
The shared core stays.
The representation does not.
The public implementation first runs the StateAwareFFN state update, then computesx_conditionandesa_conditiononce. Those conditions are reused while the current state is iteratively refined.
Hidden stream + current ESA + previous ESA + prior state + physical-depth signal create the next state.
Model-width conditions are projected once and reused across every virtual refinement.
The current state enters the same refiner with a learnedpass embeddingandpass gate.
The updated state becomes the input state of the next virtual pass.
VIRTUAL REFINEMENT
hₖ = SiLU(Wₛ Norm(sₖ) + cₓ + cₑₛₐ + pₖ)sₖ₊₁ = sₖ + σ(gₖ) ⊙ W↓ hₖRefiner weights
State projection, conditioning projections and the down projection are reused across refinement passes.
Evolving state
Pass two does not see the same state as pass one. It operates on the state already altered by the previous pass.
Pass embedding
Every refinement index has a learned embedding that injects a distinct virtual-pass identity.
Pass gate
Every refinement owns a state-width gate vector controlling how strongly that pass writes its update.
The virtual down projection is zero-initialized, so when matching base StateAwareFFN weights are copied, the virtual extension begins functionally as the base block and learns to open the refinement path during training.
[ MEMORY ARITHMETIC ]
Virtual passes are deliberately smaller than full physical duplication.
This is parameter arithmetic from the current MLBricks implementation, not an end-to-end quality or RAM benchmark. Example:d_model=384,state_dim=256,virtual_hidden_dim=128, two virtual refinements.
parameters in the base block
parameters for two virtual passes
over one base StateAwareFFN
fewer extra parameters than duplicating two additional full StateAwareFFN blocks
WHY THIS MATTERS ON EDGE
Reuse weights that are already resident.
On phones, embedded accelerators and other memory-constrained systems, model storage and weight movement can be as limiting as arithmetic. A shared virtual refiner can spend extra cycles on an already-loaded parameter set instead of requiring another full block to be stored and fetched. The exact latency/energy win is device-dependent and must be benchmarked.
[ RESEARCH CONTEXT ]
Recurrent depth is established.
Pass identity is the important twist.
These papers are context for the design space, not evidence that MLBricks has reproduced their results.
Repeat computation through depth.
Universal Transformers made recurrent application of a transformation across depth a first-class sequence-modeling idea, showing that effective computation depth does not have to equal the number of separately parameterized layers.
Read paper ↗Parameter sharing can shrink the stack.
ALBERT used cross-layer parameter sharing to reduce parameter count and memory consumption. Its ablations also show the tension: pure sharing can reduce representational diversity, especially in FFN parameters.
Read paper ↗Shared layers can matter on mobile.
SHARP studies adjacent-layer sharing plus recovery parameters and reports reduced stored MLP parameters and mobile model storage. It is a different method, but it reinforces why weight reuse is interesting for edge inference.
Read paper ↗MLBRICKS DIFFERENCEDo not repeat an indistinguishable layer.
VirtualStateAwareFFN combines a shared refiner with an evolving recurrent state and pass-specific embedding/gating. The research question is whether this can recover useful depth-like refinement without paying for the full parameter footprint of additional physical blocks.
[ WHERE IT FITS ]
Designed for models where memory is expensive but another pass is affordable.
The component is exposed directly and is also selectable inside current ESA, VESA, VisionBolt and Gaussian model paths.
Parameter-efficient depth
Explore deeper state refinement while limiting growth in stored FFN parameters.
State-on-state refinement
Fits naturally next to architectures where mixer state and persistent memory are already first-class signals.
Reuse across visual blocks
VESA and VisualBolt can select the virtual FFNBrick path within the same composable vision runtime.
Compute-for-memory trade
Potentially useful when arithmetic throughput is plentiful but RAM capacity or parameter bandwidth is constrained.
[ INSTALLATION · MLBRICKS KIT 1.0.0B1 ]
Install once.
Import frommlbricks.
This component ships inside the unifiedmlbricks-kitdistribution. Python imports continue to use themlbricksnamespace.
TERMINAL
pip install mlbricks-kit==1.0.0b1[ QUICK START ]
Two virtual state refinements.
The API exposes the number of virtual refinements and the small hidden width separately from the model width.
PYTHON
from mlbricks import VirtualStateAwareFFN
ffn = VirtualStateAwareFFN(
d_model=384,
state_dim=256,
depth_embedding_dim=64,
layer_index=0,
total_layers=6,
virtual_refinements=2,
virtual_hidden_dim=128,
backend="auto",
)
state = ffn.initial_state(x)
out, state = ffn(x, esa_now, esa_prev, state)[ API · MLBRICKS KIT 1.0.0B1 ]
Virtual State-Aware FFN
public surface.
This component extends StateAwareFFN with configurable virtual refinement passes while keeping the same state-aware forward contract.
CONSTRUCTOR
VirtualStateAwareFFN(
d_model,
state_dim=256,
depth_embedding_dim=64,
layer_index=0,
total_layers=1,
virtual_refinements=2,
virtual_hidden_dim=128,
use_native=None,
fused_cuda=True,
backend="auto",
)initial_state(x)Inherited state initializer.
ffn(x, esa_update, previous_esa, previous_state)Run the state-aware update with virtual refinement.
virtual_gate_mean()Return a compact diagnostic of the learned virtual gates.
reset_virtual_identity()Reset virtual refinement parameters toward their identity-oriented initialization.
set_backend(...) / resolved_backend()Use the inherited MLBricks backend surface.
[ VIRTUAL STATE-AWARE FFN ]
More refinement does not have to mean more physical layers.
Reuse a compact state refiner, let the representation evolve, and give every pass its own learned identity.