Где должна входить диффузия в языковую модель? Замена скрытых состояний, направляемая геометрией

Аннотация

Модели непрерывной диффузии в языковых задачах уступают авторегрессионным трансформерам, отчасти потому, что диффузия применяется в пространствах, плохо приспособленных для шумоподавления и восстановления токенов. Мы предлагаем DiHAL — гибрид диффузии и трансформера, управляемый геометрией, который позволяет определить, на каком этапе диффузия должна войти в предобученный трансформер. DiHAL оценивает слои с помощью геометрических прокси, выбирает дружественное к диффузии скрытое состояние в качестве интерфейса и заменяет нижнюю часть трансформера диффузионным мостом, сохраняя верхние слои и исходную языковую голову. Восстанавливая скрытое состояние выбранного слоя, а не токены, DiHAL избегает прямого перехода от непрерывного к дискретному. Эксперименты на моделях размером 8B показывают, что геометрическая оценка предсказывает эффективные мелкие встраиваемые слои при фиксированном протоколе обучения моста, а восстановление скрытого состояния превосходит базовые подходы непрерывной диффузии в сравнительном анализе, согласованном по бюджету обучения диффузии/восстановления. Эти результаты позволяют предположить, что геометрия скрытого состояния помогает определить, где в предобученных языковых моделях возможна замена на основе диффузии.

English

Continuous diffusion language models lag behind autoregressive transformers, partly because diffusion is applied in spaces poorly suited to language denoising and token recovery. We propose DiHAL, a geometry-guided diffusion-transformer hybrid that asks where diffusion should enter a pretrained transformer. DiHAL scores layers with geometry-based proxies, selects a diffusion-friendly hidden-state interface, and replaces the lower transformer prefix with a diffusion bridge while retaining the upper layers and original LM head. By reconstructing the selected-layer hidden state rather than tokens, DiHAL avoids direct continuous-to-discrete recovery. Experiments on 8B-scale backbones show that the geometry score predicts effective shallow insertion layers under a fixed bridge-training protocol and that hidden-state recovery improves over continuous diffusion baselines in a diagnostic comparison matching the diffusion/recovery training budget. These results suggest that hidden-state geometry helps identify where diffusion-based replacement is feasible inside pretrained language models.