TA-SAE: Untangling Temporal Polysemanticity in Dictionary Learning for Rectified Flow Transformers

As text-to-image rectified flow models become increasingly capable, understanding their internal representations is essential for interpretability, control, and model safety. While sparse autoencoders (SAEs) have emerged as a promising tool for extracting interpretable features, existing works do not shed light on how individual concepts evolve over time internally, due to a reliance on timestep-specific SAEs. We argue that the reliance on timestep-dependent SAEs is in large part due to a lack of any inductive bias in SAEs towards the generation trajectory, crippling their ability to be used across long trajectories. We show that this creates failure modes missed by standard SAE diagnostics - a feature can reconstruct well while changing semantic identity across denoising time, and a concept can be handed off across different features at different denoising stages. We explore each of these phenomena, which we refer to as temporal feature absorption and temporal feature splitting respectively. In addition, to combat them, we introduce trajectory-aware SAEs (TA-SAEs), which learn a single dictionary over denoising time while assigning each feature a learned persistence coefficient, regularizing persistent features to remain stable across adjacent timesteps while allowing transient features to capture time-local denoising computations. On FLUX.1-dev block-18 activations, we show that TA-SAE preserves standard SAE quality, reduces absorption and splitting, and improves temporal semantic stability. To our knowledge, we are the first to formalize temporal feature absorption and temporal feature splitting as feature--concept identity failures in diffusion SAEs, and propose an alternative to traditional SAEs to combat them.