diff --git a/astrai/config/model_config.py b/astrai/config/model_config.py
index ceb12db..d0c1cf2 100644
--- a/astrai/config/model_config.py
+++ b/astrai/config/model_config.py
@@ -70,10 +70,12 @@ class EncoderConfig(BaseModelConfig):
     rope_theta: Optional[float] = None
     rope_scaling: Optional[dict] = None
 
+    attn_type: str = "gqa"
     n_heads: Optional[int] = None
     n_kv_heads: Optional[int] = None
     use_qk_norm: Optional[bool] = None
     use_gated_attention: Optional[bool] = None
 
+    ffn_type: str = "mlp"
     pooling_type: Optional[str] = None
     normalize_embeddings: Optional[bool] = None
diff --git a/astrai/model/components/decoder_block.py b/astrai/model/components/decoder_block.py
index 60b263f..d5c71aa 100644
--- a/astrai/model/components/decoder_block.py
+++ b/astrai/model/components/decoder_block.py
@@ -1,3 +1,4 @@
+from dataclasses import asdict
 from typing import Optional
 
 import torch.nn as nn
@@ -10,35 +11,13 @@ from astrai.model.components.norm import RMSNorm
 
 
 class DecoderBlock(nn.Module):
-    def __init__(
-        self,
-        dim: int,
-        n_heads: int,
-        dim_ffn: int,
-        n_kv_heads: int,
-        norm_eps: float,
-        use_qk_norm: bool,
-        use_gated_attention: bool,
-        layer_id: int,
-        attn_type: str = "gqa",
-        ffn_type: str = "mlp",
-        **kwargs,
-    ):
+    def __init__(self, config, layer_id: int):
         super().__init__()
-        self.attention = AttnFactory.create(
-            attn_type,
-            dim=dim,
-            n_heads=n_heads,
-            n_kv_heads=n_kv_heads,
-            use_qk_norm=use_qk_norm,
-            norm_eps=norm_eps,
-            use_gated_attention=use_gated_attention,
-            layer_id=layer_id,
-            **kwargs,
-        )
-        self.input_norm = RMSNorm(dim, norm_eps)
-        self.post_attention_norm = RMSNorm(dim, norm_eps)
-        self.mlp = FFNFactory.create(ffn_type, dim, dim_ffn, **kwargs)
+        cfg = asdict(config)
+        self.attention = AttnFactory.create(config.attn_type, **cfg, layer_id=layer_id)
+        self.input_norm = RMSNorm(config.dim, config.norm_eps)
+        self.post_attention_norm = RMSNorm(config.dim, config.norm_eps)
+        self.mlp = FFNFactory.create(config.ffn_type, **cfg)
 
     def forward(
         self,
diff --git a/astrai/model/encoder.py b/astrai/model/encoder.py
index 2957be7..1ea1b2c 100644
--- a/astrai/model/encoder.py
+++ b/astrai/model/encoder.py
@@ -26,19 +26,7 @@ class EmbeddingEncoder(AutoModel):
         self.embed_tokens = Embedding(config.vocab_size, config.dim)
 
         self.layers = nn.ModuleList(
-            [
-                DecoderBlock(
-                    config.dim,
-                    config.n_heads,
-                    config.dim_ffn,
-                    config.n_kv_heads,
-                    config.norm_eps,
-                    config.use_qk_norm,
-                    config.use_gated_attention,
-                    layer_id,
-                )
-                for layer_id in range(config.n_layers)
-            ]
+            [DecoderBlock(config, layer_id) for layer_id in range(config.n_layers)]
         )
 
         self.norm = RMSNorm(config.dim, config.norm_eps)
diff --git a/astrai/model/transformer.py b/astrai/model/transformer.py
index 9a2a7e2..44655f8 100644
--- a/astrai/model/transformer.py
+++ b/astrai/model/transformer.py
@@ -62,28 +62,7 @@ class AutoRegressiveLM(AutoModel):
         self.embed_tokens = Embedding(config.vocab_size, config.dim)
 
         self.layers = nn.ModuleList(
-            [
-                DecoderBlock(
-                    config.dim,
-                    config.n_heads,
-                    config.dim_ffn,
-                    config.n_kv_heads,
-                    config.norm_eps,
-                    config.use_qk_norm,
-                    config.use_gated_attention,
-                    layer_id,
-                    attn_type=config.attn_type,
-                    ffn_type=config.ffn_type,
-                    n_routed_experts=config.n_routed_experts,
-                    n_shared_experts=config.n_shared_experts,
-                    n_activated_experts=config.n_activated_experts,
-                    topk_method=config.topk_method,
-                    kv_lora_rank=config.kv_lora_rank,
-                    qk_nope_head_dim=config.qk_nope_head_dim,
-                    qk_rope_head_dim=config.qk_rope_head_dim,
-                )
-                for layer_id in range(config.n_layers)
-            ]
+            [DecoderBlock(config, layer_id) for layer_id in range(config.n_layers)]
         )
 
         self.norm = RMSNorm(config.dim, config.norm_eps)