2026-06-14 | MOLTs

Goal: Review Mixture of Linear Transforms Architecture

Summary: Hand write MOLT usage for a single layer

Work sessions

In	Out	Task
16:00	17:00	Understand code base

Goals

I've been feeling lately that a lot of my work has been carried out by coding agents. For today, my goal is to build understanding using the traditional techniques I used for understanding code bases before AI (using a debugger/stepping through example inputs)

Results

Nothing pretty! Just a very step by step intuition from Georg's original molt branch

from model.jumprelu import JumpReLU
from model.molt import Molt
from model.standardize import DimensionwiseInputStandardizer, DimensionwiseOutputStandardizer

print("Testing MOLTs from scratch")

# Goal is to use the MOLT class

# d_acts is the dimensions of the model
# d_features for JumpReLU is over gates which should be active
n_layers = 1
multiplier = 1
ranks = [512, 256, 128, 64, 32]
d_act = 10
# if there are 3 ranks, there are (1 + 2 + 4) = (2 ^ 3) -1
# Then we multiply by the multiplier
num_transforms = ((2 ** (len(ranks))) - 1) * multiplier

molt_layer = Molt(
    d_acts=d_act,
    N=multiplier,
    nonlinearity=JumpReLU(
        theta=0.0,
        bandwidth=1.0,
        n_layers=n_layers,
        d_features=num_transforms
    ),
    input_standardizer=DimensionwiseInputStandardizer(n_layers, d_act),
    output_standardizer=DimensionwiseOutputStandardizer(n_layers, d_act),
).to("mps")

print("Initialized MOLT layer where standardizers are not yet initialized")