Structural Disentanglement in Bilinear MLPs via Architectural Inductive Bias

TL;DR

The paper addresses why selective unlearning and long-horizon extrapolation remain fragile, arguing that representational structure—shaped by architectural inductive bias—limits these capabilities. It introduces structural disentanglement and analyzes bilinear MLPs, where the learned operator $Q=\sum_k \alpha_k \mathbf{w}_k \mathbf{v}_k^\top$ decomposes into orthogonal modes under gradient flow, enabling independent learning and targeted edits. The authors provide a theoretical gradient-flow analysis showing mode-wise decoupling and present controlled experiments in modular arithmetic, cyclic reasoning, and Lie group dynamics demonstrating that multiplicative architectures recover ground-truth algebraic operators and enable surgical unlearning with minimal collateral damage. The study suggests that model editability and generalization hinge on representational structure induced by architecture, highlighting the central role of architectural inductive bias over post-hoc unlearning algorithms.

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.

Figures (13)

• Figure 1: Architectural Inductive Bias: Multiplicative vs. Pointwise. Comparison of unlearning dynamics for Orthogonal ($\alpha=0$, Top) and Parallel ($\alpha=1$, Bottom) tasks. Left 3 Columns (Multiplicative): Bilinear, SwiGLU, and GeGLU architectures maintain structural separation. Right 3 Columns (Pointwise): ReLU, Tanh, and Sigmoid architectures exhibit task interference.
• Figure 2: Surgical Unlearning Analysis. Detailed breakdown of how Bilinear models (Blue) outperform ReLU models (Red) in entangled scenarios. (a) Neuron specialization prevents interference. (b) Pruning follows an ideal Pareto frontier. (c) Gradient descent moves orthogonally. (d) Selectivity remains high even as interaction rank increases.
• Figure 3: Spectral entropy $H_k$ across output classes for trained models. The vertical red line denotes the true addition operator. Multiplicative architectures concentrate closer to the true value, while pointwise nonlinearities exhibit either lower entropy from procedural memorization (ReLU, Sigmoid) or higher entropy from diffuse spectral representations (Tanh).
• Figure 4: Normalized singular value decay for interaction matrices in modular multiplication. The bilinear spectra (blue) decay sharply, indicating the operators are effectively lower-dimensional.
• Figure 6: Long-horizon extrapolation on Lie Group dynamics. Left: Rigid Body Stability ($S^3$): The multiplicative models preserve the unit norm constraint better. Right: Fluid Volume Conservation ($SL(2)$): The multiplicative models preserve the determinant near 1.0 better.