[ARCH-PROP] PrefiAll

[SchweitzerGAO]

Architecture Name

PrefiAll

Parent issue

#1

Motivations

Autoregressive language models generate tokens sequentially under a causal attention mask. While this formulation provides strong generation quality and stable training, its token-by-token dependency limits decoding parallelism and can lead to high inference latency, particularly for long sequences.

Masked diffusion language models provide a complementary generation paradigm by predicting and updating multiple tokens in parallel. However, realizing this potential without sacrificing model quality remains challenging. Compared with next-token prediction, jointly recovering multiple uncertain tokens may present a harder optimization problem, and fully non-causal modeling may provide weaker inductive bias for local consistency and token-level dependencies.

PrefiAll explores whether autoregressive and diffusion-style computation can be combined within a single Transformer. The central hypothesis is that different parts of the network may benefit from different modeling constraints: diffusion-style layers can provide globally informed representations or candidate predictions, while causal layers can improve local consistency and preserve autoregressive modeling capabilities.

Proposed Architecture

PrefiAll is a hybrid Transformer architecture containing both diffusion-style and causal computation.

A subset of Transformer layers operates with a less restrictive attention pattern, such as bidirectional or block-wise attention, and processes multiple token positions jointly. These layers are intended to incorporate broader sequence context and support some degree of parallel token prediction or representation updating.

The remaining layers use causal attention and produce autoregressive predictions. They act on the intermediate representations generated by the earlier layers and may provide additional token-level refinement under a left-to-right constraint.

Several aspects of the architecture will remain configurable during the initial investigation, including:

• the number and placement of diffusion-style layers;
• the attention pattern used by those layers;
• the interface between diffusion-style and causal layers;
• the corruption or masking strategy;
• the balance between diffusion and autoregressive training objectives; and
• the generation procedure used to translate the hybrid computation into practical decoding speedups.

The initial implementation will prioritize a simple architecture that enables controlled comparison with standard autoregressive and diffusion baselines. More specialized variants will be considered only after the basic behavior of the hybrid model is understood.

Preliminary Results (if any)

No response

Experiments Plan

All primary comparisons will use matched model capacity and, where practical, aligned data, token budget, context length, optimization settings, and evaluation intervals.

Architecture and training validation

• Implement a minimal hybrid Transformer with configurable diffusion-style and causal layer allocation.
• Verify that the causal portion of the model does not access information that would invalidate autoregressive prediction.
• Compare several simple layer-allocation and attention-pattern variants.
• Evaluate alternative masking or corruption settings without committing to a single diffusion formulation.
• Study the interaction between the autoregressive and diffusion training objectives.

1B validation

• Train selected PrefiAll configurations at approximately the 1B scale using a pipeline aligned with OLMo-style training where practical.

Key metrics

• Record training loss, validation loss or perplexity, and downstream-task performance.
• Measure Tokens per Forward (TPF) and Tokens per Second (TPS) under clearly specified decoding settings.

Scaling decision

• Consider larger-scale experiments only if the 1B results show a reproducible and meaningful quality–efficiency trade-off.