Skip to content

xlm.modules.regression

RateHead

Bases: Module

Unified regression head for predicting positive scalars (rates, lengths, etc.).

Supports two output modes
  • "continuous": activation function + affine scaling -> scalar per token
  • "discretized": softmax over bins -> expected-value scalar per token

Optional AdaLN conditioning and MLP vs single-linear architecture.