Ordinal Regression Comparison¶

import argparse
from dataclasses import dataclass
import torch
from torch import Tensor
from torch.utils.data import DataLoader, TensorDataset

from torchregress.losses import CumulativeLinkLoss, OrdinalCrossEntropyLoss
from torchregress.metrics import (
    mean_absolute_class_error,
    ordinal_accuracy,
    quadratic_weighted_kappa,
)
from torchregress.utils import ordinal_predict

@dataclass(frozen=True)
class OrdinalComparisonConfig:
    seed: int = 260227
    n_train: int = 512
    n_test: int = 256
    n_features: int = 6
    num_classes: int = 5
    hidden: int = 32
    epochs: int = 30
    batch_size: int = 64
    lr: float = 1e-2

class MLP(torch.nn.Module):
    def __init__(self, n_features: int, hidden: int, out_dim: int) -> None:
        super().__init__()
        self.net = torch.nn.Sequential(
            torch.nn.Linear(n_features, hidden),
            torch.nn.ReLU(),
            torch.nn.Linear(hidden, hidden),
            torch.nn.ReLU(),
            torch.nn.Linear(hidden, out_dim),
        )
    def forward(self, x: Tensor) -> Tensor:
        return self.net(x)

def make_data(cfg: OrdinalComparisonConfig) -> tuple[Tensor, Tensor, Tensor, Tensor]:
    torch.manual_seed(cfg.seed)
    x = torch.randn(cfg.n_train + cfg.n_test, cfg.n_features)

    # Latent linear model with non-linear addition
    w = torch.tensor([0.9, -0.7, 0.4, 0.2, -0.3, 0.1])[: cfg.n_features]
    latent = x @ w + 0.35 * x[:, 0] * x[:, 1] - 0.2 * x[:, 2] ** 2 + 0.4 * torch.randn(x.shape[0])

    # Bin into classes using thresholds
    cutpoints = torch.tensor([-1.0, -0.25, 0.4, 1.1])
    y = torch.bucketize(latent, cutpoints).long()

    split = cfg.n_train
    return x[:split], y[:split], x[split:], y[split:]

def train_model(model: torch.nn.Module, loss_fn: torch.nn.Module, x_tr: Tensor, y_tr: Tensor, cfg: OrdinalComparisonConfig) -> None:
    model.train()
    optimizer = torch.optim.Adam(model.parameters(), lr=cfg.lr)
    loader = DataLoader(TensorDataset(x_tr, y_tr), batch_size=cfg.batch_size, shuffle=True)
    for _ in range(cfg.epochs):
        for xb, yb in loader:
            optimizer.zero_grad(set_to_none=True)
            logits = model(xb)
            loss = loss_fn(logits, yb)
            loss.backward()
            optimizer.step()

def evaluate(model: torch.nn.Module, x_te: Tensor, y_te: Tensor, encoding: str) -> dict[str, float]:
    model.eval()
    with torch.no_grad():
        logits = model(x_te)
    pred = ordinal_predict(logits, encoding=encoding)
    return {
        "Accuracy": float(ordinal_accuracy(pred, y_te, encoding="labels").item()),
        "OrdinalMAE": float(mean_absolute_class_error(pred, y_te, encoding="labels").item()),
        "QWK": float(quadratic_weighted_kappa(pred, y_te, encoding="labels").item()),
    }

def main() -> None:
    cfg = OrdinalComparisonConfig()
    x_tr, y_tr, x_te, y_te = make_data(cfg)

    # Compare baseline independent class logits vs the ordinal cumulative-link objective
    methods = [
        ("OrdinalCrossEntropy", MLP(cfg.n_features, cfg.hidden, cfg.num_classes), OrdinalCrossEntropyLoss(), "class_logits"),
        ("CumulativeLink", MLP(cfg.n_features, cfg.hidden, cfg.num_classes - 1), CumulativeLinkLoss(), "cumulative_logits"),
    ]

    for name, model, loss_fn, encoding in methods:
        train_model(model, loss_fn, x_tr, y_tr, cfg)
        metrics = evaluate(model, x_te, y_te, encoding)
        print(f"Method: {name}")
        print(f"  Accuracy: {metrics['Accuracy']:.4f}, OrdinalMAE: {metrics['OrdinalMAE']:.4f}, QWK: {metrics['QWK']:.4f}")

if __name__ == "__main__":
    main()