lightly.models

The lightly.models package provides model implementations.

Example implementations for all models can be found here: Model Examples

The package contains an implementation of the commonly used ResNet and adaptations of the architecture which make self-supervised learning simpler.

The package also hosts the Lightly model zoo - a list of downloadable ResNet checkpoints.

.resnet

Custom ResNet Implementation

Note that the architecture we present here differs from the one used in torchvision. We replace the first 7x7 convolution by a 3x3 convolution to make the model faster and run better on smaller input image resolutions.

Furthermore, we introduce a resnet-9 variant for extra small models. These can run for example on a microcontroller with 100kBytes of storage.

class lightly.models.resnet.BasicBlock(in_planes: int, planes: int, stride: int = 1, num_splits: int = 0)

Implementation of the ResNet Basic Block.

in_planes

Number of input channels.

planes

Number of channels.

stride

Stride of the first convolutional.

forward(x: Tensor) → Tensor

Forward pass through basic ResNet block.

Parameters

x – Tensor of shape bsz x channels x W x H

Returns

Tensor of shape bsz x channels x W x H

class lightly.models.resnet.Bottleneck(in_planes: int, planes: int, stride: int = 1, num_splits: int = 0)

Implementation of the ResNet Bottleneck Block.

in_planes

Number of input channels.

planes

Number of channels.

stride

Stride of the first convolutional.

forward(x: Tensor) → Tensor

Forward pass through bottleneck ResNet block.

Parameters

x – Tensor of shape bsz x channels x W x H

Returns

Tensor of shape bsz x channels x W x H

class lightly.models.resnet.ResNet(block: type[lightly.models.resnet.BasicBlock] = <class 'lightly.models.resnet.BasicBlock'>, layers: ~typing.List[int] = [2, 2, 2, 2], num_classes: int = 10, width: float = 1.0, num_splits: int = 0)

ResNet implementation.

[1] Kaiming He, Xiangyu Zhang, Shaoqing Ren, Jian Sun Deep Residual Learning for Image Recognition. arXiv:1512.03385

block

ResNet building block type.

layers

List of blocks per layer.

num_classes

Number of classes in final softmax layer.

width

Multiplier for ResNet width.

forward(x: Tensor) → Tensor

Forward pass through ResNet.

Parameters

x – Tensor of shape bsz x channels x W x H

Returns

Output tensor of shape bsz x num_classes

lightly.models.resnet.ResNetGenerator(name: str = 'resnet-18', width: float = 1, num_classes: int = 10, num_splits: int = 0) → ResNet

Builds and returns the specified ResNet.

Parameters

  • name – ResNet version from resnet-{9, 18, 34, 50, 101, 152}.

  • width – ResNet width.

  • num_classes – Output dim of the last layer.

  • num_splits – Number of splits to use for SplitBatchNorm (for MoCo model). Increase this number to simulate multi-gpu behavior. E.g. num_splits=8 simulates a 8-GPU cluster. num_splits=0 uses normal PyTorch BatchNorm.

Returns

ResNet as nn.Module.

Examples

>>> # binary classifier with ResNet-34
>>> from lightly.models import ResNetGenerator
>>> resnet = ResNetGenerator("resnet-34", num_classes=2)

.zoo

Lightly Model Zoo

lightly.models.zoo.checkpoints() → List[str]

Returns the Lightly model zoo as a list of checkpoints.

Checkpoints:
ResNet-9:

SimCLR with width = 0.0625 and num_ftrs = 16

ResNet-9:

SimCLR with width = 0.125 and num_ftrs = 16

ResNet-18:

SimCLR with width = 1.0 and num_ftrs = 16

ResNet-18:

SimCLR with width = 1.0 and num_ftrs = 32

ResNet-34:

SimCLR with width = 1.0 and num_ftrs = 16

ResNet-34:

SimCLR with width = 1.0 and num_ftrs = 32

Returns

A list of available checkpoints as URLs.

The lightly.models.modules package provides reusable modules.

This package contains reusable modules such as the NNmemoryBankModule which can be combined with any lightly model.

.nn_memory_bank

Nearest Neighbour Memory Bank Module

class lightly.models.modules.nn_memory_bank.NNMemoryBankModule(size: Union[int, Sequence[int]])

Nearest Neighbour Memory Bank implementation

This class implements a nearest neighbour memory bank as described in the NNCLR paper[0]. During the forward pass we return the nearest neighbour from the memory bank.

[0] NNCLR, 2021, https://arxiv.org/abs/2104.14548

size

Size of the memory bank as (num_features, dim) tuple. If num_features is 0 then the memory bank is disabled.

Examples

>>> criterion = NTXentLoss(temperature=0.1)
>>>
>>> nn_replacer = NNMemoryBankModule(size=(2**16, 128))
>>>
>>> # forward pass
>>> (z0, p0), (z1, p1) = model(x0), model(x1)
>>> z0 = nn_replacer(z0.detach(), update=False)
>>> z1 = nn_replacer(z1.detach(), update=True)
>>>
>>> loss = 0.5 * (criterion(z0, p1) + criterion(z1, p0))
forward(output: Tensor, update: bool = False) → Tensor

Returns nearest neighbour of output tensor from memory bank

Parameters

  • output – The torch tensor for which you want the nearest neighbour

  • update – If True updated the memory bank by adding output to it

.heads

Projection and Prediction Heads for Self-supervised Learning

class lightly.models.modules.heads.BYOLPredictionHead(input_dim: int = 256, hidden_dim: int = 4096, output_dim: int = 256)

Prediction head used for BYOL.

“This MLP consists in a linear layer with output size 4096 followed by batch normalization, rectified linear units (ReLU), and a final linear layer with output dimension 256.” [0]

class lightly.models.modules.heads.BYOLProjectionHead(input_dim: int = 2048, hidden_dim: int = 4096, output_dim: int = 256)

Projection head used for BYOL.

“This MLP consists in a linear layer with output size 4096 followed by batch normalization, rectified linear units (ReLU), and a final linear layer with output dimension 256.” [0]

class lightly.models.modules.heads.BarlowTwinsProjectionHead(input_dim: int = 2048, hidden_dim: int = 8192, output_dim: int = 8192)

Projection head used for Barlow Twins.

“The projector network has three linear layers, each with 8192 output units. The first two layers of the projector are followed by a batch normalization layer and rectified linear units.” [0]

class lightly.models.modules.heads.CAPIProjectionHead(input_dim: int = 384, num_clusters: int = 16384, bias: bool = False, weight_norm: bool = False)

Projection head for CAPI [0].

L2-normalizes the input features and projects them with a single linear layer to cluster logits. The linear layer weights are the cluster prototypes. The same head is used for the student and, with its own weights, for the teacher’s online clustering.

input_dim

The input dimension of the head.

num_clusters

Number of clusters, i.e. the number of output prototypes.

bias

If True, the linear layer uses a bias term.

weight_norm

If True, the prototypes are weight-normalized to unit norm with a fixed scale, as used for the student head in the reference. The teacher’s clustering head leaves this disabled.

Examples

>>> # project 384-dimensional features to 16384 cluster logits
>>> head = CAPIProjectionHead(input_dim=384, num_clusters=16384)
>>>
>>> # features has shape (batch_size, sequence_length, 384)
>>> logits = head(features)  # (batch_size, sequence_length, 16384)
forward(x: Tensor) → Tensor

Projects L2-normalized features to cluster logits.

Parameters

x – Input features of shape (…, input_dim).

Returns

Cluster logits of shape (…, num_clusters).

class lightly.models.modules.heads.DINOProjectionHead(input_dim: int = 2048, hidden_dim: int = 2048, bottleneck_dim: int = 256, output_dim: int = 65536, batch_norm: bool = False, freeze_last_layer: int = -1, norm_last_layer: bool = True)

Projection head used in DINO.

“The projection head consists of a 3-layer multi-layer perceptron (MLP) with hidden dimension 2048 followed by l2 normalization and a weight normalized fully connected layer with K dimensions, which is similar to the design from SwAV [1].” [0]

input_dim

The input dimension of the head.

hidden_dim

The hidden dimension.

bottleneck_dim

Dimension of the bottleneck in the last layer of the head.

output_dim

The output dimension of the head.

batch_norm

Whether to use batch norm or not. Should be set to False when using a vision transformer backbone.

freeze_last_layer

Number of epochs during which we keep the output layer fixed. Typically doing so during the first epoch helps training. Try increasing this value if the loss does not decrease.

norm_last_layer

Whether or not to weight normalize the last layer of the DINO head. Not normalizing leads to better performance but can make the training unstable.

cancel_last_layer_gradients(current_epoch: int) → None

Cancel last layer gradients to stabilize the training.

forward(x: Tensor) → Tensor

Computes one forward pass through the head.

class lightly.models.modules.heads.DINOv2ProjectionHead(input_dim: int = 2048, hidden_dim: int = 2048, bottleneck_dim: int = 256, output_dim: int = 65536, batch_norm: bool = False)
forward(x: Tensor) → Tensor

Computes one forward pass through the projection head.

Parameters

x – Input of shape bsz x num_ftrs.

class lightly.models.modules.heads.DenseCLProjectionHead(input_dim: int = 2048, hidden_dim: int = 2048, output_dim: int = 128)

Projection head for DenseCL [0].

The projection head consists of a 2-layer MLP. It can be used for global and local features.

class lightly.models.modules.heads.LeJEPAProjectionHead(input_dim: int, hidden_dim: int = 2048, output_dim: int = 64, num_layers: int = 3)

Projection head used for LeJEPA.

LeJEPA applies a projector (a multi-layer perceptron) on top of the backbone and computes the SIGReg and invariance losses on the projected embeddings. [0] The default is a 3-layer MLP with BatchNorm and ReLU, matching the reference implementation. [1]

class lightly.models.modules.heads.LeWMProjectionHead(input_dim: int, hidden_dim: int = 2048, output_dim: int = 192, batch_norm: bool = True)

Projection head used for LeWM.

Warning

Experimental. This head may change in a minor release.

LeWM projects the backbone features to the width of the predictor and computes the prediction and SIGReg losses on the projected embeddings. The head is an MLP with one hidden layer: a linear layer, batch normalization, a GELU and a second linear layer, as in the official implementation [1].

class lightly.models.modules.heads.MMCRProjectionHead(input_dim: int = 2048, hidden_dim: int = 8192, output_dim: int = 512, num_layers: int = 2, batch_norm: bool = True, use_bias: bool = False)

Projection head used for MMCR.

“Following Chen et al. (14), we append a small perceptron to the output of the average pooling layer of the ResNet so that zi = g(h(xi)), where h is the ResNet and g is the MLP.” [0]

class lightly.models.modules.heads.MSNProjectionHead(input_dim: int = 768, hidden_dim: int = 2048, output_dim: int = 256)

Projection head for MSN [0].

“We train with a 3-layer projection head with output dimension 256 and batch-normalization at the input and hidden layers..” [0]

Code inspired by [1].

input_dim

Input dimension, default value 768 is for a ViT base model.

hidden_dim

Hidden dimension.

output_dim

Output dimension.

class lightly.models.modules.heads.MoCoProjectionHead(input_dim: int = 2048, hidden_dim: int = 2048, output_dim: int = 128, num_layers: int = 2, batch_norm: bool = False)

Projection head used for MoCo.

“(…) we replace the fc head in MoCo with a 2-layer MLP head (hidden layer 2048-d, with ReLU)” [1]

“The projection head is a 3-layer MLP. The prediction head is a 2-layer MLP. The hidden layers of both MLPs are 4096-d and are with ReLU; the output layers of both MLPs are 256-d, without ReLU. In MoCo v3, all layers in both MLPs have BN” [2]

class lightly.models.modules.heads.NNCLRPredictionHead(input_dim: int = 256, hidden_dim: int = 4096, output_dim: int = 256)

Prediction head used for NNCLR.

“The architecture of the prediction MLP g is 2 fully-connected layers of size [4096,d]. The hidden layer of the prediction MLP is followed by batch-norm and ReLU. The last layer has no batch-norm or activation.” [0]

class lightly.models.modules.heads.NNCLRProjectionHead(input_dim: int = 2048, hidden_dim: int = 2048, output_dim: int = 256)

Projection head used for NNCLR.

“The architectureof the projection MLP is 3 fully connected layers of sizes [2048,2048,d] where d is the embedding size used to apply the loss. We use d = 256 in the experiments unless otherwise stated. All fully-connected layers are followed by batch-normalization [36]. All the batch-norm layers except the last layer are followed by ReLU activation.” [0]

class lightly.models.modules.heads.ProjectionHead(blocks: Sequence[Union[Tuple[int, int, Optional[Module], Optional[Module]], Tuple[int, int, Optional[Module], Optional[Module], bool]]])

Base class for all projection and prediction heads.

Parameters

blocks – List of tuples, each denoting one block of the projection head MLP. Each tuple reads (in_features, out_features, batch_norm_layer, non_linearity_layer, use_bias (optional)).

Examples

>>> # the following projection head has two blocks
>>> # the first block uses batch norm an a ReLU non-linearity
>>> # the second block is a simple linear layer
>>> projection_head = ProjectionHead([
>>>     (256, 256, nn.BatchNorm1d(256), nn.ReLU()),
>>>     (256, 128, None, None)
>>> ])
forward(x: Tensor) → Tensor

Computes one forward pass through the projection head.

Parameters

x – Input of shape bsz x num_ftrs.

class lightly.models.modules.heads.SMoGPredictionHead(input_dim: int = 128, hidden_dim: int = 2048, output_dim: int = 128)

Prediction head used for SMoG.

“The two kinds of head are both a two-layer MLP and their hidden layer is followed by a BatchNorm [28] and an activation function. (…) The output layer of projection head also has BN” [0]

class lightly.models.modules.heads.SMoGProjectionHead(input_dim: int = 2048, hidden_dim: int = 2048, output_dim: int = 128)

Projection head used for SMoG.

“The two kinds of head are both a two-layer MLP and their hidden layer is followed by a BatchNorm [28] and an activation function. (…) The output layer of projection head also has BN” [0]

class lightly.models.modules.heads.SMoGPrototypes(group_features: Tensor, beta: float)

SMoG prototypes module for synchronous momentum grouping.

Parameters

  • group_features – Tensor containing the group features.

  • beta – Beta parameter for momentum updating.

assign_groups(x: Tensor) → Tensor

Assigns each representation in x to a group based on cosine similarity.

Parameters

x – Tensor of shape (bsz, dim).

Returns

Tensor of shape (bsz,) indicating group assignments.

forward(x: Tensor, group_features: Tensor, temperature: float = 0.1) → Tensor

Computes the logits for given model outputs and group features.

Parameters

  • x – Tensor of shape bsz x dim.

  • group_features – Momentum updated group features of shape n_groups x dim.

  • temperature – Temperature parameter for calculating the logits.

Returns

The computed logits.

get_updated_group_features(x: Tensor) → Tensor

Performs the synchronous momentum update of the group vectors.

Parameters

x – Tensor of shape bsz x dim.

Returns

The updated group features.

set_group_features(x: Tensor) → None

Sets the group features and asserts they don’t require gradient.

class lightly.models.modules.heads.SimCLRProjectionHead(input_dim: int = 2048, hidden_dim: int = 2048, output_dim: int = 128, num_layers: int = 2, batch_norm: bool = True)

Projection head used for SimCLR.

“We use a MLP with one hidden layer to obtain zi = g(h) = W_2 * σ(W_1 * h) where σ is a ReLU non-linearity.” [0]

“We use a 3-layer MLP projection head on top of a ResNet encoder.” [1]

class lightly.models.modules.heads.SimSiamPredictionHead(input_dim: int = 2048, hidden_dim: int = 512, output_dim: int = 2048)

Prediction head used for SimSiam.

“The prediction MLP (h) has BN applied to its hidden fc layers. Its output fc does not have BN (…) or ReLU. This MLP has 2 layers.” [0]

class lightly.models.modules.heads.SimSiamProjectionHead(input_dim: int = 2048, hidden_dim: int = 2048, output_dim: int = 2048)

Projection head used for SimSiam.

“The projection MLP (in f) has BN applied to each fully-connected (fc) layer, including its output fc. Its output fc has no ReLU. The hidden fc is 2048-d. This MLP has 3 layers.” [0]

class lightly.models.modules.heads.SwaVProjectionHead(input_dim: int = 2048, hidden_dim: int = 2048, output_dim: int = 128)

Projection head used for SwaV.

class lightly.models.modules.heads.SwaVPrototypes(input_dim: int = 128, n_prototypes: Union[List[int], int] = 3000, n_steps_frozen_prototypes: int = 0)

Multihead Prototypes used for SwaV.

Each output feature is assigned to a prototype, SwaV solves the swapped prediction problem where the features of one augmentation are used to predict the assigned prototypes of the other augmentation.

input_dim

The input dimension of the head.

n_prototypes

Number of prototypes.

n_steps_frozen_prototypes

Number of steps during which we keep the prototypes fixed.

Examples

>>> # use features with 128 dimensions and 512 prototypes
>>> prototypes = SwaVPrototypes(128, 512)
>>>
>>> # pass batch through backbone and projection head.
>>> features = model(x)
>>> features = nn.functional.normalize(features, dim=1, p=2)
>>>
>>> # logits has shape bsz x 512
>>> logits = prototypes(features)
forward(x: Tensor, step: Optional[int] = None) → Union[Tensor, List[Tensor]]

Forward pass of the SwaVPrototypes module.

Parameters

  • x – Input tensor.

  • step – Current training step.

Returns

The logits after passing through the prototype heads. Returns a single tensor if there’s one prototype head, otherwise returns a list of tensors.

normalize() → None

Normalizes the prototypes so that they are on the unit sphere.

class lightly.models.modules.heads.TiCoProjectionHead(input_dim: int = 2048, hidden_dim: int = 4096, output_dim: int = 256)

Projection head used for TiCo.

“This MLP consists in a linear layer with output size 4096 followed by batch normalization, rectified linear units (ReLU), and a final linear layer with output dimension 256.” [0]

class lightly.models.modules.heads.VICRegProjectionHead(input_dim: int = 2048, hidden_dim: int = 8192, output_dim: int = 8192, num_layers: int = 3)

Projection head used for VICReg.

“The projector network has three linear layers, each with 8192 output units. The first two layers of the projector are followed by a batch normalization layer and rectified linear units.” [0]

class lightly.models.modules.heads.VicRegLLocalProjectionHead(input_dim: int = 2048, hidden_dim: int = 8192, output_dim: int = 8192)

Projection head used for the local head of VICRegL.

“The projector network has three linear layers. The first two layers of the projector are followed by a batch normalization layer and rectified linear units.” [0]

.world_model

Building blocks for world models.

A world model predicts how an environment evolves. A latent world model does it in embedding space: it predicts the embedding of the next frame from the embeddings of the past frames, optionally conditioned on the action that was taken.

The modules here are parts that latent world models share. ActionEncoder turns what the agent did into an embedding, LatentDynamicsPredictor maps embeddings and optional actions to future embeddings, PredictorBlock is its conditioned transformer block, and AdaLNZero is the AdaLN-Zero conditioning projection – both reusable on their own.

The image encoder is deliberately absent. A predictor reads embeddings, not pixels, so any backbone in lightly.models works as an encoder, trained from scratch or frozen and pretrained. Reconstruction decoders and reward, value or continue heads belong here too when a method trains them jointly with the dynamics.

Out of scope: environments, trajectory collection and planning solvers. They need an environment, an episode or a discount factor to be meaningful, which makes them the job of a control library.

Warning

This subpackage is new. The shapes and signatures here – including the conditional and causal options, PredictorBlock and AdaLNZero – may change in a minor release while the remaining world model methods are added.

class lightly.models.modules.world_model.conditioning.ActionEncoder(*, action_dim: int, output_dim: int, hidden_dim: int | None = None)

Maps continuous action vectors to conditioning embeddings.

The module applies the same MLP to every timestep of a batch of action sequences. Concatenate the action with the proprioceptive state before calling it to condition on both.

action_dim

Dimension of a single action vector.

output_dim

Dimension of the conditioning embedding.

mlp

The multi-layer perceptron applied to each action vector.

Examples

>>> action_encoder = ActionEncoder(action_dim=2, output_dim=192)
>>> actions = torch.randn(8, 4, 2)  # (batch, time, action_dim)
>>> action_emb = action_encoder(actions)  # (8, 4, 192)
forward(actions: Tensor) → Tensor

Encode a batch of action sequences.

Parameters

actions – Actions of shape (B, T, action_dim).

Returns

Action embeddings of shape (B, T, output_dim).

class lightly.models.modules.world_model.predictor.LatentDynamicsPredictor(*, num_frames: int, input_dim: int, hidden_dim: int, depth: int, num_heads: int, output_dim: int | None = None, mlp_ratio: float = 4.0, dropout: float = 0.0, conditional: bool = True, causal: bool = True)

Predicts the next frame embedding from past embeddings and optional actions.

Warning

Experimental. The shapes and signatures here may change in a minor release while the remaining world model methods are added.

The predictor is a transformer over a sequence of frame embeddings. With causal=True frame t attends to the frames 0 to t, and position information is a learned embedding over frames. With conditional=True (LeWM) the action embedding of frame t produces the shift, scale and gate parameters of every transformer block for that frame, which is the AdaLN-Zero conditioning that LeWM uses. With conditional=False the predictor takes no actions and its blocks are plain transformer blocks – the actionless predictor.

Reference:

LeWorldModel, 2026, https://arxiv.org/abs/2603.19312

num_frames

Maximum number of frames in the context window.

input_dim

Dimension of an input frame embedding.

output_dim

Dimension of a predicted frame embedding.

conditional

Whether the predictor conditions on actions.

causal

Whether frame t attends only to frames 0 to t.

Examples

>>> predictor = LatentDynamicsPredictor(
...     num_frames=4, input_dim=192, hidden_dim=384, depth=6, num_heads=6
... )
>>> emb = torch.randn(8, 4, 192)  # (batch, time, dim)
>>> action_emb = torch.randn(8, 4, 192)
>>> pred = predictor(emb, action_emb=action_emb)  # (8, 4, 192)
>>> # pred[:, t] is the prediction of emb[:, t + 1]
>>>
>>> # An actionless predictor takes no action embeddings:
>>> predictor = LatentDynamicsPredictor(
...     num_frames=4,
...     input_dim=192,
...     hidden_dim=384,
...     depth=6,
...     num_heads=6,
...     conditional=False,
... )
>>> pred = predictor(emb)  # (8, 4, 192)
forward(embeddings: Tensor, action_emb: torch.Tensor | None = None) → Tensor

Predict the next frame embedding for every frame in the context.

Parameters

  • embeddings – Frame embeddings of shape (B, T, input_dim).

  • action_emb – Action embeddings of shape (B, T, input_dim), required when the predictor is conditional and None otherwise.

    action_emb[:, t] is the action taken at frame t, which leads to frame t + 1. A shape check cannot catch a phase error here, so a training step reads:

    predicted = predictor(emb[:, :-1], action_emb=act[:, :-1])
target = emb[:, 1:]
Returns

Predicted embeddings of shape (B, T, output_dim). Entry t predicts the frame embedding at t + 1, so the last entry has no target inside the clip.

rollout(embeddings: Tensor, action_emb: torch.Tensor | None = None, steps: int = 1) → Tensor

Roll the dynamics forward in latent space.

The predictor consumes embeddings as context and then feeds its own predictions back as input for steps further frames. The context window slides, so the predictor never sees more than num_frames frames at a time. This is the operation that a planner calls to score candidate action sequences.

Parameters

  • embeddings – Context embeddings of shape (B, T_ctx, input_dim).

  • action_emb – Action embeddings of shape (B, T_ctx + steps - 1, input_dim), required when the predictor is conditional and None otherwise.

  • steps – Number of frames to predict beyond the context.

Returns

The predicted embeddings of shape (B, steps, output_dim).

class lightly.models.modules.world_model.predictor.PredictorBlock(*, hidden_dim: int, num_heads: int, mlp_ratio: float, dropout: float, conditional: bool = True)

Transformer block with optional AdaLN-Zero conditioning.

Warning

Experimental. This block may change in a minor release.

With conditional=True the block follows the AdaLN-Zero scheme of DiT: a per-timestep conditioning embedding cond produces the shift, scale and gate parameters, and the gate projections start at zero so the block is an identity at initialization. cond is opaque – an action embedding for LeWM, a diffusion timestep or a goal embedding for another method.

With conditional=False the block is a plain pre-norm transformer block that ignores cond, which is the actionless / unconditional predictor.

forward(x: Tensor, attn_mask: torch.Tensor | None = None, cond: torch.Tensor | None = None) → Tensor

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class lightly.models.modules.world_model.predictor.AdaLNZero(hidden_dim: int)

AdaLN-Zero conditioning projection.

Maps a conditioning embedding of shape (..., hidden_dim) to the six per-sublayer modulation parameters (shift, scale and gate for the attention and the MLP sublayers), returned concatenated as (..., 6 * hidden_dim); the caller splits them with chunk(6, dim=-1). The projection is zero-initialized, so a block conditioned by it is the identity at initialization (AdaLN-Zero).

AdaLNZero subclasses Sequential so the linear layer stays at child index 1, which keeps the parameter names stable.

Warning

Experimental. This module may change in a minor release.