dew.nn.gemma3n

Gemma 3n’s residual stream: AltUp, the LAuReL block and activation sparsity.

Gemma3nTextModel (modeling_gemma3n.py) carries altup_num_inputs copies of the residual stream. The embeddings are the first; each other copy is the embeddings through its own projection, rescaled to the embeddings’ RMS magnitude.

Every layer predicts all copies from the active one (Gemma3nTextAltUp.predict), runs the transformer block on that prediction, and corrects every copy by the block’s innovation (Gemma3nTextAltUp.correct). It then adds the per-layer input’s contribution to the copies past the first. After the last layer those copies are projected back and rescaled to the first’s magnitude, and their mean is what the final norm reads.

Gemma3nTextLaurelBlock is the learned augmented residual: a rank-laurel_rank map of the block’s normed input, normed and added back, which the block averages with the attention residual over sqrt(2).

Gemma3nTextMLP._gaussian_topk keeps, per token, the gate activations above the mean by norm.ppf(sparsity) standard deviations, so a sparsity of 0.95 zeros about 95% of them before the activation.

Name	Summary
MAGNITUDE_EPSILON
AltUp	How many copies of the residual stream a model carries and which one its blocks run on, under the reference’s names (configuration_gemma3n.py).
AltUpLayer	One layer’s AltUp: predict before its block and correct after.
LaurelBlock	x + post_laurel_norm(linear_right(linear_left(x))).
gaussian_topk	Zero all but the top 1 - sparsity fraction of each row, assuming the row is Gaussian: the cutoff is the row’s mean plus norm.ppf(sparsity) of its population standard deviation, and what is above it is kept as its distance above (modeling_gemma3n.py, Gemma3nTextMLP._gaussian_topk).
rescale_to	x scaled to target’s RMS magnitude per token, the magnitude of x floored at MAGNITUDE_EPSILON (the reference’s new_magnitude).

MAGNITUDE_EPSILON

attribute source

MAGNITUDE_EPSILON = 1e-05

AltUp

dataclass source

class AltUp(
    num_inputs: int = 4,
    active_idx: int = 0,
    coef_clip: float | None = 120.0,
    correct_scale: bool = True,
)

How many copies of the residual stream a model carries and which one its blocks run on, under the reference’s names (configuration_gemma3n.py).

AltUpLayer

Flax module source

class AltUpLayer(
    norm_eps: float = 1e-06,
    dtype: Dtype | None = None,
    precision: PrecisionLike = None,
)

One layer’s AltUp: predict before its block and correct after.

The stream is [num_inputs, B, S, D]. Both steps read the modalities of a token, the tanh of its normed, 1 / D-scaled active copy through modality_router. predict maps them to a num_inputs by num_inputs matrix per token (prediction_coefs) that mixes the copies, added to the copies; correct maps them to one coefficient per copy (correction_coefs, plus one) that scales the block’s innovation, the activated output minus the active prediction, added to every prediction. coef_clip bounds both coefficient weights while training, as the reference clamps them (modeling_gemma3n.py, Gemma3nTextAltUp).

AltUpLayer.setup

def setup()

AltUpLayer.modalities

def modalities(x)

AltUpLayer.predict

def predict(stream, train: bool = False)

AltUpLayer.correct

def correct(predictions, activated, train: bool = False)

AltUpLayer.scale_corrected_output

def scale_corrected_output(corrected)

LaurelBlock

Flax module source

class LaurelBlock(
    norm_eps: float = 1e-06,
    dtype: Dtype | None = None,
    precision: PrecisionLike = None,
)

x + post_laurel_norm(linear_right(linear_left(x))).

gaussian_topk

function source

def gaussian_topk(x, sparsity: float)

Zero all but the top 1 - sparsity fraction of each row, assuming the row is Gaussian: the cutoff is the row’s mean plus norm.ppf(sparsity) of its population standard deviation, and what is above it is kept as its distance above (modeling_gemma3n.py, Gemma3nTextMLP._gaussian_topk).

rescale_to

function source

def rescale_to(x, target)

x scaled to target’s RMS magnitude per token, the magnitude of x floored at MAGNITUDE_EPSILON (the reference’s new_magnitude).

A zero target (a zero pad embedding at a padded slot) has magnitude zero with a zero gradient; sqrt’s infinite derivative there would turn the masked slot’s zero upstream gradient into NaN for every parameter it touches.