Another wide transformation.
d → 4dactivation4d → d
A typical Transformer-style FFN uses a large expansion dimension. Repeating that block increases parameter storage and weight traffic substantially.
VIRTUAL DEPTH · MICRO COMPUTE
Micro Virtual
FFN
ActiveMore transformations. Much smaller per-pass footprint.
MicroVirtualFFN creates lightweight virtual depth with narrow, pass-specific SwiGLU-like refinements. Instead of adding another full-width FFN block, each virtual pass computes a compact gated update and adds it back to the current representation.
[ THE IDEA ]
Do not pay for a full layer
when a small refinement may be enough.
Conventional depth often means storing a complete new feed-forward block. MicroVirtualFFN instead inserts one or more narrow residual transformations between physical stages.
A narrow gated update.
Wgate / Wuphidden=64Wdown → residual
The micro path keeps the hidden width small and can perform several residual refinements without duplicating a full-width FFN.
IMPORTANT DISTINCTIONMicroVirtualFFN does not literally share one weight matrix across all passes.
The current implementation stores pass-specific gate, up and down tensors. Its parameter advantage comes from making each virtual passsmall. Shared-core recurrence is the design used by VirtualStateAwareFFN.
[ MULTI-PASS REFINEMENT ]
Each pass sees what the previous pass changed.
Every refinement owns its own small parameter slice. The current representation is updated after each pass, so later passes operate on a representation that has already moved.
Current normalized model representation.
Pass-specific gated nonlinear transformation.
The next pass receives the refined representation.
A different parameter slice operates on x₁.
Effective extra computation depth without another full physical FFN.
PASS k
uₖ = W↓ₖ [ SiLU(Wgateₖ xₖ) ⊙ (Wupₖ xₖ) ]xₖ₊₁ = xₖ + uₖGate weights
Each refinement index selects its own gate projection.
Value/up weights
Each refinement owns a separate narrow value projection.
Down weights
Each pass projects its micro hidden state back to model width with its own matrix.
Zero update at initialization
All down matrices start at zero, sorefine(x)begins as an exact identity path and learns virtual updates from there.
[ PARAMETER FOOTPRINT ]
Virtual depth can be narrow.
Illustrative arithmetic from the current implementation atd_model=384,hidden_dim=64, two micro refinements. A conventional comparison FFN uses384 → 1536 → 384with biases.
trainable parameters
trainable parameters
standard FFN / two micro passes
vs 1,536 in the illustrated 4× FFN
LOW-MEMORY COMPUTE
Small weights can matter as much as small arithmetic.
Edge devices often face hard limits on model storage, available RAM and memory bandwidth. A narrow virtual FFN reduces the amount of FFN parameter data that must be stored compared with adding another full-width block. Repeated passes still cost compute, so the intended trade ismore reuse / small refinement for less stored capacity, not “free depth.”
[ NATIVE EXECUTION ]
Small virtual passes also have an optimized inference path.
The optional FFNBrick native backend keeps the heavy matrix multiplies in optimized ATen/cuBLAS paths while reducing avoidable intermediate work around the multi-pass loop.
One native refine call
All configured passes can run inside one native call during eager inference.
Gate + value stay together
The CUDA path can consume packed gate/value activations without an extra intermediate copy.
Accumulate while projecting down
During multi-pass refinement, residual accumulation can be fused into the down-projection GEMM path.
Readable PyTorch loop
Autograd andtorch.compileretain the original PyTorch refinement semantics.
[ TWO FORMS OF VIRTUAL DEPTH ]
Micro when you want tiny passes.
State-aware when you want a shared controller.
MLBricks exposes both ideas because they optimize different constraints.
Small, distinct passes.
Each refinement gets its own narrow gate/up/down weights. Best aligned with experiments where the goal is to add inexpensive residual transformations.
- Pass-specific parameters
- No recurrent FFN state required
- Zero-initialized down projections
- Very small parameter footprint
Shared refiner, evolving state.
A StateAwareFFN base plus a shared condition-aware state refiner. Pass embeddings and gates give each virtual step a distinct identity.
- Shared virtual-refiner weights
- Persistent state across physical depth
- ESA/current-input conditioning reused
- Pass-specific embedding + state gate
[ RESEARCH CONTEXT ]
Parameter reuse is a real scaling axis.
These works motivate the broader memory-vs-compute design space. They do not validate MLBricks quality by themselves.
Computation can recur through depth.
Universal Transformer showed that repeating transformations through computational depth is a viable alternative to a strictly feed-forward stack of unrelated layers.
Read paper ↗Sharing parameters reduces the stored model.
ALBERT demonstrates large parameter reductions through cross-layer sharing, while its ablations also make clear that aggressively sharing FFN parameters can cost representational diversity.
Read paper ↗Weight reuse targets mobile memory traffic.
SHARP specifically studies layer sharing for LLM inference and reports lower stored MLP parameters and mobile model storage. MLBricks uses a different architecture, but the system motivation is closely related.
Read paper ↗[ 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 micro virtual passes.
The direct API is deliberately small. Userefine()when you want all configured residual refinements applied sequentially.
PYTHON
from mlbricks import MicroVirtualFFN
micro = MicroVirtualFFN(
d_model=384,
hidden_dim=64,
refinements=2,
backend="auto",
)
# sequential residual virtual passes
x = micro.refine(x)[ API · MLBRICKS KIT 1.0.0B1 ]
MicroVirtualFFN
public surface.
A compact refinement block with a small constructor and explicit single-pass or multi-pass entry points.
CONSTRUCTOR
MicroVirtualFFN(
d_model,
hidden_dim=64,
refinements=1,
use_native=None,
fused_cuda=True,
backend="auto",
)ffn(x, refinement_index=0)Run one configured refinement pass.
refine(x)Apply all configured refinements sequentially.
reset_identity()Reset the refinement parameters toward the component's identity-oriented initialization.
set_backend(...) / resolved_backend()Select and inspect the MLBricks execution route.
[ MICRO VIRTUAL FFN ]
Spend a little more compute instead of storing another full FFN.
MicroVirtualFFN is an experiment in narrow virtual depth for parameter-sensitive models, edge hardware and composable MLBricks architectures.