Optimize memory efficiency in adaptive model architecture

[csmangum]

Related to #6

Optimize memory efficiency in adaptive models by implementing conditional computation, parameter sharing, and low-rank approximations.

• Conditional Computation Architecture:

  • Modify AdaptiveEntropyBottleneck in meaning_transform/src/models/adaptive_entropy_bottleneck.py to create projection layers only if compression exceeds a threshold.
  • Add low-rank approximations for large projections in AdaptiveEntropyBottleneck.

• Parameter Sharing in FeatureGroupedVAE:

  • Implement parameter sharing across feature groups in FeatureGroupedVAE in meaning_transform/src/models/feature_grouped_vae.py.
  • Update FeatureGroupedVAE to use shared components for each feature group.

• Documentation Update:

  • Update docs/agent_memory_architecture.md to reflect the new architecture with conditional computation, parameter sharing, and low-rank approximations.

---

For more details, open the Copilot Workspace session.

[Copilot]

Pull Request Overview

This PR optimizes memory efficiency for adaptive models by introducing conditional computation, low-rank approximations, and parameter sharing.

• Modified AdaptiveEntropyBottleneck to conditionally create projection layers based on a compression threshold and to use low-rank approximations for large projections.
• Updated FeatureGroupedVAE to share a common compressor across feature groups and replaced group-specific bottlenecks with shared components.
• Revised documentation in docs/agent_memory_architecture.md to reflect the new architectural changes.

Reviewed Changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated no comments.

File	Description
meaning_transform/src/models/feature_grouped_vae.py	Implemented shared compressor and updated loss and rate computations for groups.
meaning_transform/src/models/adaptive_entropy_bottleneck.py	Added conditional logic for projection layers and integrated low-rank approximations.
docs/agent_memory_architecture.md	Updated documentation to include details on conditional computation and sharing.

Comments suppressed due to low confidence (3)

meaning_transform/src/models/feature_grouped_vae.py:74

• Consider defining a dedicated nn.Module subclass for the shared compressor to encapsulate the mu and scale networks, as it improves clarity and maintainability.

self.shared_compressor = nn.Module()

meaning_transform/src/models/feature_grouped_vae.py:246

• Review the compression loss computation to ensure it scales appropriately for each feature group and maintains numerical stability; consider extracting the constant into a predefined variable.

compression_loss += 0.5 * log_scale_group.mul(2).exp() + 0.5 * torch.log(2 * torch.tensor(torch.pi, device=z.device))

meaning_transform/src/models/adaptive_entropy_bottleneck.py:57

• Ensure that latent_dim is large enough so that latent_dim // 4 is non-zero; otherwise, the projection layers may not function as intended.

self.proj_up = nn.Sequential(
                    nn.Linear(self.effective_dim, latent_dim // 4),
                    nn.LeakyReLU(),
                    nn.Linear(latent_dim // 4, latent_dim * 2)
                )