deep_delta_learning/

transformer.rs

use burn::module::{Ignored, Module};
use burn::nn::{Embedding, EmbeddingConfig, Linear, LinearConfig, RmsNorm, RmsNormConfig};
use burn::prelude::*;

use crate::attention::{MultiHeadAttention, MultiHeadAttentionConfig};
use crate::compressor::{ChannelConvCompressor, TokenConvCompressor};
use crate::config::{CompressionVariant, DdlConfig};
use crate::delta_res_block::DeltaResBlock;
use crate::embed_expander::EmbeddingExpander;
use crate::generation::{
    AutoregressiveModel, GenerationConfig, GenerationError, GenerationResult, generate_tokens,
};
use crate::mlp::{SwiGluMlp, SwiGluMlpConfig};
use crate::spectral::{LayerDiagnostics, ModelDiagnostics};
use crate::utils::create_causal_mask;

#[cfg(feature = "spectral")]
use crate::spectral::{LayerTrace, SpectralDiagnostics, summarize_layer_traces};

#[derive(Module, Debug)]
pub struct DdlTransformerBlock<B: Backend> {
    attn_norm: RmsNorm<B>,
    attn: MultiHeadAttention<B>,
    attn_delta: DeltaResBlock<B>,
    mlp_norm: RmsNorm<B>,
    mlp: SwiGluMlp<B>,
    mlp_delta: DeltaResBlock<B>,
    token_compressor: Option<TokenConvCompressor<B>>,
    channel_compressor: Option<ChannelConvCompressor<B>>,
}

impl<B: Backend> DdlTransformerBlock<B> {
    pub fn new(config: &DdlConfig, device: &B::Device) -> Self {
        let attn = MultiHeadAttentionConfig::new(
            config.d_model,
            config.num_heads,
            config.effective_head_dim(),
            config.max_seq_len,
        )
        .with_rope_theta(config.rope_theta)
        .init(device);
        let mlp = SwiGluMlpConfig::new(config.d_model, config.effective_mlp_hidden()).init(device);
        let attn_norm = RmsNormConfig::new(config.d_model)
            .with_epsilon(1e-6)
            .init(device);
        let mlp_norm = RmsNormConfig::new(config.d_model)
            .with_epsilon(1e-6)
            .init(device);
        let token_compressor = (config.uses_matrix_state()
            && matches!(config.compression, CompressionVariant::TokenConv))
        .then(|| {
            TokenConvCompressor::<B>::new(
                config.d_model,
                config.d_value,
                config.shortconv_kernel_size,
                device,
            )
        });
        let channel_compressor = (config.uses_matrix_state()
            && matches!(config.compression, CompressionVariant::ChannelConv))
        .then(|| ChannelConvCompressor::<B>::new(config.d_model, config.d_value, device));

        Self {
            attn_norm,
            attn,
            attn_delta: DeltaResBlock::new(config, device),
            mlp_norm,
            mlp,
            mlp_delta: DeltaResBlock::new(config, device),
            token_compressor,
            channel_compressor,
        }
    }

    fn compress(&self, state: Tensor<B, 4>) -> Tensor<B, 3> {
        match (&self.token_compressor, &self.channel_compressor) {
            (Some(compressor), None) => compressor.forward(state),
            (None, Some(compressor)) => compressor.forward(state),
            _ => panic!("matrix-state block requires exactly one compressor"),
        }
    }

    pub fn forward_vector(
        &self,
        x: Tensor<B, 3>,
        mask: Option<&Tensor<B, 3>>,
    ) -> (Tensor<B, 3>, LayerDiagnostics) {
        let residual = x.clone();
        let x_ctx = self.attn_norm.forward(residual.clone());
        let attn_out = self.attn.forward(x_ctx.clone(), mask);
        let (x, attention) = self
            .attn_delta
            .forward_vector(residual, attn_out, x_ctx, x.clone());

        let residual = x.clone();
        let x_ctx = self.mlp_norm.forward(residual.clone());
        let mlp_out = self.mlp.forward(x_ctx.clone());
        let (x, mlp) = self
            .mlp_delta
            .forward_vector(residual, mlp_out, x_ctx, x.clone());

        (x, LayerDiagnostics { attention, mlp })
    }

    fn forward_vector_hidden(&self, x: Tensor<B, 3>, mask: Option<&Tensor<B, 3>>) -> Tensor<B, 3> {
        let residual = x.clone();
        let x_ctx = self.attn_norm.forward(residual.clone());
        let attn_out = self.attn.forward(x_ctx.clone(), mask);
        let x = self
            .attn_delta
            .forward_vector_hidden(residual, attn_out, x_ctx, x.clone());

        let residual = x.clone();
        let x_ctx = self.mlp_norm.forward(residual.clone());
        let mlp_out = self.mlp.forward(x_ctx.clone());
        self.mlp_delta
            .forward_vector_hidden(residual, mlp_out, x_ctx, x)
    }

    #[cfg(feature = "spectral")]
    fn forward_vector_traced(
        &self,
        x: Tensor<B, 3>,
        mask: Option<&Tensor<B, 3>>,
    ) -> (Tensor<B, 3>, LayerTrace) {
        let residual = x.clone();
        let x_ctx = self.attn_norm.forward(residual.clone());
        let attn_out = self.attn.forward(x_ctx.clone(), mask);
        let (x, attention) =
            self.attn_delta
                .forward_vector_traced(residual, attn_out, x_ctx, x.clone());

        let residual = x.clone();
        let x_ctx = self.mlp_norm.forward(residual.clone());
        let mlp_out = self.mlp.forward(x_ctx.clone());
        let (x, mlp) = self
            .mlp_delta
            .forward_vector_traced(residual, mlp_out, x_ctx, x.clone());

        (x, LayerTrace { attention, mlp })
    }

    pub fn forward_matrix(
        &self,
        state: Tensor<B, 4>,
        mask: Option<&Tensor<B, 3>>,
    ) -> (Tensor<B, 4>, LayerDiagnostics) {
        let x_lin = self.compress(state.clone());
        let x_ctx = self.attn_norm.forward(x_lin.clone());
        let attn_out = self.attn.forward(x_ctx.clone(), mask);
        let (state, attention) = self
            .attn_delta
            .forward_matrix(state, attn_out, x_ctx, x_lin);

        let x_lin = self.compress(state.clone());
        let x_ctx = self.mlp_norm.forward(x_lin.clone());
        let mlp_out = self.mlp.forward(x_ctx.clone());
        let (state, mlp) = self.mlp_delta.forward_matrix(state, mlp_out, x_ctx, x_lin);

        (state, LayerDiagnostics { attention, mlp })
    }

    fn forward_matrix_hidden(
        &self,
        state: Tensor<B, 4>,
        mask: Option<&Tensor<B, 3>>,
    ) -> Tensor<B, 4> {
        let x_lin = self.compress(state.clone());
        let x_ctx = self.attn_norm.forward(x_lin.clone());
        let attn_out = self.attn.forward(x_ctx.clone(), mask);
        let state = self
            .attn_delta
            .forward_matrix_hidden(state, attn_out, x_ctx, x_lin);

        let x_lin = self.compress(state.clone());
        let x_ctx = self.mlp_norm.forward(x_lin.clone());
        let mlp_out = self.mlp.forward(x_ctx.clone());
        self.mlp_delta
            .forward_matrix_hidden(state, mlp_out, x_ctx, x_lin)
    }

    #[cfg(feature = "spectral")]
    fn forward_matrix_traced(
        &self,
        state: Tensor<B, 4>,
        mask: Option<&Tensor<B, 3>>,
    ) -> (Tensor<B, 4>, LayerTrace) {
        let x_lin = self.compress(state.clone());
        let x_ctx = self.attn_norm.forward(x_lin.clone());
        let attn_out = self.attn.forward(x_ctx.clone(), mask);
        let (state, attention) = self
            .attn_delta
            .forward_matrix_traced(state, attn_out, x_ctx, x_lin);

        let x_lin = self.compress(state.clone());
        let x_ctx = self.mlp_norm.forward(x_lin.clone());
        let mlp_out = self.mlp.forward(x_ctx.clone());
        let (state, mlp) = self
            .mlp_delta
            .forward_matrix_traced(state, mlp_out, x_ctx, x_lin);

        (state, LayerTrace { attention, mlp })
    }
}

#[derive(Module, Debug)]
pub struct DdlTransformer<B: Backend> {
    embedding: Embedding<B>,
    embed_expander: EmbeddingExpander<B>,
    blocks: Vec<DdlTransformerBlock<B>>,
    final_token_compressor: Option<TokenConvCompressor<B>>,
    final_channel_compressor: Option<ChannelConvCompressor<B>>,
    final_norm: RmsNorm<B>,
    lm_head: Linear<B>,
    config: Ignored<DdlConfig>,
}

impl<B: Backend> DdlTransformer<B> {
    pub fn new(config: DdlConfig, device: &B::Device) -> Self {
        let blocks = (0..config.num_layers)
            .map(|_| DdlTransformerBlock::new(&config, device))
            .collect();
        let final_token_compressor = (config.uses_matrix_state()
            && matches!(config.compression, CompressionVariant::TokenConv))
        .then(|| {
            TokenConvCompressor::<B>::new(
                config.d_model,
                config.d_value,
                config.shortconv_kernel_size,
                device,
            )
        });
        let final_channel_compressor = (config.uses_matrix_state()
            && matches!(config.compression, CompressionVariant::ChannelConv))
        .then(|| ChannelConvCompressor::<B>::new(config.d_model, config.d_value, device));
        Self {
            embedding: EmbeddingConfig::new(config.vocab_size, config.d_model).init(device),
            embed_expander: EmbeddingExpander::new(
                config.d_model,
                config.d_value,
                config.embed_conv,
                config.shortconv_kernel_size,
                device,
            ),
            blocks,
            final_token_compressor,
            final_channel_compressor,
            final_norm: RmsNormConfig::new(config.d_model)
                .with_epsilon(1e-6)
                .init(device),
            lm_head: LinearConfig::new(config.d_model, config.vocab_size)
                .with_bias(false)
                .init(device),
            config: Ignored(config),
        }
    }

    fn compress_final(&self, state: Tensor<B, 4>) -> Tensor<B, 3> {
        match (&self.final_token_compressor, &self.final_channel_compressor) {
            (Some(compressor), None) => compressor.forward(state),
            (None, Some(compressor)) => compressor.forward(state),
            _ => panic!("matrix-state model requires exactly one final compressor"),
        }
    }

    fn resolve_mask(
        &self,
        input_ids: &Tensor<B, 2, Int>,
        mask: Option<&Tensor<B, 3>>,
    ) -> Tensor<B, 3> {
        match mask {
            Some(mask) => mask.clone(),
            None => {
                let [batch_size, seq_len] = input_ids.dims();
                create_causal_mask::<B>(batch_size, seq_len, &input_ids.device())
            }
        }
    }

    pub fn forward_logits(
        &self,
        input_ids: Tensor<B, 2, Int>,
        mask: Option<&Tensor<B, 3>>,
    ) -> Tensor<B, 3> {
        let mask = self.resolve_mask(&input_ids, mask);

        if self.config.d_value == 1 {
            let mut hidden = self.embedding.forward(input_ids);
            for block in &self.blocks {
                hidden = block.forward_vector_hidden(hidden, Some(&mask));
            }
            return self.lm_head.forward(self.final_norm.forward(hidden));
        }

        let mut state = self
            .embed_expander
            .forward(self.embedding.forward(input_ids));
        for block in &self.blocks {
            state = block.forward_matrix_hidden(state, Some(&mask));
        }
        let hidden = self.compress_final(state);
        self.lm_head.forward(self.final_norm.forward(hidden))
    }

    pub fn forward(
        &self,
        input_ids: Tensor<B, 2, Int>,
        mask: Option<&Tensor<B, 3>>,
    ) -> (Tensor<B, 3>, ModelDiagnostics) {
        let mask = self.resolve_mask(&input_ids, mask);
        let mut diagnostics = Vec::with_capacity(self.blocks.len());

        if self.config.d_value == 1 {
            let mut hidden = self.embedding.forward(input_ids);
            for block in &self.blocks {
                let (next_hidden, layer_diag) = block.forward_vector(hidden, Some(&mask));
                hidden = next_hidden;
                diagnostics.push(layer_diag);
            }
            let logits = self.lm_head.forward(self.final_norm.forward(hidden));
            return (
                logits,
                ModelDiagnostics {
                    layers: diagnostics,
                },
            );
        }

        let mut state = self
            .embed_expander
            .forward(self.embedding.forward(input_ids));
        for block in &self.blocks {
            let (next_state, layer_diag) = block.forward_matrix(state, Some(&mask));
            state = next_state;
            diagnostics.push(layer_diag);
        }
        let hidden = self.compress_final(state);
        let logits = self.lm_head.forward(self.final_norm.forward(hidden));
        (
            logits,
            ModelDiagnostics {
                layers: diagnostics,
            },
        )
    }

    #[cfg(feature = "spectral")]
    pub fn forward_with_spectral_diagnostics(
        &self,
        input_ids: Tensor<B, 2, Int>,
        mask: Option<&Tensor<B, 3>>,
    ) -> (Tensor<B, 3>, ModelDiagnostics, SpectralDiagnostics) {
        let mask = self.resolve_mask(&input_ids, mask);
        let mut layer_traces = Vec::with_capacity(self.blocks.len());

        if self.config.d_value == 1 {
            let mut hidden = self.embedding.forward(input_ids);
            for block in &self.blocks {
                let (next_hidden, layer_trace) = block.forward_vector_traced(hidden, Some(&mask));
                hidden = next_hidden;
                layer_traces.push(layer_trace);
            }
            let logits = self.lm_head.forward(self.final_norm.forward(hidden));
            let (diagnostics, spectral) = summarize_layer_traces(layer_traces);
            return (logits, diagnostics, spectral);
        }

        let mut state = self
            .embed_expander
            .forward(self.embedding.forward(input_ids));
        for block in &self.blocks {
            let (next_state, layer_trace) = block.forward_matrix_traced(state, Some(&mask));
            state = next_state;
            layer_traces.push(layer_trace);
        }
        let hidden = self.compress_final(state);
        let logits = self.lm_head.forward(self.final_norm.forward(hidden));
        let (diagnostics, spectral) = summarize_layer_traces(layer_traces);
        (logits, diagnostics, spectral)
    }

    pub fn max_seq_len(&self) -> usize {
        self.config.max_seq_len
    }

    pub fn generate(
        &self,
        prompt_tokens: &[usize],
        generation_config: &GenerationConfig,
        device: &B::Device,
    ) -> Result<GenerationResult, GenerationError> {
        generate_tokens(self, prompt_tokens, generation_config, device)
    }
}

impl<B: Backend> AutoregressiveModel<B> for DdlTransformer<B> {
    fn forward_logits(
        &self,
        input_ids: Tensor<B, 2, Int>,
        mask: Option<&Tensor<B, 3>>,
    ) -> Tensor<B, 3> {
        DdlTransformer::forward_logits(self, input_ids, mask)
    }

    fn max_seq_len(&self) -> usize {
        DdlTransformer::max_seq_len(self)
    }
}