Structural Disentanglement in Bilinear MLPs via Architectural Inductive Bias

TL;DR

This paper demonstrates that bilinear MLPs with multiplicative interactions have a non-mixing property that facilitates structural disentanglement, improving model interpretability and unlearning capabilities.

Contribution

The study provides a mathematical foundation showing bilinear parameterizations enable orthogonal subspace representations, aiding structural disentanglement and model editability.

Findings

Bilinear MLPs exhibit a non-mixing property under gradient flow.

Multiplicative architectures recover true algebraic operators.

Architectural bias enhances unlearning and generalization.

Abstract

Selective unlearning and long-horizon extrapolation remain fragile in modern neural networks, even when tasks have underlying algebraic structure. In this work, we argue that these failures arise not solely from optimization or unlearning algorithms, but from how models structure their internal representations during training. We explore if having explicit multiplicative interactions as an architectural inductive bias helps in structural disentanglement, through Bilinear MLPs. We show analytically that bilinear parameterizations possess a `non-mixing' property under gradient flow conditions, where functional components separate into orthogonal subspace representations. This provides a mathematical foundation for surgical model modification. We validate this hypothesis through a series of controlled experiments spanning modular arithmetic, cyclic reasoning, Lie group dynamics, and targeted unlearning benchmarks. Unlike pointwise nonlinear networks, multiplicative architectures are able to recover true operators aligned with the underlying algebraic structure. Our results suggest that model editability and generalization are constrained by representational structure, and that architectural inductive bias plays a central role in enabling reliable unlearning.